fix cleanup and test

awslabs · Jan 13, 2025 · 96e1ab2 · 96e1ab2
1 parent d01fab7
commit 96e1ab2
Show file tree

Hide file tree

Showing 5 changed files with 77 additions and 13 deletions.
diff --git a/source/s3_copy_object.c b/source/s3_copy_object.c
@@ -82,7 +82,8 @@ struct aws_s3_meta_request *aws_s3_meta_request_copy_object_new(
             copy_object,
             &s_s3_copy_object_vtable,
             &copy_object->base)) {
-        goto on_error;
+        aws_mem_release(allocator, copy_object);
+        return NULL;
     }
 
     aws_array_list_init_dynamic(
@@ -105,7 +106,7 @@ struct aws_s3_meta_request *aws_s3_meta_request_copy_object_new(
 
     return &copy_object->base;
 on_error:
-    aws_mem_release(allocator, copy_object);
+    aws_s3_meta_request_release(&copy_object->base);
     return NULL;
 }
 

diff --git a/source/s3_request_messages.c b/source/s3_request_messages.c
@@ -460,12 +460,13 @@ struct aws_http_message *aws_s3_get_source_object_size_message_new(
     struct aws_allocator *allocator,
     struct aws_http_message *base_message,
     struct aws_uri *source_uri) {
-    struct aws_http_message *message = NULL;
 
-    message = aws_http_message_new_request(allocator);
+    struct aws_http_message *message = aws_http_message_new_request(allocator);
     if (message == NULL) {
         goto error_cleanup;
     }
+    struct aws_byte_buf head_object_host_header;
+    AWS_ZERO_STRUCT(head_object_host_header);
 
     if (aws_http_message_set_request_method(message, g_head_method)) {
         goto error_cleanup;
@@ -474,6 +475,10 @@ struct aws_http_message *aws_s3_get_source_object_size_message_new(
         /* Parse source host header and path from the provided URI */
         struct aws_byte_cursor host = *aws_uri_host_name(source_uri);
         struct aws_byte_cursor path = *aws_uri_path(source_uri);
+        if (host.len == 0 || path.len == 0) {
+            aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
+            goto error_cleanup;
+        }
         struct aws_http_header host_header = {
             .name = g_host_header_name,
             .value = host,
@@ -489,28 +494,26 @@ struct aws_http_message *aws_s3_get_source_object_size_message_new(
     }
 
     /* Parse the source host header and path from the x-amz-copy-source header and the destination URI */
-    struct aws_byte_buf head_object_host_header;
-    AWS_ZERO_STRUCT(head_object_host_header);
 
     AWS_PRECONDITION(allocator);
 
     /* Find the x-amz-copy-source header, to extract source bucket/key information. */
     struct aws_http_headers *headers = aws_http_message_get_headers(base_message);
     if (!headers) {
         AWS_LOGF_ERROR(AWS_LS_S3_GENERAL, "CopyRequest is missing headers");
-        return NULL;
+        goto error_cleanup;
     }
 
     struct aws_byte_cursor source_header;
     const struct aws_byte_cursor copy_source_header = AWS_BYTE_CUR_INIT_FROM_STRING_LITERAL("x-amz-copy-source");
     if (aws_http_headers_get(headers, copy_source_header, &source_header) != AWS_OP_SUCCESS) {
         AWS_LOGF_ERROR(AWS_LS_S3_GENERAL, "CopyRequest is missing the x-amz-copy-source header");
-        return NULL;
+        goto error_cleanup;
     }
     struct aws_byte_cursor host;
     if (aws_http_headers_get(headers, g_host_header_name, &host) != AWS_OP_SUCCESS) {
         AWS_LOGF_ERROR(AWS_LS_S3_GENERAL, "CopyRequest is missing the Host header");
-        return NULL;
+        goto error_cleanup;
     }
 
     struct aws_byte_cursor request_path = source_header;

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -302,6 +302,7 @@ add_net_test_case(test_s3_multipart_copy_large_object_special_char)
 add_net_test_case(test_s3_multipart_copy_large_object)
 add_net_test_case(test_s3_copy_object_invalid_source_key)
 add_net_test_case(test_s3_copy_source_prefixed_by_slash)
+add_net_test_case(test_s3_copy_invalid_source_uri)
 add_net_test_case(test_s3_copy_source_prefixed_by_slash_multipart)
 add_net_test_case(test_s3_put_pause_resume_happy_path)
 add_net_test_case(test_s3_put_pause_resume_all_parts_done)

diff --git a/tests/s3_data_plane_tests.c b/tests/s3_data_plane_tests.c
@@ -6712,6 +6712,58 @@ static int s_test_s3_copy_source_prefixed_by_slash(struct aws_allocator *allocat
     return AWS_OP_SUCCESS;
 }
 
+AWS_TEST_CASE(test_s3_copy_invalid_source_uri, s_test_s3_copy_invalid_source_uri)
+static int s_test_s3_copy_invalid_source_uri(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_byte_cursor source_key = AWS_BYTE_CUR_INIT_FROM_STRING_LITERAL("pre-existing-1MB");
+    struct aws_byte_cursor destination_key = AWS_BYTE_CUR_INIT_FROM_STRING_LITERAL("copies/destination_1MB");
+
+    struct aws_byte_cursor source_bucket = g_test_bucket_name;
+
+    char copy_source_value[1024];
+    snprintf(
+        copy_source_value,
+        sizeof(copy_source_value),
+        "/%.*s/%.*s",
+        (int)source_bucket.len,
+        source_bucket.ptr,
+        (int)source_key.len,
+        source_key.ptr);
+
+    struct aws_byte_cursor x_amz_copy_source = aws_byte_cursor_from_c_str(copy_source_value);
+    struct aws_byte_cursor destination_bucket = g_test_bucket_name;
+
+    char endpoint[1024];
+    snprintf(
+        endpoint,
+        sizeof(endpoint),
+        "%.*s.s3.%s.amazonaws.com",
+        (int)destination_bucket.len,
+        destination_bucket.ptr,
+        g_test_s3_region.ptr);
+
+    struct aws_byte_cursor copy_source_uri = aws_byte_cursor_from_c_str("http://invalid-uri.com:80:80/path");
+
+    ASSERT_SUCCESS(aws_test_s3_copy_object_from_x_amz_copy_source(
+        allocator,
+        &tester,
+        x_amz_copy_source,
+        aws_byte_cursor_from_c_str(endpoint),
+        destination_key,
+        AWS_ERROR_MALFORMED_INPUT_STRING,
+        0,
+        MB_TO_BYTES(1),
+        false,
+        copy_source_uri));
+
+    aws_s3_tester_clean_up(&tester);
+    return AWS_OP_SUCCESS;
+}
+
 /**
  * Test multipart Copy Object meta request using a slash prefix in the x_amz_copy_source header.
  * S3 supports both bucket/key and /bucket/key

diff --git a/tests/s3_tester.c b/tests/s3_tester.c
@@ -2497,7 +2497,14 @@ int aws_test_s3_copy_object_from_x_amz_copy_source(
     }
 
     struct aws_s3_meta_request *meta_request = aws_s3_client_make_meta_request(client, &meta_request_options);
-    ASSERT_NOT_NULL(meta_request);
+    if (meta_request == NULL) {
+        if (expected_error_code == AWS_OP_SUCCESS) {
+            AWS_FATAL_ASSERT(false && "meta_request is NULL");
+        } else {
+            ASSERT_INT_EQUALS(expected_error_code, aws_last_error());
+            goto cleanup;
+        }
+    }
 
     /* wait completion of the meta request */
     aws_mutex_lock(&test_data.mutex);
@@ -2512,11 +2519,11 @@ int aws_test_s3_copy_object_from_x_amz_copy_source(
     if (test_data.meta_request_error_code == AWS_ERROR_SUCCESS) {
         ASSERT_UINT_EQUALS(expected_size, test_data.progress_callback_total_bytes_transferred);
         ASSERT_UINT_EQUALS(expected_size, test_data.progress_callback_content_length);
+        /* assert headers callback was invoked */
+        ASSERT_TRUE(test_data.headers_callback_was_invoked);
     }
 
-    /* assert headers callback was invoked */
-    ASSERT_TRUE(test_data.headers_callback_was_invoked);
-
+cleanup:
     aws_s3_meta_request_release(meta_request);
     aws_mutex_clean_up(&test_data.mutex);
     aws_http_message_destroy(message);