don't regenerate the checksum

TingDaoK · TingDaoK · commit 9d53af3d4439 · 2025-06-26T14:20:28.000-07:00
diff --git a/include/aws/s3/private/s3_checksums.h b/include/aws/s3/private/s3_checksums.h
@@ -93,15 +93,18 @@ struct aws_input_stream *aws_checksum_stream_new(
  * @param existing_stream   The data to be chunkified prepended by information on the stream length followed by a final
  *                          chunk and a trailing chunk containing a checksum of the existing stream. Destroying the
  *                          chunk stream will destroy the existing stream.
- * @param checksum_output   Optional argument, if provided the buffer will be initialized to the appropriate size and
+ * @param checksum_buffer   Required.
+ *                          - Empty buffer, the buffer will be initialized to the appropriate size and
  *                          filled with the checksum result when calculated. Callers responsibility to cleanup.
+ *                          - Otherwise, the buffer will be used directly.
+ *                          Caller takes the ownership of the buffer, error or not.
  */
 AWS_S3_API
 struct aws_input_stream *aws_chunk_stream_new(
     struct aws_allocator *allocator,
     struct aws_input_stream *existing_stream,
     enum aws_s3_checksum_algorithm algorithm,
-    struct aws_byte_buf *checksum_output);
+    struct aws_byte_buf *checksum_buffer);
 
 /**
  * Get the size of the checksum output corresponding to the aws_s3_checksum_algorithm enum value.
diff --git a/include/aws/s3/private/s3_client_impl.h b/include/aws/s3/private/s3_client_impl.h
@@ -174,6 +174,8 @@ struct aws_s3_client_vtable {
     struct aws_http_stream *(*http_connection_make_request)(
         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);
 };
 
 struct aws_s3_upload_part_timeout_stats {
diff --git a/include/aws/s3/private/s3_request_messages.h b/include/aws/s3/private/s3_request_messages.h
@@ -53,7 +53,7 @@ struct aws_input_stream *aws_s3_message_util_assign_body(
     struct aws_byte_buf *byte_buf,
     struct aws_http_message *out_message,
     const struct checksum_config_storage *checksum_config,
-    struct aws_byte_buf *out_checksum);
+    struct aws_byte_buf *checksum_buf);
 
 /* Create an HTTP request for an S3 Ranged Get Object Request, using the given request as a basis */
 AWS_S3_API
@@ -90,7 +90,7 @@ struct aws_http_message *aws_s3_upload_part_message_new(
     const struct aws_string *upload_id,
     bool should_compute_content_md5,
     const struct checksum_config_storage *checksum_config,
-    struct aws_byte_buf *encoded_checksum_output);
+    struct aws_byte_buf *encoded_checksum);
 
 /* Create an HTTP request for an S3 UploadPartCopy request, using the original request as a basis.
  * If multipart is not needed, part number and upload_id can be 0 and NULL,
diff --git a/source/s3_auto_ranged_put.c b/source/s3_auto_ranged_put.c
@@ -1115,6 +1115,7 @@ static void s_s3_prepare_upload_part_on_read_done(void *user_data) {
 static void s_s3_prepare_upload_part_finish(struct aws_s3_prepare_upload_part_job *part_prep, int error_code) {
     struct aws_s3_request *request = part_prep->request;
     struct aws_s3_meta_request *meta_request = request->meta_request;
+    struct aws_s3_client *client = meta_request->client;
     struct aws_s3_auto_ranged_put *auto_ranged_put = meta_request->impl;
 
     if (error_code != AWS_ERROR_SUCCESS) {
@@ -1141,8 +1142,10 @@ static void s_s3_prepare_upload_part_finish(struct aws_s3_prepare_upload_part_jo
             aws_array_list_get_at(&auto_ranged_put->synced_data.part_list, &part, request->part_number - 1);
             AWS_ASSERT(part != NULL);
             checksum_buf = &part->checksum_base64;
-            /* Clean up the buffer in case of it's initialized before and retry happens. */
-            aws_byte_buf_clean_up(checksum_buf);
+            /* If checksum buf is not empty, it means either the part being retried or the part resumed from list parts.
+             * Keep reusing the old checksum in case of the request body in memory mangled */
+            AWS_ASSERT(
+                checksum_buf->len == 0 || request->num_times_prepared > 0 || auto_ranged_put->resume_token != NULL);
             aws_s3_meta_request_unlock_synced_data(meta_request);
         }
         /* END CRITICAL SECTION */
@@ -1173,6 +1176,10 @@ static void s_s3_prepare_upload_part_finish(struct aws_s3_prepare_upload_part_jo
     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_chunk_stream.c b/source/s3_chunk_stream.c
@@ -22,14 +22,15 @@ typedef int(set_stream_fn)(struct aws_chunk_stream *parent_stream);
 struct aws_chunk_stream {
     struct aws_input_stream base;
     struct aws_allocator *allocator;
-
     /* aws_input_stream_byte_cursor provides our actual functionality  */
     /* Pointing to the stream we read from */
     struct aws_input_stream *current_stream;
 
-    struct aws_input_stream *checksum_stream;
-    struct aws_byte_buf checksum_result;
-    struct aws_byte_buf *checksum_result_output;
+    /* The passed-in buffer, owned by the caller. If passed-in buffer is empty
+     * it will be created by the chunk stream, but still caller owns its lifetime. Error or not. */
+    struct aws_byte_buf *checksum_buffer;
+
+    struct aws_input_stream *chunk_body_stream;
     struct aws_byte_buf pre_chunk_buffer;
     struct aws_byte_buf post_chunk_buffer;
     struct aws_byte_cursor checksum_header_name;
@@ -63,16 +64,12 @@ static int s_set_post_chunk_stream(struct aws_chunk_stream *parent_stream) {
     struct aws_byte_cursor post_trailer_cursor = aws_byte_cursor_from_string(s_post_trailer);
     struct aws_byte_cursor colon_cursor = aws_byte_cursor_from_string(s_colon);
 
-    if (parent_stream->checksum_result.len == 0) {
+    if (parent_stream->checksum_buffer->len == 0) {
         AWS_LOGF_ERROR(AWS_LS_S3_META_REQUEST, "Failed to extract base64 encoded checksum of stream");
         return aws_raise_error(AWS_ERROR_S3_CHECKSUM_CALCULATION_FAILED);
     }
-    struct aws_byte_cursor checksum_result_cursor = aws_byte_cursor_from_buf(&parent_stream->checksum_result);
-    if (parent_stream->checksum_result_output &&
-        aws_byte_buf_init_copy_from_cursor(
-            parent_stream->checksum_result_output, parent_stream->allocator, checksum_result_cursor)) {
-        return AWS_OP_ERR;
-    }
+    struct aws_byte_cursor checksum_result_cursor = aws_byte_cursor_from_buf(parent_stream->checksum_buffer);
+
     if (aws_byte_buf_init(
             &parent_stream->post_chunk_buffer,
             parent_stream->allocator,
@@ -92,16 +89,15 @@ static int s_set_post_chunk_stream(struct aws_chunk_stream *parent_stream) {
     parent_stream->set_current_stream_fn = s_set_null_stream;
     return AWS_OP_SUCCESS;
 error:
-    aws_byte_buf_clean_up(parent_stream->checksum_result_output);
     aws_byte_buf_clean_up(&parent_stream->post_chunk_buffer);
     return AWS_OP_ERR;
 }
 
 static int s_set_chunk_stream(struct aws_chunk_stream *parent_stream) {
     aws_input_stream_release(parent_stream->current_stream);
-    parent_stream->current_stream = parent_stream->checksum_stream;
+    parent_stream->current_stream = parent_stream->chunk_body_stream;
     aws_byte_buf_clean_up(&parent_stream->pre_chunk_buffer);
-    parent_stream->checksum_stream = NULL;
+    parent_stream->chunk_body_stream = NULL;
     parent_stream->set_current_stream_fn = s_set_post_chunk_stream;
     return AWS_OP_SUCCESS;
 }
@@ -169,11 +165,10 @@ static void s_aws_input_chunk_stream_destroy(struct aws_chunk_stream *impl) {
         if (impl->current_stream) {
             aws_input_stream_release(impl->current_stream);
         }
-        if (impl->checksum_stream) {
-            aws_input_stream_release(impl->checksum_stream);
+        if (impl->chunk_body_stream) {
+            aws_input_stream_release(impl->chunk_body_stream);
         }
         aws_byte_buf_clean_up(&impl->pre_chunk_buffer);
-        aws_byte_buf_clean_up(&impl->checksum_result);
         aws_byte_buf_clean_up(&impl->post_chunk_buffer);
         aws_mem_release(impl->allocator, impl);
     }
@@ -190,13 +185,17 @@ struct aws_input_stream *aws_chunk_stream_new(
     struct aws_allocator *allocator,
     struct aws_input_stream *existing_stream,
     enum aws_s3_checksum_algorithm algorithm,
-    struct aws_byte_buf *checksum_output) {
+    struct aws_byte_buf *checksum_buffer) {
+    AWS_PRECONDITION(allocator);
+    AWS_PRECONDITION(existing_stream);
+    AWS_PRECONDITION(checksum_buffer);
 
     struct aws_chunk_stream *impl = aws_mem_calloc(allocator, 1, sizeof(struct aws_chunk_stream));
 
     impl->allocator = allocator;
     impl->base.vtable = &s_aws_input_chunk_stream_vtable;
-    impl->checksum_result_output = checksum_output;
+    impl->checksum_buffer = checksum_buffer;
+
     int64_t stream_length = 0;
     int64_t final_chunk_len = 0;
     if (aws_input_stream_get_length(existing_stream, &stream_length)) {
@@ -225,15 +224,29 @@ struct aws_input_stream *aws_chunk_stream_new(
     if (aws_base64_compute_encoded_len(checksum_len, &encoded_checksum_len)) {
         goto error;
     }
-    if (aws_byte_buf_init(&impl->checksum_result, allocator, encoded_checksum_len)) {
-        goto error;
-    }
-
-    impl->checksum_stream = aws_checksum_stream_new(allocator, existing_stream, algorithm, &impl->checksum_result);
-    if (impl->checksum_stream == NULL) {
-        goto error;
+    if (checksum_buffer->len == 0) {
+        /* Empty passed-in buffer, calculate the checksum during reading from the stream. */
+        if (aws_byte_buf_init(impl->checksum_buffer, allocator, encoded_checksum_len)) {
+            goto error;
+        }
+        /* Wrap the existing stream with checksum stream to calculate the checksum when reading from it. */
+        impl->chunk_body_stream = aws_checksum_stream_new(allocator, existing_stream, algorithm, impl->checksum_buffer);
+        if (impl->chunk_body_stream == NULL) {
+            goto error;
+        }
+    } else {
+        if (checksum_buffer->len != encoded_checksum_len) {
+            AWS_LOGF_ERROR(
+                AWS_LS_S3_GENERAL,
+                "Mismatched checksum buffer and algorithm. checksum_buf->len is %zu, but encoded_checksum_len is %zu",
+                checksum_buffer->len,
+                encoded_checksum_len);
+            aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
+            goto error;
+        }
+        /* No need to calculate the checksum during read, use the existing stream directly. */
+        impl->chunk_body_stream = aws_input_stream_acquire(existing_stream);
     }
-
     int64_t prechunk_stream_len = 0;
     int64_t colon_len = s_colon->len;
     int64_t post_trailer_len = s_post_trailer->len;
@@ -248,9 +261,9 @@ struct aws_input_stream *aws_chunk_stream_new(
         }
         impl->set_current_stream_fn = s_set_chunk_stream;
     } else {
-        impl->current_stream = impl->checksum_stream;
+        impl->current_stream = impl->chunk_body_stream;
         final_chunk_len = s_empty_chunk->len;
-        impl->checksum_stream = NULL;
+        impl->chunk_body_stream = NULL;
         impl->set_current_stream_fn = s_set_post_chunk_stream;
     }
 
@@ -268,10 +281,9 @@ struct aws_input_stream *aws_chunk_stream_new(
     return &impl->base;
 
 error:
-    aws_input_stream_release(impl->checksum_stream);
+    aws_input_stream_release(impl->chunk_body_stream);
     aws_input_stream_release(impl->current_stream);
     aws_byte_buf_clean_up(&impl->pre_chunk_buffer);
-    aws_byte_buf_clean_up(&impl->checksum_result);
     aws_mem_release(impl->allocator, impl);
     return NULL;
 }
diff --git a/source/s3_request_messages.c b/source/s3_request_messages.c
@@ -352,8 +352,10 @@ struct aws_http_message *aws_s3_create_multipart_upload_message_new(
     return NULL;
 }
 
-/* Create a new put object request from an existing put object request.  Currently just optionally adds part information
- * for a multipart upload. */
+/**
+ * Create a new put object request from an existing put object request.  Currently just optionally adds part information
+ * for a multipart upload.
+ **/
 struct aws_http_message *aws_s3_upload_part_message_new(
     struct aws_allocator *allocator,
     struct aws_http_message *base_message,
@@ -362,7 +364,7 @@ struct aws_http_message *aws_s3_upload_part_message_new(
     const struct aws_string *upload_id,
     bool should_compute_content_md5,
     const struct checksum_config_storage *checksum_config,
-    struct aws_byte_buf *encoded_checksum_output) {
+    struct aws_byte_buf *encoded_checksum) {
     AWS_PRECONDITION(allocator);
     AWS_PRECONDITION(base_message);
     AWS_PRECONDITION(part_number > 0);
@@ -383,7 +385,7 @@ struct aws_http_message *aws_s3_upload_part_message_new(
         goto error_clean_up;
     }
 
-    if (aws_s3_message_util_assign_body(allocator, buffer, message, checksum_config, encoded_checksum_output) == NULL) {
+    if (aws_s3_message_util_assign_body(allocator, buffer, message, checksum_config, encoded_checksum) == NULL) {
         goto error_clean_up;
     }
 
@@ -836,19 +838,30 @@ static int s_calculate_in_memory_checksum_helper(
     struct aws_allocator *allocator,
     struct aws_byte_cursor data,
     const struct checksum_config_storage *checksum_config,
-    struct aws_byte_buf *out_checksum) {
+    struct aws_byte_buf *checksum_buf) {
     AWS_ASSERT(checksum_config->checksum_algorithm != AWS_SCA_NONE);
-    AWS_ASSERT(out_checksum != NULL);
-    AWS_ZERO_STRUCT(*out_checksum);
+    AWS_ASSERT(checksum_buf != NULL);
 
     int ret_code = AWS_OP_ERR;
     size_t digest_size = aws_get_digest_size_from_checksum_algorithm(checksum_config->checksum_algorithm);
     size_t encoded_checksum_len = 0;
     if (aws_base64_compute_encoded_len(digest_size, &encoded_checksum_len)) {
         return AWS_OP_ERR;
     }
-
-    aws_byte_buf_init(out_checksum, allocator, encoded_checksum_len);
+    if (checksum_buf->len > 0) {
+        if (checksum_buf->len == encoded_checksum_len) {
+            return AWS_OP_SUCCESS;
+        } else {
+            AWS_LOGF_ERROR(
+                AWS_LS_S3_GENERAL,
+                "Mismatched checksum buffer and algorithm. checksum_buf->len is %zu, but encoded_checksum_len is %zu",
+                checksum_buf->len,
+                encoded_checksum_len);
+            aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
+            return AWS_OP_ERR;
+        }
+    }
+    aws_byte_buf_init(checksum_buf, allocator, encoded_checksum_len);
 
     struct aws_byte_buf raw_checksum;
     aws_byte_buf_init(&raw_checksum, allocator, digest_size);
@@ -857,14 +870,14 @@ static int s_calculate_in_memory_checksum_helper(
         goto done;
     }
     struct aws_byte_cursor raw_checksum_cursor = aws_byte_cursor_from_buf(&raw_checksum);
-    if (aws_base64_encode(&raw_checksum_cursor, out_checksum)) {
+    if (aws_base64_encode(&raw_checksum_cursor, checksum_buf)) {
         goto done;
     }
 
     ret_code = AWS_OP_SUCCESS;
 done:
     if (ret_code) {
-        aws_byte_buf_clean_up(out_checksum);
+        aws_byte_buf_clean_up(checksum_buf);
     }
     aws_byte_buf_clean_up(&raw_checksum);
     return ret_code;
@@ -880,19 +893,19 @@ static int s_calculate_and_add_checksum_to_header_helper(
     struct aws_byte_cursor data,
     const struct checksum_config_storage *checksum_config,
     struct aws_http_message *out_message,
-    struct aws_byte_buf *out_checksum) {
+    struct aws_byte_buf *checksum_buf) {
     AWS_ASSERT(checksum_config->checksum_algorithm != AWS_SCA_NONE);
     AWS_ASSERT(out_message != NULL);
     int ret_code = AWS_OP_ERR;
 
     struct aws_byte_buf local_encoded_checksum_buf;
+    AWS_ZERO_STRUCT(local_encoded_checksum_buf);
     struct aws_byte_buf *local_encoded_checksum;
-    if (out_checksum == NULL) {
+    if (checksum_buf == NULL) {
         local_encoded_checksum = &local_encoded_checksum_buf;
     } else {
-        local_encoded_checksum = out_checksum;
+        local_encoded_checksum = checksum_buf;
     }
-    AWS_ZERO_STRUCT(*local_encoded_checksum);
     if (s_calculate_in_memory_checksum_helper(allocator, data, checksum_config, local_encoded_checksum)) {
         goto done;
     }
@@ -908,8 +921,8 @@ static int s_calculate_and_add_checksum_to_header_helper(
 
     ret_code = AWS_OP_SUCCESS;
 done:
-    if (ret_code || out_checksum == NULL) {
-        /* In case of error happen or out_checksum is not set, clean up the encoded checksum. Otherwise, the caller will
+    if (ret_code || checksum_buf == NULL) {
+        /* In case of error happen or checksum_buf is not set, clean up the encoded checksum. Otherwise, the caller will
          * own the encoded checksum. */
         aws_byte_buf_clean_up(local_encoded_checksum);
     }
@@ -922,7 +935,7 @@ struct aws_input_stream *aws_s3_message_util_assign_body(
     struct aws_byte_buf *byte_buf,
     struct aws_http_message *out_message,
     const struct checksum_config_storage *checksum_config,
-    struct aws_byte_buf *out_checksum) {
+    struct aws_byte_buf *checksum_buf) {
     AWS_PRECONDITION(allocator);
     AWS_PRECONDITION(out_message);
     AWS_PRECONDITION(byte_buf);
@@ -992,7 +1005,7 @@ struct aws_input_stream *aws_s3_message_util_assign_body(
             }
             /* set input stream to chunk stream */
             struct aws_input_stream *chunk_stream =
-                aws_chunk_stream_new(allocator, input_stream, checksum_config->checksum_algorithm, out_checksum);
+                aws_chunk_stream_new(allocator, input_stream, checksum_config->checksum_algorithm, checksum_buf);
             if (!chunk_stream) {
                 goto error_clean_up;
             }
@@ -1001,14 +1014,14 @@ struct aws_input_stream *aws_s3_message_util_assign_body(
         } else if (checksum_config->location == AWS_SCL_HEADER) {
             /* Calculate the checksum directly from memory and add it to the header. */
             if (s_calculate_and_add_checksum_to_header_helper(
-                    allocator, buffer_byte_cursor, checksum_config, out_message, out_checksum)) {
+                    allocator, buffer_byte_cursor, checksum_config, out_message, checksum_buf)) {
                 goto error_clean_up;
             }
 
-        } else if (checksum_config->checksum_algorithm != AWS_SCA_NONE && out_checksum != NULL) {
-            /* In case checksums still wanted, and we can calculate it directly from the buffer in memory to
-             * out_checksum */
-            if (s_calculate_in_memory_checksum_helper(allocator, buffer_byte_cursor, checksum_config, out_checksum)) {
+        } else if (checksum_config->checksum_algorithm != AWS_SCA_NONE && checksum_buf != NULL) {
+            /* In case checksums still wanted, but not sent to the server, and we can calculate it directly from the
+             * buffer in memory to out_checksum */
+            if (s_calculate_in_memory_checksum_helper(allocator, buffer_byte_cursor, checksum_config, checksum_buf)) {
                 goto error_clean_up;
             }
         }
diff --git a/tests/s3_mock_server_tests.c b/tests/s3_mock_server_tests.c
