Fix: Get object with checksum leak when retry happens

.github/workflows/ci.yml

@@ -16,6 +16,7 @@ env:
   AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
   AWS_DEFAULT_REGION: ${{ secrets.AWS_DEFAULT_REGION }}
   AWS_REGION: us-east-1
+  CTEST_PARALLEL_LEVEL: 4
 
 jobs:
   linux-compat:

source/s3_meta_request.c

@@ -982,12 +982,12 @@ static void s_get_response_part_finish_checksum_helper(struct aws_s3_connection
         request->validation_algorithm = request->request_level_running_response_sum->algorithm;
         aws_byte_buf_clean_up(&response_body_sum);
         aws_byte_buf_clean_up(&encoded_response_body_sum);
-        aws_checksum_destroy(request->request_level_running_response_sum);
-        aws_byte_buf_clean_up(&request->request_level_response_header_checksum);
-        request->request_level_running_response_sum = NULL;
     } else {
         request->did_validate = false;
     }
+    aws_checksum_destroy(request->request_level_running_response_sum);
+    aws_byte_buf_clean_up(&request->request_level_response_header_checksum);
+    request->request_level_running_response_sum = NULL;
 }
 
 static int s_s3_meta_request_incoming_headers(

tests/CMakeLists.txt

@@ -254,6 +254,7 @@ add_net_test_case(test_s3_list_bucket_valid)
 if(ENABLE_MOCK_SERVER_TESTS)
     add_net_test_case(multipart_upload_mock_server)
     add_net_test_case(multipart_upload_checksum_with_retry_mock_server)
+    add_net_test_case(multipart_download_checksum_with_retry_mock_server)
     add_net_test_case(async_internal_error_from_complete_multipart_mock_server)
     add_net_test_case(async_access_denied_from_complete_multipart_mock_server)
     add_net_test_case(get_object_modified_mock_server)

tests/mock_s3_server/GetObject/get_object_checksum_retry.json

@@ -0,0 +1,15 @@
+{
+    "status": 200,
+    "headers": {
+      "ETag": "b54357faf0632cce46e942fa68356b38",
+      "Date": "Thu, 12 Jan 2023 00:04:21 GMT",
+      "Last-Modified": "Tue, 10 Jan 2023 23:39:32 GMT",
+      "Accept-Ranges": "bytes",
+      "Content-Range": "bytes 0-65535/65536",
+      "Content-Type": "binary/octet-stream",
+      "x-amz-checksum-crc32": "q1875w=="
+    },
+    "body": [
+      "less data"
+    ]
+}

tests/mock_s3_server/mock_s3_server.py

@@ -19,6 +19,8 @@
 VERBOSE = False
 SHOULD_THROTTLE = True
 
+COUNT = 0
+RETRY_TEST = False
 
 class S3Opts(Enum):
     CreateMultipartUpload = 1
@@ -193,18 +195,29 @@ async def send_response_from_json(wrapper, response_json_path, chunked=False, ge
                  "response with", len(body), "bytes")
 
     headers = wrapper.basic_headers()
+    content_length_set = False
     for header in data['headers'].items():
-        headers.append((header[0], header[1]))
+        headers.append((header[0], str(header[1])))
+        if header[0].lower() == "content-length":
+            content_length_set = True
 
     if chunked:
         headers.append(('Transfer-Encoding', "chunked"))
         res = h11.Response(status_code=status_code, headers=headers)
         await wrapper.send(res)
         await wrapper.send(h11.Data(data=b"%X\r\n%s\r\n" % (len(body), body.encode())))
     else:
-        headers.append(("Content-Length", str(len(body))))
-        res = h11.Response(status_code=status_code, headers=headers)
-        await wrapper.send(res)
+        if COUNT <= 1 and RETRY_TEST:
+            headers.append(("Content-Length", str(123456)))
+
+        elif content_length_set is False:
+            headers.append(("Content-Length", str(len(body))))
+
+        try:
+            res = h11.Response(status_code=status_code, headers=headers)
+            await wrapper.send(res)
+        except Exception as e:
+            print(e)
         if head_request:
             await wrapper.send(h11.EndOfMessage())
             return
@@ -230,7 +243,7 @@ async def send_mock_s3_response(wrapper, request_type, path, generate_body=False
             wrapper.info("Throttling")
         # Flip the flag
         wrapper.should_throttle = not wrapper.should_throttle
-    await send_response_from_json(wrapper, response_file, generate_body=generate_body, generate_body_size=generate_body_size)
+    await send_response_from_json(wrapper, response_file, generate_body=generate_body, generate_body_size=generate_body_size, head_request=head_request)
 
 
 async def maybe_send_error_response(wrapper, exc):
@@ -279,7 +292,15 @@ def handle_get_object_modified(start_range, end_range, request):
             return "/get_object_modified_failure", data_length, False
 
 
-def handle_get_object(request, parsed_path):
+def handle_get_object(wrapper, request, parsed_path, head_request=False):
+    global COUNT
+    global RETRY_TEST
+    if parsed_path.path == "/get_object_checksum_retry" and not head_request:
+        COUNT = COUNT + 1
+        RETRY_TEST = True
+    else:
+        COUNT = 0
+        RETRY_TEST = False
 
     body_range_value = get_request_header_value(request, "range")
 
@@ -340,7 +361,7 @@ async def handle_mock_s3_request(wrapper, request):
         else:
             request_type = S3Opts.GetObject
             response_path, generate_body_size, generate_body = handle_get_object(
-                request, parsed_path)
+                wrapper, request, parsed_path, head_request=method == "HEAD" )
     else:
         # TODO: support more type.
         wrapper.info("unsupported request:", request)

tests/s3_mock_server_tests.c

@@ -194,6 +194,50 @@ TEST_CASE(multipart_upload_checksum_with_retry_mock_server) {
     return AWS_OP_SUCCESS;
 }
 
+TEST_CASE(multipart_download_checksum_with_retry_mock_server) {
+    (void)ctx;
+    /**
+     * We had a memory leak after the header of the request received successfully, the request failed.
+     * We have allocated memory that never frees.
+     */
+    struct aws_s3_tester tester;
+    ASSERT_SUCCESS(aws_s3_tester_init(allocator, &tester));
+    struct aws_s3_tester_client_options client_options = {
+        .part_size = MB_TO_BYTES(5),
+        .tls_usage = AWS_S3_TLS_DISABLED,
+    };
+
+    struct aws_s3_client *client = NULL;
+    ASSERT_SUCCESS(aws_s3_tester_client_new(&tester, &client_options, &client));
+    /* Mock server will response without fake checksum for the body */
+    struct aws_byte_cursor object_path = aws_byte_cursor_from_c_str("/get_object_checksum_retry");
+
+    struct aws_s3_tester_meta_request_options get_options = {
+        .allocator = allocator,
+        .meta_request_type = AWS_S3_META_REQUEST_TYPE_GET_OBJECT,
+        .client = client,
+        .expected_validate_checksum_alg = AWS_SCA_CRC32,
+        .validate_get_response_checksum = true,
+        .get_options =
+            {
+                .object_path = object_path,
+            },
+        .default_type_options =
+            {
+                .mode = AWS_S3_TESTER_DEFAULT_TYPE_MODE_GET,
+            },
+        .mock_server = true,
+        .validate_type = AWS_S3_TESTER_VALIDATE_TYPE_EXPECT_FAILURE,
+    };
+
+    ASSERT_SUCCESS(aws_s3_tester_send_meta_request_with_options(&tester, &get_options, NULL));
+
+    aws_s3_client_release(client);
+    aws_s3_tester_clean_up(&tester);
+
+    return AWS_OP_SUCCESS;
+}
+
 TEST_CASE(async_internal_error_from_complete_multipart_mock_server) {
     (void)ctx;