Fix get-object with partNumber (#401)

Co-authored-by: Michael Graeb <[email protected]>
awslabs · Dec 28, 2023 · 7dd72a9 · 7dd72a9
1 parent c01ebda
commit 7dd72a9
Show file tree

Hide file tree

Showing 5 changed files with 109 additions and 14 deletions.
diff --git a/source/s3_client.c b/source/s3_client.c
@@ -1168,20 +1168,34 @@ static struct aws_s3_meta_request *s_s3_client_meta_request_factory_default(
     /* Call the appropriate meta-request new function. */
     switch (options->type) {
         case AWS_S3_META_REQUEST_TYPE_GET_OBJECT: {
-            /* If the initial request already has partNumber, the request is not
-             * splittable(?). Treat it as a Default request.
-             * TODO: Still need tests to verify that the request of a part is
-             * splittable or not */
-            if (aws_http_headers_has(initial_message_headers, aws_byte_cursor_from_c_str("partNumber"))) {
-                return aws_s3_meta_request_default_new(
-                    client->allocator,
-                    client,
-                    AWS_S3_REQUEST_TYPE_GET_OBJECT,
-                    content_length,
-                    false /*should_compute_content_md5*/,
-                    options);
+            struct aws_byte_cursor path_and_query;
+
+            if (aws_http_message_get_request_path(options->message, &path_and_query) == AWS_OP_SUCCESS) {
+                /* If the initial request already has partNumber, the request is not
+                 * splittable(?). Treat it as a Default request.
+                 * TODO: Still need tests to verify that the request of a part is
+                 * splittable or not */
+                struct aws_byte_cursor sub_string;
+                AWS_ZERO_STRUCT(sub_string);
+                /* The first split on '?' for path and query is path, the second is query */
+                if (aws_byte_cursor_next_split(&path_and_query, '?', &sub_string) == true) {
+                    aws_byte_cursor_next_split(&path_and_query, '?', &sub_string);
+                    struct aws_uri_param param;
+                    AWS_ZERO_STRUCT(param);
+                    struct aws_byte_cursor part_number_query_str = aws_byte_cursor_from_c_str("partNumber");
+                    while (aws_query_string_next_param(sub_string, &param)) {
+                        if (aws_byte_cursor_eq(&param.key, &part_number_query_str)) {
+                            return aws_s3_meta_request_default_new(
+                                client->allocator,
+                                client,
+                                AWS_S3_REQUEST_TYPE_GET_OBJECT,
+                                content_length,
+                                false /*should_compute_content_md5*/,
+                                options);
+                        }
+                    }
+                }
             }
-
             return aws_s3_meta_request_auto_ranged_get_new(client->allocator, client, part_size, options);
         }
         case AWS_S3_META_REQUEST_TYPE_PUT_OBJECT: {

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -73,6 +73,7 @@ add_net_test_case(test_s3_get_object_sse_aes256)
 add_net_test_case(test_s3_get_object_backpressure_small_increments)
 add_net_test_case(test_s3_get_object_backpressure_big_increments)
 add_net_test_case(test_s3_get_object_backpressure_initial_size_zero)
+add_net_test_case(test_s3_get_object_part)
 add_net_test_case(test_s3_no_signing)
 add_net_test_case(test_s3_signing_override)
 add_net_test_case(test_s3_put_object_tls_disabled)

diff --git a/tests/s3_data_plane_tests.c b/tests/s3_data_plane_tests.c
@@ -1614,6 +1614,78 @@ static int s_test_s3_get_object_backpressure_initial_size_zero(struct aws_alloca
     return s_test_s3_get_object_backpressure_helper(allocator, part_size, window_initial_size, window_increment_size);
 }
 
+AWS_TEST_CASE(test_s3_get_object_part, s_test_s3_get_object_part)
+static int s_test_s3_get_object_part(struct aws_allocator *allocator, void *ctx) {
+    (void)ctx;
+
+    struct aws_s3_tester tester;
+    AWS_ZERO_STRUCT(tester);
+    ASSERT_SUCCESS(aws_s3_tester_init(allocator, &tester));
+
+    struct aws_s3_client_config client_config = {
+        .part_size = MB_TO_BYTES(8),
+    };
+
+    ASSERT_SUCCESS(aws_s3_tester_bind_client(
+        &tester, &client_config, AWS_S3_TESTER_BIND_CLIENT_REGION | AWS_S3_TESTER_BIND_CLIENT_SIGNING));
+
+    struct aws_s3_client *client = aws_s3_client_new(allocator, &client_config);
+    struct aws_s3_meta_request_test_results meta_request_test_results;
+    aws_s3_meta_request_test_results_init(&meta_request_test_results, allocator);
+
+    /*** PUT FILE ***/
+
+    struct aws_byte_buf path_buf;
+    AWS_ZERO_STRUCT(path_buf);
+
+    ASSERT_SUCCESS(
+        aws_s3_tester_upload_file_path_init(allocator, &path_buf, aws_byte_cursor_from_c_str("/get_object_part_test")));
+
+    struct aws_byte_cursor object_path = aws_byte_cursor_from_buf(&path_buf);
+
+    struct aws_s3_tester_meta_request_options put_options = {
+        .allocator = allocator,
+        .meta_request_type = AWS_S3_META_REQUEST_TYPE_PUT_OBJECT,
+        .client = client,
+        .put_options =
+            {
+                .object_size_mb = 10,
+                .object_path_override = object_path,
+            },
+    };
+
+    ASSERT_SUCCESS(aws_s3_tester_send_meta_request_with_options(&tester, &put_options, NULL));
+
+    /* GET FILE */
+
+    struct aws_s3_tester_meta_request_options options = {
+        .allocator = allocator,
+        .client = client,
+        .meta_request_type = AWS_S3_META_REQUEST_TYPE_GET_OBJECT,
+        .validate_type = AWS_S3_TESTER_VALIDATE_TYPE_NO_VALIDATE,
+        .get_options =
+            {
+                .object_path = object_path,
+                .part_number = 2,
+            },
+    };
+
+    ASSERT_SUCCESS(aws_s3_tester_send_meta_request_with_options(&tester, &options, &meta_request_test_results));
+
+    ASSERT_UINT_EQUALS(AWS_ERROR_SUCCESS, meta_request_test_results.finished_error_code);
+    /* Only one request was made to get the second part of the object */
+    ASSERT_UINT_EQUALS(1, aws_array_list_length(&meta_request_test_results.synced_data.metrics));
+
+    aws_s3_meta_request_test_results_clean_up(&meta_request_test_results);
+
+    client = aws_s3_client_release(client);
+
+    aws_byte_buf_clean_up(&path_buf);
+    aws_s3_tester_clean_up(&tester);
+
+    return 0;
+}
+
 static int s_test_s3_put_object_helper(
     struct aws_allocator *allocator,
     enum aws_s3_client_tls_usage tls_usage,

diff --git a/tests/s3_tester.c b/tests/s3_tester.c
@@ -136,7 +136,7 @@ static int s_s3_test_meta_request_body_callback(
 
     /* If this is an auto-ranged-get meta request, then grab the object range start so that the expected_range_start can
      * be properly offset.*/
-    if (meta_request->type == AWS_S3_META_REQUEST_TYPE_GET_OBJECT) {
+    if (meta_request->type == AWS_S3_META_REQUEST_TYPE_GET_OBJECT && meta_request->part_size != 0) {
 
         aws_s3_meta_request_lock_synced_data(meta_request);
 
@@ -1499,6 +1499,12 @@ int aws_s3_tester_send_meta_request_with_options(
 
             struct aws_http_message *message = aws_s3_test_get_object_request_new(
                 allocator, aws_byte_cursor_from_string(host_name), options->get_options.object_path);
+            ASSERT_SUCCESS(aws_s3_message_util_set_multipart_request_path(
+                allocator,
+                NULL /*upload_id*/,
+                options->get_options.part_number,
+                false /*append_uploads_suffix*/,
+                message));
 
             if (options->get_options.object_range.ptr != NULL) {
                 struct aws_http_header range_header = {

diff --git a/tests/s3_tester.h b/tests/s3_tester.h
@@ -183,6 +183,8 @@ struct aws_s3_tester_meta_request_options {
     struct {
         struct aws_byte_cursor object_path;
         struct aws_byte_cursor object_range;
+        /* Get the part from S3, starts from 1. 0 means not set. */
+        int part_number;
     } get_options;
 
     /* Put Object Meta request specific options. */