see if it's reproducing the issue

TingDaoK · TingDaoK · commit d5b12ab150b4 · 2025-08-08T14:27:45.000-07:00
diff --git a/include/aws/s3/private/s3_client_impl.h b/include/aws/s3/private/s3_client_impl.h
@@ -175,7 +175,7 @@ struct aws_s3_client_vtable {
         struct aws_http_connection *client_connection,
         const struct aws_http_make_request_options *options);
 
-    void (*after_prepare_upload_part_finish)(struct aws_s3_request *request);
+    void (*after_prepare_upload_part_finish)(struct aws_s3_request *request, struct aws_http_message *message);
 };
 
 struct aws_s3_upload_part_timeout_stats {
diff --git a/include/aws/s3/s3_client.h b/include/aws/s3/s3_client.h
@@ -884,7 +884,7 @@ struct aws_s3_meta_request_options {
      * Optional.
      * Callback for reviewing an upload before it completes.
      * WARNING: experimental/unstable
-     * See `aws_s3_upload_review_fn`
+     * See `aws_s3_meta_request_upload_review_fn`
      */
     aws_s3_meta_request_upload_review_fn *upload_review_callback;
 
diff --git a/source/s3_auto_ranged_put.c b/source/s3_auto_ranged_put.c
@@ -1184,15 +1184,15 @@ static void s_s3_prepare_upload_part_finish(struct aws_s3_prepare_upload_part_jo
         aws_future_http_message_set_error(part_prep->on_complete, aws_last_error());
         goto on_done;
     }
+    if (client->vtable->after_prepare_upload_part_finish) {
+        /* TEST ONLY, allow test to stub here. */
+        client->vtable->after_prepare_upload_part_finish(request, message);
+    }
 
     /* Success! */
     aws_future_http_message_set_result_by_move(part_prep->on_complete, &message);
 
 on_done:
-    if (client->vtable->after_prepare_upload_part_finish) {
-        /* TEST ONLY, allow test to stub here. */
-        client->vtable->after_prepare_upload_part_finish(request);
-    }
     AWS_FATAL_ASSERT(aws_future_http_message_is_done(part_prep->on_complete));
     aws_future_bool_release(part_prep->asyncstep_read_part);
     aws_future_http_message_release(part_prep->on_complete);
diff --git a/source/s3_checksum_stream.c b/source/s3_checksum_stream.c
@@ -113,7 +113,7 @@ static void s_aws_input_checksum_stream_destroy(struct aws_checksum_stream *impl
     }
 
     /* Compute the checksum of whatever was read, if we didn't reach the end of the underlying stream. */
-    // s_finalize_checksum(impl);
+    s_finalize_checksum(impl);
 
     aws_checksum_destroy(impl->checksum);
     aws_input_stream_release(impl->old_stream);
diff --git a/source/s3_chunk_stream.c b/source/s3_chunk_stream.c
@@ -167,6 +167,8 @@ static void s_aws_input_chunk_stream_destroy(struct aws_chunk_stream *impl) {
         }
         aws_byte_buf_clean_up(&impl->pre_chunk_buffer);
         aws_byte_buf_clean_up(&impl->post_chunk_buffer);
+        /* Either we calculated the checksum, or we the checksum is empty. Otherwise, something was wrong. */
+        AWS_ASSERT(impl->checksum_context->checksum_calculated || impl->checksum_context->base64_checksum.len == 0);
         aws_s3_upload_request_checksum_context_release(impl->checksum_context);
         aws_mem_release(impl->allocator, impl);
     }
diff --git a/tests/mock_s3_server/mock_s3_server.py b/tests/mock_s3_server/mock_s3_server.py
@@ -22,7 +22,6 @@
 
 # Flags to keep between requests
 SHOULD_THROTTLE = True
-SHOULD_SKIP_WAIT = True
 RETRY_REQUEST_COUNT = 0
 
 
@@ -55,16 +54,37 @@ class ResponseConfig:
     disconnect_after_headers = False
     generate_body_size: Optional[int] = None
     json_path: str = None
-    throttle: bool = False
+    forced_throttle: bool = False
     force_retry: bool = False
+    should_skip_wait = False
     request_headers: Optional[List[Tuple[bytes, bytes]]] = None
 
+    def __init__(self, path, request=None):
+        self.path = path
+        self.request_headers = request.headers
+        self.json_path = None
+        self.generate_body_size = None
+        self.disconnect_after_headers = False
+        self.forced_throttle = False
+        self.force_retry = False
+        self.should_skip_wait = False
+
+        if get_request_header_value(request, "before_finish") is not None:
+            self.should_skip_wait = True
+        if get_request_header_value(request, "force_throttle") is not None:
+            self.forced_throttle = True
+
     def _resolve_file_path(self, wrapper, request_type):
         global SHOULD_THROTTLE
         if self.json_path is None:
+            if self.forced_throttle:
+                # force the throttle to happend, instead of just 50%.
+                response_file = os.path.join(
+                    base_dir, request_type.name, f"throttle.json")
+                self.json_path = response_file
+                return
             response_file = os.path.join(
                 base_dir, request_type.name, f"{self.path[1:]}.json")
-            print(response_file)
             if os.path.exists(response_file) == False:
                 wrapper.info(
                     response_file, "not exist, using the default response")
@@ -77,7 +97,6 @@ def _resolve_file_path(self, wrapper, request_type):
                     response_file = os.path.join(
                         base_dir, request_type.name, f"default.json")
                 else:
-                    print("throttle")
                     wrapper.info("Throttling")
                 # Flip the flag
                 SHOULD_THROTTLE = not SHOULD_THROTTLE
@@ -483,27 +502,21 @@ async def handle_mock_s3_request(wrapper, request):
         # TODO: support more type.
         wrapper.info("unsupported request:", request)
         request_type = S3Opts.CreateMultipartUpload
-    global SHOULD_SKIP_WAIT
-    if "before_finish" not in parsed_path.path or request_type is not S3Opts.UploadPart:
-        # restore the state
-        SHOULD_SKIP_WAIT = True
-    if "before_finish" in parsed_path.path and SHOULD_SKIP_WAIT and request_type is S3Opts.UploadPart:
-        print("not wait for the request to complete")
-        SHOULD_SKIP_WAIT = False
-    else:
+
+    if response_config is None:
+        response_config = ResponseConfig(parsed_path.path, request)
+
+    if not response_config.should_skip_wait:
         while True:
             event = await wrapper.next_event()
             if type(event) is h11.EndOfMessage:
                 break
             assert type(event) is h11.Data
-
-    if response_config is None:
-        response_config = ResponseConfig(parsed_path.path)
-    response_config.request_headers = request.headers
+    else:
+        print("Skipping waiting for request body")
 
     response = response_config.resolve_response(
         wrapper, request_type, head_request=method == "HEAD")
-
     await send_response(wrapper, response)
 
 
diff --git a/tests/s3_mock_server_tests.c b/tests/s3_mock_server_tests.c
@@ -435,13 +435,31 @@ TEST_CASE(multipart_upload_with_network_interface_names_mock_server) {
 }
 
 /* Total hack to flip the bytes. */
-static void s_after_prepare_upload_part_finish(struct aws_s3_request *request) {
+static void s_after_prepare_upload_part_finish(struct aws_s3_request *request, struct aws_http_message *message) {
+    (void)message;
     if (request->num_times_prepared > 1) {
         /* mock that the body buffer was messed up in memory */
         request->request_body.buffer[1]++;
     }
 }
 
+static void s_after_prepare_upload_part_finish_retry_before_finish_sending(
+    struct aws_s3_request *request,
+    struct aws_http_message *message) {
+    if (request->num_times_prepared == 0 && message != NULL) {
+        struct aws_http_header before_finish_header = {
+            .name = aws_byte_cursor_from_c_str("before_finish"),
+            .value = aws_byte_cursor_from_c_str("true"),
+        };
+        aws_http_message_add_header(message, before_finish_header);
+        struct aws_http_header throttle_header = {
+            .name = aws_byte_cursor_from_c_str("force_throttle"),
+            .value = aws_byte_cursor_from_c_str("true"),
+        };
+        aws_http_message_add_header(message, throttle_header);
+    }
+}
+
 /**
  * This test is built for
  * 1. The retry happens before the upload has finished.
@@ -458,9 +476,10 @@ TEST_CASE(multipart_upload_checksum_with_retry_before_finish_mock_server) {
     struct aws_s3_client *client = NULL;
     ASSERT_SUCCESS(aws_s3_tester_client_new(&tester, &client_options, &client));
     struct aws_s3_client_vtable *patched_client_vtable = aws_s3_tester_patch_client_vtable(&tester, client, NULL);
-    patched_client_vtable->after_prepare_upload_part_finish = s_after_prepare_upload_part_finish;
+    patched_client_vtable->after_prepare_upload_part_finish =
+        s_after_prepare_upload_part_finish_retry_before_finish_sending;
 
-    struct aws_byte_cursor object_path = aws_byte_cursor_from_c_str("/throttle_before_finish");
+    struct aws_byte_cursor object_path = aws_byte_cursor_from_c_str("/throttle");
     {
         /* 1. Trailer checksum */
         struct aws_s3_tester_meta_request_options put_options = {
@@ -473,7 +492,6 @@ TEST_CASE(multipart_upload_checksum_with_retry_before_finish_mock_server) {
                 {
                     .object_size_mb = 10,
                     .object_path_override = object_path,
-                    .sleep_before_read_secs = 1,
                 },
             .mock_server = true,
         };
diff --git a/tests/s3_tester.c b/tests/s3_tester.c
@@ -1681,7 +1681,6 @@ int aws_s3_tester_send_meta_request_with_options(
                         .autogen_length = upload_size_bytes,
                         .eof_requires_extra_read = options->put_options.eof_requires_extra_read,
                         .max_bytes_per_read = options->put_options.max_bytes_per_read,
-                        .sleep_before_read_secs = options->put_options.sleep_before_read_secs,
                     },
             };
             if (options->put_options.invalid_input_stream) {
diff --git a/tests/s3_tester.h b/tests/s3_tester.h
@@ -229,7 +229,6 @@ struct aws_s3_tester_meta_request_options {
         struct aws_byte_cursor content_encoding;
         enum aws_s3_tester_full_object_checksum full_object_checksum;
         struct aws_byte_cursor if_none_match_header;
-        size_t sleep_before_read_secs;
     } put_options;
 
     enum aws_s3_tester_sse_type sse_type;

Original file line number	Diff line number	Diff line change
`@@ -113,7 +113,7 @@ static void s_aws_input_checksum_stream_destroy(struct aws_checksum_stream *impl`
`113`	`113`	`}`
`114`	`114`
`115`	`115`	`/* Compute the checksum of whatever was read, if we didn't reach the end of the underlying stream. */`
`116`		`- // s_finalize_checksum(impl);`
	`116`	`+ s_finalize_checksum(impl);`
`117`	`117`
`118`	`118`	`aws_checksum_destroy(impl->checksum);`
`119`	`119`	`aws_input_stream_release(impl->old_stream);`
Original file line number	Diff line number	Diff line change
`@@ -167,6 +167,8 @@ static void s_aws_input_chunk_stream_destroy(struct aws_chunk_stream *impl) {`
`167`	`167`	`}`
`168`	`168`	`aws_byte_buf_clean_up(&impl->pre_chunk_buffer);`
`169`	`169`	`aws_byte_buf_clean_up(&impl->post_chunk_buffer);`
	`170`	`+ /* Either we calculated the checksum, or we the checksum is empty. Otherwise, something was wrong. */`
	`171`	`+ AWS_ASSERT(impl->checksum_context->checksum_calculated \|\| impl->checksum_context->base64_checksum.len == 0);`
`170`	`172`	`aws_s3_upload_request_checksum_context_release(impl->checksum_context);`
`171`	`173`	`aws_mem_release(impl->allocator, impl);`
`172`	`174`	`}`