Add test and update mock server

Author: waahm7

source/s3express_credentials_provider.c

@@ -397,16 +397,20 @@ static void s_on_request_finished(
 static struct aws_http_message *s_create_session_request_new(
     struct aws_allocator *allocator,
     struct aws_byte_cursor host_value,
-    struct aws_http_headers *headers) {
+    struct aws_http_headers *headers,
+    const struct aws_uri *endpoint_override) {
     struct aws_http_message *request = aws_http_message_new_request(allocator);
 
     struct aws_http_header host_header = {
         .name = g_host_header_name,
         .value = host_value,
     };
 
-    if (aws_http_message_add_header(request, host_header)) {
-        goto error;
+    /* NOTE: ONLY FOR TESTS. Don't add the host header for endpoint override. */
+    if (endpoint_override == NULL) {
+        if (aws_http_message_add_header(request, host_header)) {
+            goto error;
+        }
     }
 
     struct aws_http_header user_agent_header = {
@@ -436,7 +440,14 @@ static struct aws_http_message *s_create_session_request_new(
         goto error;
     }
 
-    if (aws_http_message_set_request_path(request, s_create_session_path_query)) {
+    struct aws_byte_cursor path_and_query = s_create_session_path_query;
+    if (endpoint_override != NULL) {
+        const struct aws_byte_cursor *override_path_query = aws_uri_path_and_query(endpoint_override);
+        if (override_path_query->len > 0) {
+            path_and_query = *override_path_query;
+        }
+    }
+    if (aws_http_message_set_request_path(request, path_and_query)) {
         goto error;
     }
     return request;
@@ -528,15 +539,14 @@ static struct aws_s3express_session_creator *s_session_creator_new(
     const struct aws_credentials_properties_s3express *s3express_properties) {
 
     struct aws_s3express_credentials_provider_impl *impl = provider->impl;
-    struct aws_http_message *request =
-        s_create_session_request_new(provider->allocator, s3express_properties->host, s3express_properties->headers);
+    struct aws_http_message *request = s_create_session_request_new(
+        provider->allocator,
+        s3express_properties->host,
+        s3express_properties->headers,
+        impl->mock_test.endpoint_override);
     if (!request) {
         return NULL;
     }
-    if (impl->mock_test.endpoint_override) {
-        /* NOTE: ONLY FOR TESTS. Erase the host header for endpoint override. */
-        aws_http_headers_erase(aws_http_message_get_headers(request), g_host_header_name);
-    }
 
     struct aws_s3express_session_creator *session_creator =
         aws_mem_calloc(provider->allocator, 1, sizeof(struct aws_s3express_session_creator));

tests/CMakeLists.txt

@@ -351,6 +351,7 @@ if(ENABLE_MOCK_SERVER_TESTS)
     add_net_test_case(s3express_provider_background_refresh_mock_server)
     add_net_test_case(s3express_provider_background_refresh_remove_inactive_creds_mock_server)
     add_net_test_case(s3express_provider_stress_mock_server)
+    add_net_test_case(s3express_provider_get_credentials_kms_headers_mock_server)
 
     add_net_test_case(s3express_client_sanity_test_mock_server)
     add_net_test_case(s3express_client_sanity_override_test_mock_server)

tests/mock_s3_server/CompleteMultipartUpload/sse_kms.json

@@ -0,0 +1,14 @@
+{
+  "status": 200,
+  "headers": {"Connection": "close",
+   "x-amz-server-side-encryption": "aws:kms"},
+  "body": [
+    "<?xml version=\"1.0\" encoding=\"UTF-8\"?>",
+    "<CompleteMultipartUploadResult xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\">",
+    "<Location>http://default.s3.us-west-2.amazonaws.com/default</Location>",
+    "<Bucket>default</Bucket>",
+    "<Key>default</Key>",
+    "<ETag>\"3858f62230ac3c915f300c664312c11f-9\"</ETag>",
+   "</CompleteMultipartUploadResult>"
+  ]
+}

tests/mock_s3_server/CreateSession/sse_kms.json

@@ -0,0 +1,21 @@
+{
+  "status": 200,
+  "headers": {
+    "x-amz-request-id": "12345",
+    "x-amz-server-side-encryption": "aws:kms"
+  },
+  "request_headers": {
+    "x-amz-server-side-encryption": "aws:kms"
+  },
+  "body": [
+    "<?xml version=\"1.0\" encoding=\"UTF-8\"?>",
+    "<CreateSessionResult xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\">",
+    "<Credentials>",
+    "<SessionToken>sessionToken</SessionToken>",
+    "<SecretAccessKey>secretKey</SecretAccessKey>",
+    "<AccessKeyId>accessKeyId</AccessKeyId>",
+    "<Expiration>2023-06-26T17:33:30Z</Expiration>",
+    "</Credentials>",
+    "</CreateSessionResult>"
+  ]
+}

tests/mock_s3_server/mock_s3_server.py

@@ -56,7 +56,7 @@ class ResponseConfig:
     json_path: str = None
     throttle: bool = False
     force_retry: bool = False
-
+    request_headers = None
     def _resolve_file_path(self, wrapper, request_type):
         global SHOULD_THROTTLE
         if self.json_path is None:
@@ -85,6 +85,16 @@ def resolve_response(self, wrapper, request_type, chunked=False, head_request=Fa
                      ".\n generate_body_size: ", self.generate_body_size)
         with open(self.json_path, 'r') as f:
             data = json.load(f)
+        headers = wrapper.basic_headers()
+
+        # If request_headers is present, validate that the request contains all required headers
+        if 'request_headers' in data:
+                for header in data['request_headers']:
+                    header_bytes = header.encode('utf-8')
+                    if not any(header_bytes == h[0] for h in self.request_headers):
+                        response =  Response(status_code=500, delay=0, headers=headers,
+                                             data=json.dumps({'error': f"Missing required header: {header}"}), chunked=chunked, head_request=head_request)
+                        return response
 
         # if response has delay, then sleep before sending it
         delay = data.get('delay', 0)
@@ -95,7 +105,6 @@ def resolve_response(self, wrapper, request_type, chunked=False, head_request=Fa
         else:
             body = "\n".join(data['body'])
 
-        headers = wrapper.basic_headers()
         content_length_set = False
         for header in data['headers'].items():
             headers.append((header[0], str(header[1])))
@@ -499,6 +508,7 @@ async def handle_mock_s3_request(wrapper, request):
 
     if response_config is None:
         response_config = ResponseConfig(parsed_path.path)
+    response_config.request_headers = request.headers
 
     response = response_config.resolve_response(
         wrapper, request_type, head_request=method == "HEAD")

tests/s3_mock_server_s3express_provider_test.c

@@ -45,6 +45,8 @@ struct aws_s3express_provider_tester {
     struct aws_uri mock_server;
     struct aws_s3_client *client;
 
+    aws_simple_completion_callback *on_provider_shutdown_callback;
+    void *shutdown_user_data;
     int error_code;
 };
 
@@ -252,6 +254,87 @@ TEST_CASE(s3express_provider_get_credentials_mock_server) {
     return AWS_OP_SUCCESS;
 }
 
+struct aws_s3express_credentials_provider *s_s3express_credentials_provider_factory(
+    struct aws_allocator *allocator,
+    struct aws_s3_client *client,
+    aws_simple_completion_callback on_provider_shutdown_callback,
+    void *shutdown_user_data,
+    void *factory_user_data) {
+    (void)allocator;
+    (void)client;
+    (void)on_provider_shutdown_callback;
+    (void)shutdown_user_data;
+    (void)factory_user_data;
+    s_s3express_tester.on_provider_shutdown_callback = on_provider_shutdown_callback;
+    s_s3express_tester.shutdown_user_data = shutdown_user_data;
+    struct aws_s3express_credentials_provider_default_options options = {
+        .client = client,
+        // .shutdown_complete_callback = s_on_shutdown_complete,
+        // .shutdown_user_data = &s_s3express_tester,
+        .mock_test.bg_refresh_secs_override = s_bg_refresh_secs_override,
+    };
+    struct aws_s3express_credentials_provider *provider =
+        aws_s3express_credentials_provider_new_default(allocator, &options);
+    struct aws_s3express_credentials_provider_impl *impl = provider->impl;
+    impl->mock_test.endpoint_override = &s_s3express_tester.mock_server;
+    impl->mock_test.s3express_session_is_valid_override = s_s3express_session_always_true;
+
+    return provider;
+}
+
+TEST_CASE(s3express_provider_get_credentials_kms_headers_mock_server) {
+    (void)ctx;
+
+    struct aws_s3_tester tester;
+    ASSERT_SUCCESS(aws_s3_tester_init(allocator, &tester));
+    ASSERT_SUCCESS(s_s3express_tester_init(allocator));
+
+    struct aws_s3_tester_client_options client_options = {
+        .part_size = MB_TO_BYTES(5),
+        .tls_usage = AWS_S3_TLS_DISABLED,
+        .s3express_provider_override_factory = s_s3express_credentials_provider_factory,
+        .factory_user_data = NULL,
+    };
+
+    struct aws_s3_client *client = NULL;
+    ASSERT_SUCCESS(aws_s3_tester_client_new(&tester, &client_options, &client));
+
+    struct aws_byte_cursor object_path = aws_byte_cursor_from_c_str("/sse_kms");
+    char uri[1024] = {'\0'};
+    snprintf(uri, sizeof(uri), "" PRInSTR "sse_kms?session=", AWS_BYTE_CURSOR_PRI(g_mock_server_uri));
+    struct aws_byte_cursor uri_cursor = aws_byte_cursor_from_c_str(uri);
+    aws_uri_clean_up(&s_s3express_tester.mock_server);
+    ASSERT_SUCCESS(aws_uri_init_parse(&s_s3express_tester.mock_server, allocator, &uri_cursor));
+
+    struct aws_s3_tester_meta_request_options put_options = {
+        .allocator = allocator,
+        .meta_request_type = AWS_S3_META_REQUEST_TYPE_PUT_OBJECT,
+        .client = client,
+        .checksum_algorithm = AWS_SCA_CRC32,
+        .validate_get_response_checksum = false,
+        .put_options =
+            {
+                .object_size_mb = 10,
+                .object_path_override = object_path,
+            },
+        .mock_server = true,
+        .use_s3express_signing = true,
+        .sse_type = AWS_S3_TESTER_SSE_KMS,
+    };
+    struct aws_s3_meta_request_test_results out_results;
+    aws_s3_meta_request_test_results_init(&out_results, allocator);
+    ASSERT_SUCCESS(aws_s3_tester_send_meta_request_with_options(&tester, &put_options, &out_results));
+    aws_s3_meta_request_test_results_clean_up(&out_results);
+    aws_s3_client_release(client);
+
+    /* Call the provider shutdown callback to finish cleanup */
+    s_s3express_tester.on_provider_shutdown_callback(s_s3express_tester.shutdown_user_data);
+    ASSERT_SUCCESS(s_s3express_tester_cleanup());
+    aws_s3_tester_clean_up(&tester);
+
+    return AWS_OP_SUCCESS;
+}
+
 TEST_CASE(s3express_provider_get_credentials_multiple_mock_server) {
     (void)ctx;
 

tests/s3_tester.c

@@ -1387,6 +1387,9 @@ int aws_s3_tester_client_new(
     struct aws_s3_client_config client_config = {
         .part_size = options->part_size,
         .max_part_size = options->max_part_size,
+        .s3express_provider_override_factory = options->s3express_provider_override_factory,
+        .factory_user_data = options->factory_user_data,
+        .enable_s3express = options->s3express_provider_override_factory != NULL,
     };
     struct aws_http_proxy_options proxy_options = {
         .connection_type = AWS_HPCT_HTTP_FORWARD,

tests/s3_tester.h

@@ -140,6 +140,8 @@ struct aws_s3_tester_client_options {
     size_t num_network_interface_names;
     uint32_t setup_region : 1;
     uint32_t use_proxy : 1;
+    aws_s3express_provider_factory_fn *s3express_provider_override_factory;
+    void *factory_user_data;
 };
 
 /* should really break this up to a client setup, and a meta_request sending */