add lock and don't direct access the synced data

TingDaoK · TingDaoK · commit cbe0c9e7852e · 2025-08-11T11:46:00.000-07:00
diff --git a/include/aws/s3/private/s3_checksum_context.h b/include/aws/s3/private/s3_checksum_context.h
@@ -8,6 +8,7 @@
 
 #include "aws/s3/s3_client.h"
 #include <aws/common/byte_buf.h>
+#include <aws/common/mutex.h>
 #include <aws/common/ref_count.h>
 
 struct aws_s3_meta_request_checksum_config_storage;
@@ -28,9 +29,13 @@ struct aws_s3_upload_request_checksum_context {
     enum aws_s3_checksum_algorithm algorithm;
     enum aws_s3_checksum_location location;
 
-    struct aws_byte_buf base64_checksum;
-    /* The checksum already be calculated or not. */
-    bool checksum_calculated;
+    struct {
+        /* Lock to make sure the checksum context is safe to be access from different threads. */
+        struct aws_mutex lock;
+        struct aws_byte_buf base64_checksum;
+        /* The checksum already be calculated or not. */
+        bool checksum_calculated;
+    } synced_data;
 
     /* Validation */
     size_t encoded_checksum_size;
@@ -96,8 +101,7 @@ struct aws_s3_upload_request_checksum_context *aws_s3_upload_request_checksum_co
  * @return true if checksum calculation is needed, false otherwise
  */
 AWS_S3_API
-bool aws_s3_upload_request_checksum_context_should_calculate(
-    const struct aws_s3_upload_request_checksum_context *context);
+bool aws_s3_upload_request_checksum_context_should_calculate(struct aws_s3_upload_request_checksum_context *context);
 
 /**
  * Check if checksum should be added to HTTP headers.
@@ -122,15 +126,18 @@ bool aws_s3_upload_request_checksum_context_should_add_trailer(
     const struct aws_s3_upload_request_checksum_context *context);
 
 /**
- * Get the checksum buffer to use for output.
- * Returns the internal buffer for storing the calculated checksum.
+ * Encode the checksum to base64 and store it in the context.
+ * This function is thread-safe and can be called from multiple threads.
+ * Returns AWS_OP_SUCCESS on success, AWS_OP_ERR otherwise
  *
  * @param context The checksum context
- * @return Pointer to the checksum buffer, or NULL if context is invalid
+ * @param raw_checksum_cursor the byte cursor to the raw checksum value.
+ * @return AWS_OP_SUCCESS on success, AWS_OP_ERR otherwise
  */
 AWS_S3_API
-struct aws_byte_buf *aws_s3_upload_request_checksum_context_get_output_buffer(
-    struct aws_s3_upload_request_checksum_context *context);
+int aws_s3_upload_request_checksum_context_finalize_checksum(
+    struct aws_s3_upload_request_checksum_context *context,
+    struct aws_byte_cursor raw_checksum_cursor);
 
 /**
  * Get a cursor to the current base64 encoded checksum value (for use in headers/trailers).
@@ -141,7 +148,7 @@ struct aws_byte_buf *aws_s3_upload_request_checksum_context_get_output_buffer(
  */
 AWS_S3_API
 struct aws_byte_cursor aws_s3_upload_request_checksum_context_get_checksum_cursor(
-    const struct aws_s3_upload_request_checksum_context *context);
+    struct aws_s3_upload_request_checksum_context *context);
 
 AWS_EXTERN_C_END
 
diff --git a/include/aws/s3/private/s3_checksums.h b/include/aws/s3/private/s3_checksums.h
@@ -85,6 +85,15 @@ struct aws_input_stream *aws_checksum_stream_new(
 AWS_S3_API
 int aws_checksum_stream_finalize_checksum(struct aws_input_stream *checksum_stream, struct aws_byte_buf *checksum_buf);
 
+/**
+ * Finalize the checksum has read so far to the checksum context.
+ * Not thread safe.
+ */
+AWS_S3_API
+int aws_checksum_stream_finalize_checksum_context(
+    struct aws_input_stream *checksum_stream,
+    struct aws_s3_upload_request_checksum_context *checksum_context);
+
 /**
  * TODO: properly support chunked encoding.
  * Creates a chunked encoding stream that wraps an existing stream and adds checksum trailers.
diff --git a/source/s3_auto_ranged_put.c b/source/s3_auto_ranged_put.c
@@ -1099,15 +1099,20 @@ static void s_s3_prepare_upload_part_on_read_done(void *user_data) {
             goto on_done;
         }
         struct aws_s3_upload_request_checksum_context *context = previously_uploaded_info->checksum_context;
-        /* if previously uploaded part had a checksum, compare it to what we just skipped */
-        if (context != NULL && context->checksum_calculated == true &&
-            s_verify_part_matches_checksum(
-                meta_request->allocator,
-                aws_byte_cursor_from_buf(&request->request_body),
-                meta_request->checksum_config.checksum_algorithm,
-                aws_byte_cursor_from_buf(&context->base64_checksum))) {
-            error_code = aws_last_error_or_unknown();
-            goto on_done;
+        if (context) {
+            if (!aws_s3_upload_request_checksum_context_should_calculate(context)) {
+                struct aws_byte_cursor previous_calculated_checksum =
+                    aws_s3_upload_request_checksum_context_get_checksum_cursor(context);
+                /* if previously uploaded part had a checksum, compare it to what we just skipped */
+                if (s_verify_part_matches_checksum(
+                        meta_request->allocator,
+                        aws_byte_cursor_from_buf(&request->request_body),
+                        meta_request->checksum_config.checksum_algorithm,
+                        previous_calculated_checksum) != AWS_OP_SUCCESS) {
+                    error_code = aws_last_error_or_unknown();
+                    goto on_done;
+                }
+            }
         }
     }
 
@@ -1157,9 +1162,6 @@ static void s_s3_prepare_upload_part_finish(struct aws_s3_prepare_upload_part_jo
             checksum_context = part->checksum_context;
             /* If checksum already calculated, 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_context->checksum_calculated || request->num_times_prepared > 0 ||
-                auto_ranged_put->resume_token != NULL);
             aws_s3_meta_request_unlock_synced_data(meta_request);
         }
         /* END CRITICAL SECTION */
diff --git a/source/s3_checksum_context.c b/source/s3_checksum_context.c
@@ -8,9 +8,18 @@
 #include <aws/common/encoding.h>
 #include <aws/common/logging.h>
 
+static void s_lock_synced_data(struct aws_s3_upload_request_checksum_context *context) {
+    aws_mutex_lock(&context->synced_data.lock);
+}
+
+static void s_unlock_synced_data(struct aws_s3_upload_request_checksum_context *context) {
+    aws_mutex_unlock(&context->synced_data.lock);
+}
+
 static void s_aws_s3_upload_request_checksum_context_destroy(void *context) {
     struct aws_s3_upload_request_checksum_context *checksum_context = context;
-    aws_byte_buf_clean_up(&checksum_context->base64_checksum);
+    aws_byte_buf_clean_up(&checksum_context->synced_data.base64_checksum);
+    aws_mutex_clean_up(&checksum_context->synced_data.lock);
     aws_mem_release(checksum_context->allocator, checksum_context);
 }
 
@@ -44,6 +53,10 @@ static struct aws_s3_upload_request_checksum_context *s_s3_upload_request_checks
         return NULL;
     }
     context->encoded_checksum_size = encoded_size;
+    if (!aws_mutex_init(&context->synced_data.lock)) {
+        aws_s3_upload_request_checksum_context_release(context);
+        return NULL;
+    }
     return context;
 }
 
@@ -54,7 +67,7 @@ struct aws_s3_upload_request_checksum_context *aws_s3_upload_request_checksum_co
         s_s3_upload_request_checksum_context_new_base(allocator, checksum_config);
     if (context && context->encoded_checksum_size > 0) {
         /* Initial the buffer for checksum */
-        aws_byte_buf_init(&context->base64_checksum, allocator, context->encoded_checksum_size);
+        aws_byte_buf_init(&context->synced_data.base64_checksum, allocator, context->encoded_checksum_size);
     }
     return context;
 }
@@ -79,8 +92,8 @@ struct aws_s3_upload_request_checksum_context *aws_s3_upload_request_checksum_co
             aws_s3_upload_request_checksum_context_release(context);
             return NULL;
         }
-        aws_byte_buf_init_copy_from_cursor(&context->base64_checksum, allocator, existing_base64_checksum);
-        context->checksum_calculated = true;
+        aws_byte_buf_init_copy_from_cursor(&context->synced_data.base64_checksum, allocator, existing_base64_checksum);
+        context->synced_data.checksum_calculated = true;
     }
     return context;
 }
@@ -101,14 +114,18 @@ struct aws_s3_upload_request_checksum_context *aws_s3_upload_request_checksum_co
     return NULL;
 }
 
-bool aws_s3_upload_request_checksum_context_should_calculate(
-    const struct aws_s3_upload_request_checksum_context *context) {
+bool aws_s3_upload_request_checksum_context_should_calculate(struct aws_s3_upload_request_checksum_context *context) {
     if (!context || context->algorithm == AWS_SCA_NONE) {
         return false;
     }
 
+    bool should_calculate = false;
+    s_lock_synced_data(context);
     /* If not previous calculated */
-    return !context->checksum_calculated;
+    should_calculate = !context->synced_data.checksum_calculated;
+    s_unlock_synced_data(context);
+
+    return should_calculate;
 }
 
 bool aws_s3_upload_request_checksum_context_should_add_header(
@@ -129,19 +146,46 @@ bool aws_s3_upload_request_checksum_context_should_add_trailer(
     return context->location == AWS_SCL_TRAILER && context->algorithm != AWS_SCA_NONE;
 }
 
-struct aws_byte_buf *aws_s3_upload_request_checksum_context_get_output_buffer(
-    struct aws_s3_upload_request_checksum_context *context) {
+int aws_s3_upload_request_checksum_context_finalize_checksum(
+    struct aws_s3_upload_request_checksum_context *context,
+    struct aws_byte_cursor raw_checksum_cursor) {
     if (!context) {
-        return NULL;
+        return aws_raise_error(AWS_ERROR_INVALID_ARGUMENT);
     }
-    return &context->base64_checksum;
+    s_lock_synced_data(context);
+    /* If not previous calculated */
+    if (!context->synced_data.checksum_calculated) {
+        AWS_ASSERT(context->synced_data.base64_checksum.len == 0);
+
+        if (aws_base64_encode(&raw_checksum_cursor, &context->synced_data.base64_checksum)) {
+            aws_byte_buf_reset(&context->synced_data.base64_checksum, false);
+            AWS_LOGF_ERROR(
+                AWS_LS_S3_GENERAL,
+                "Failed to base64 encode for the checksum. Raw checksum length: %zu. Output buffer capacity: %zu "
+                "length %zu",
+                raw_checksum_cursor.len,
+                context->synced_data.base64_checksum.capacity,
+                context->synced_data.base64_checksum.len);
+            s_unlock_synced_data(context);
+            return AWS_OP_ERR;
+        }
+        context->synced_data.checksum_calculated = true;
+    }
+    s_unlock_synced_data(context);
+    return AWS_OP_SUCCESS;
 }
 
 struct aws_byte_cursor aws_s3_upload_request_checksum_context_get_checksum_cursor(
-    const struct aws_s3_upload_request_checksum_context *context) {
+    struct aws_s3_upload_request_checksum_context *context) {
     struct aws_byte_cursor checksum_cursor = {0};
-    if (!context || !context->checksum_calculated) {
+    if (!context) {
         return checksum_cursor;
     }
-    return aws_byte_cursor_from_buf(&context->base64_checksum);
+    s_lock_synced_data(context);
+    /* If not previous calculated */
+    if (context->synced_data.checksum_calculated) {
+        checksum_cursor = aws_byte_cursor_from_buf(&context->synced_data.base64_checksum);
+    }
+    s_unlock_synced_data(context);
+    return checksum_cursor;
 }
diff --git a/source/s3_checksum_stream.c b/source/s3_checksum_stream.c
@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0.
  */
 
+#include "aws/s3/private/s3_checksum_context.h"
 #include "aws/s3/private/s3_checksums.h"
 #include <aws/common/encoding.h>
 #include <aws/io/stream.h>
@@ -135,3 +136,31 @@ int aws_checksum_stream_finalize_checksum(struct aws_input_stream *checksum_stre
 
     return AWS_OP_SUCCESS;
 }
+
+int aws_checksum_stream_finalize_checksum_context(
+    struct aws_input_stream *checksum_stream,
+    struct aws_s3_upload_request_checksum_context *checksum_context) {
+    struct aws_checksum_stream *impl = AWS_CONTAINER_OF(checksum_stream, struct aws_checksum_stream, base);
+
+    if (aws_checksum_finalize(impl->checksum, &impl->checksum_result) != AWS_OP_SUCCESS) {
+        AWS_LOGF_ERROR(
+            AWS_LS_S3_CLIENT,
+            "Failed to calculate checksum with error code %d (%s).",
+            aws_last_error(),
+            aws_error_str(aws_last_error()));
+        aws_byte_buf_reset(&impl->checksum_result, true);
+        return aws_raise_error(AWS_ERROR_S3_CHECKSUM_CALCULATION_FAILED);
+    }
+    struct aws_byte_cursor checksum_result_cursor = aws_byte_cursor_from_buf(&impl->checksum_result);
+    if (aws_s3_upload_request_checksum_context_finalize_checksum(checksum_context, checksum_result_cursor) !=
+        AWS_OP_SUCCESS) {
+        AWS_LOGF_ERROR(
+            AWS_LS_S3_CLIENT,
+            "Failed to finalize checksum context with error code %d (%s).",
+            aws_last_error(),
+            aws_error_str(aws_last_error()));
+        aws_byte_buf_reset(&impl->checksum_result, true);
+        return aws_raise_error(AWS_ERROR_S3_CHECKSUM_CALCULATION_FAILED);
+    }
+    return AWS_OP_SUCCESS;
+}
diff --git a/source/s3_chunk_stream.c b/source/s3_chunk_stream.c
@@ -62,19 +62,15 @@ 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);
-    /* After the checksum stream released, the checksum will be calculated. */
-    if (parent_stream->checksum_context->checksum_calculated == false) {
-        AWS_ASSERT(parent_stream->checksum_context->base64_checksum.len == 0);
-        AWS_ASSERT(parent_stream->checksum_stream != NULL);
-        /* Get the checksum from the checksum stream of the data read from it. */
-        if (aws_checksum_stream_finalize_checksum(
-                parent_stream->checksum_stream, &parent_stream->checksum_context->base64_checksum)) {
+    if (parent_stream->checksum_stream) {
+        /* If we have the checksum stream, finalize the checksum now as we finished reading from it. */
+        if (aws_checksum_stream_finalize_checksum_context(
+                parent_stream->checksum_stream, parent_stream->checksum_context)) {
             return AWS_OP_ERR;
         }
-        parent_stream->checksum_context->checksum_calculated = true;
     }
     struct aws_byte_cursor checksum_result_cursor =
-        aws_byte_cursor_from_buf(&parent_stream->checksum_context->base64_checksum);
+        aws_s3_upload_request_checksum_context_get_checksum_cursor(&parent_stream->checksum_context);
     AWS_ASSERT(parent_stream->checksum_context->encoded_checksum_size == checksum_result_cursor.len);
 
     if (aws_byte_buf_init(
@@ -175,7 +171,6 @@ 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);
     }
@@ -207,7 +202,7 @@ struct aws_input_stream *aws_chunk_stream_new(
     impl->checksum_context = aws_s3_upload_request_checksum_context_acquire(checksum_context);
 
     algorithm = checksum_context->algorithm;
-    bool checksum_calculated = checksum_context->checksum_calculated;
+    bool should_calculate_checksum = aws_s3_upload_request_checksum_context_should_calculate(impl->checksum_context);
 
     int64_t stream_length = 0;
     int64_t final_chunk_len = 0;
@@ -230,7 +225,7 @@ struct aws_input_stream *aws_chunk_stream_new(
     if (aws_byte_buf_append(&impl->pre_chunk_buffer, &pre_chunk_cursor)) {
         goto error;
     }
-    if (!checksum_calculated) {
+    if (should_calculate_checksum) {
         /* Wrap the existing stream with checksum stream to calculate the checksum when reading from it. */
         impl->checksum_stream = aws_checksum_stream_new(allocator, existing_stream, algorithm);
         if (impl->checksum_stream == NULL) {
diff --git a/source/s3_request_messages.c b/source/s3_request_messages.c
@@ -838,13 +838,9 @@ static int s_calculate_in_memory_checksum_helper(
     struct aws_byte_cursor data,
     struct aws_s3_upload_request_checksum_context *checksum_context) {
     AWS_ASSERT(checksum_context);
-    if (checksum_context->checksum_calculated) {
-        return AWS_OP_SUCCESS;
-    }
 
     int ret_code = AWS_OP_ERR;
     /* Calculate checksum for output buffer only (no header/trailer) */
-    struct aws_byte_buf *output_buffer = aws_s3_upload_request_checksum_context_get_output_buffer(checksum_context);
     struct aws_byte_buf raw_checksum;
     size_t digest_size = aws_get_digest_size_from_checksum_algorithm(checksum_context->algorithm);
     aws_byte_buf_init(&raw_checksum, allocator, digest_size);
@@ -855,18 +851,14 @@ static int s_calculate_in_memory_checksum_helper(
     }
 
     struct aws_byte_cursor raw_checksum_cursor = aws_byte_cursor_from_buf(&raw_checksum);
-    if (aws_base64_encode(&raw_checksum_cursor, output_buffer)) {
+    if (aws_s3_upload_request_checksum_context_finalize_checksum(checksum_context, raw_checksum_cursor)) {
         aws_byte_buf_clean_up(&raw_checksum);
         goto done;
     }
     aws_byte_buf_clean_up(&raw_checksum);
-    checksum_context->checksum_calculated = true;
 
     ret_code = AWS_OP_SUCCESS;
 done:
-    if (ret_code) {
-        aws_byte_buf_clean_up(output_buffer);
-    }
     aws_byte_buf_clean_up(&raw_checksum);
     return ret_code;
 }
@@ -890,7 +882,8 @@ static int s_calculate_and_add_checksum_to_header_helper(
     /* Add the encoded checksum to header. */
     const struct aws_byte_cursor header_name =
         aws_get_http_header_name_from_checksum_algorithm(checksum_context->algorithm);
-    struct aws_byte_cursor encoded_checksum_val = aws_byte_cursor_from_buf(&checksum_context->base64_checksum);
+    struct aws_byte_cursor encoded_checksum_val =
+        aws_s3_upload_request_checksum_context_get_checksum_cursor(checksum_context);
     struct aws_http_headers *headers = aws_http_message_get_headers(out_message);
     if (aws_http_headers_set(headers, header_name, encoded_checksum_val)) {
         goto done;
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -294,7 +294,6 @@ add_test_case(crc32c_test_invalid_buffer)
 add_test_case(crc32c_test_oneshot)
 add_test_case(crc32c_test_invalid_state)
 
-add_test_case(test_upload_request_checksum_context_get_output_buffer)
 add_test_case(test_upload_request_checksum_context_get_checksum_cursor)
 add_test_case(test_upload_request_checksum_context_error_cases)
 add_test_case(test_upload_request_checksum_context_different_algorithms)
diff --git a/tests/s3_checksum_context_test.c b/tests/s3_checksum_context_test.c
