Revamp checksum - retry will reuse the checksum (#532)

Co-authored-by: Dmitriy Musatkin <63878209+DmitriyMusatkin@users.noreply.github.com>

Author: TingDaoK

include/aws/s3/private/s3_checksum_context.h

@@ -0,0 +1,148 @@
+#ifndef AWS_S3_CHECKSUM_CONTEXT_H
+#define AWS_S3_CHECKSUM_CONTEXT_H
+
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "aws/s3/s3_client.h"
+#include <aws/common/byte_buf.h>
+#include <aws/common/ref_count.h>
+
+struct aws_s3_meta_request_checksum_config_storage;
+
+AWS_EXTERN_C_BEGIN
+
+/**
+ * Upload request checksum context that encapsulates all checksum-related state and behavior
+ * for individual upload part requests. This replaces the complex tri-state buffer logic
+ * with a cleaner approach. Uses reference counting for lifetime management since context
+ * is transferred between functions.
+ */
+struct aws_s3_upload_request_checksum_context {
+    struct aws_allocator *allocator;
+    struct aws_ref_count ref_count;
+
+    /* Configuration */
+    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;
+
+    /* Validation */
+    size_t encoded_checksum_size;
+};
+
+/**
+ * Create a new upload request checksum context from configuration and buffer parameters.
+ * This function encapsulates all the complex logic for determining buffer state.
+ * Returns with reference count of 1.
+ *
+ * @param allocator Memory allocator
+ * @param checksum_config Meta request level checksum configuration (can be NULL)
+ * @return New checksum context or NULL on error
+ */
+AWS_S3_API
+struct aws_s3_upload_request_checksum_context *aws_s3_upload_request_checksum_context_new(
+    struct aws_allocator *allocator,
+    const struct aws_s3_meta_request_checksum_config_storage *checksum_config);
+
+/**
+ * Create a new upload request checksum context with an existing base64 encoded checksum value.
+ * This is useful when resuming uploads or when the checksum is pre-calculated.
+ * Returns with reference count of 1.
+ *
+ * @param allocator Memory allocator
+ * @param checksum_config Meta request level checksum configuration (can be NULL)
+ * @param existing_base64_checksum Pre-calculated checksum value as a byte cursor
+ * @return New checksum context or NULL on error (e.g., if checksum size doesn't match algorithm)
+ */
+AWS_S3_API
+struct aws_s3_upload_request_checksum_context *aws_s3_upload_request_checksum_context_new_with_existing_base64_checksum(
+    struct aws_allocator *allocator,
+    const struct aws_s3_meta_request_checksum_config_storage *checksum_config,
+    struct aws_byte_cursor existing_base64_checksum);
+
+/**
+ * Acquire a reference to the upload request checksum context.
+ * Use this when transferring ownership to another function/structure.
+ *
+ * @param context The checksum context to acquire
+ * @return The same context pointer (for convenience)
+ */
+AWS_S3_API
+struct aws_s3_upload_request_checksum_context *aws_s3_upload_request_checksum_context_acquire(
+    struct aws_s3_upload_request_checksum_context *context);
+
+/**
+ * Release a reference to the upload request checksum context.
+ * When the reference count reaches zero, the context will be destroyed.
+ * Always returns NULL.
+ *
+ * @param context The checksum context to release
+ */
+AWS_S3_API
+struct aws_s3_upload_request_checksum_context *aws_s3_upload_request_checksum_context_release(
+    struct aws_s3_upload_request_checksum_context *context);
+
+/**
+ * Check if checksum calculation is needed based on context state.
+ * Returns true if the context has a valid algorithm and the checksum has not been calculated yet.
+ *
+ * @param context The checksum context to check
+ * @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);
+
+/**
+ * Check if checksum should be added to HTTP headers.
+ * Returns true if the context has a valid algorithm and the location is set to header.
+ *
+ * @param context The checksum context to check
+ * @return true if checksum should be added to headers, false otherwise
+ */
+AWS_S3_API
+bool aws_s3_upload_request_checksum_context_should_add_header(
+    const struct aws_s3_upload_request_checksum_context *context);
+
+/**
+ * Check if checksum should be added as trailer (aws-chunked encoding).
+ * Returns true if the context has a valid algorithm and the location is set to trailer.
+ *
+ * @param context The checksum context to check
+ * @return true if checksum should be added as trailer, false otherwise
+ */
+AWS_S3_API
+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.
+ *
+ * @param context The checksum context
+ * @return Pointer to the checksum buffer, or NULL if context is invalid
+ */
+AWS_S3_API
+struct aws_byte_buf *aws_s3_upload_request_checksum_context_get_output_buffer(
+    struct aws_s3_upload_request_checksum_context *context);
+
+/**
+ * Get a cursor to the current base64 encoded checksum value (for use in headers/trailers).
+ * Returns an empty cursor if the checksum has not been calculated yet.
+ *
+ * @param context The checksum context
+ * @return Byte cursor to the calculated checksum, or empty cursor if not available
+ */
+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);
+
+AWS_EXTERN_C_END
+
+#endif /* AWS_S3_CHECKSUM_CONTEXT_H */

include/aws/s3/private/s3_checksums.h

@@ -10,6 +10,7 @@
  * aws-c-sdkutil. */
 
 struct aws_s3_checksum;
+struct aws_s3_upload_request_checksum_context;
 
 /* List to check the checksum algorithm to use based on the priority. */
 static const enum aws_s3_checksum_algorithm s_checksum_algo_priority_list[] = {
@@ -40,7 +41,7 @@ struct aws_s3_checksum {
     } impl;
 };
 
-struct checksum_config_storage {
+struct aws_s3_meta_request_checksum_config_storage {
     struct aws_allocator *allocator;
     struct aws_byte_buf full_object_checksum;
     bool has_full_object_checksum;
@@ -83,25 +84,46 @@ struct aws_input_stream *aws_checksum_stream_new(
 
 /**
  * TODO: properly support chunked encoding.
+ * Creates a chunked encoding stream that wraps an existing stream and adds checksum trailers.
  *
- * A stream that takes in a stream, encodes it to aws_chunked. Computes a running checksum as it is read and add the
- * checksum as trailer at the end of the stream. All of the added bytes will be counted to the length of the stream.
- * Note: seek this stream will immediately fail, as it would prevent an accurate calculation of the
- * checksum.
+ * This function creates a stream that:
+ * 1. Encodes the input stream wraps the existing_stream with aws-chunked encoded.
+ * 2. Calculates a checksum of the stream content (if not already calculated)
+ * 3. Appends the checksum as a trailer at the end of the aws-chunked stream
  *
- * @param allocator
- * @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
- *                          filled with the checksum result when calculated. Callers responsibility to cleanup.
+ * Note: This stream does not support seeking operations, as seeking would prevent
+ * accurate checksum calculation and corrupt the chunked encoding format.
+ *
+ * @param allocator         Memory allocator to use for stream creation and internal buffers
+ * @param existing_stream   The input stream to be chunked and checksummed. This stream
+ *                          will be acquired by the chunk stream and released when the
+ *                          chunk stream is destroyed. Must not be NULL.
+ * @param checksum_context  Context containing checksum configuration and state. Must not be NULL.
+ *                          The context contains:
+ *                          - algorithm: The checksum algorithm to use (CRC32, CRC32C, etc.)
+ *                          - base64_checksum: Buffer for the calculated checksum result
+ *                          - checksum_calculated: Whether checksum is pre-calculated or needs calculation
+ *                          - encoded_checksum_size: Expected size of the base64-encoded checksum
+ *
+ *                          If checksum_calculated is false, the stream will wrap existing_stream
+ *                          with a checksum stream to calculate the checksum during reading.
+ *                          If checksum_calculated is true, the existing checksum value will be used.
+ *
+ * @return A new input stream that provides chunked encoding with checksum trailers,
+ *         or NULL if creation fails. The returned stream must be released with
+ *         aws_input_stream_release() when no longer needed.
+ *
+ * @note The total length of the returned stream includes:
+ *       - Chunk size header (hex representation + \r\n)
+ *       - Original stream content
+ *       - Final chunk marker (0\r\n or \r\n0\r\n)
+ *       - Checksum trailer (header name + : + base64 checksum + \r\n\r\n)
  */
 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_s3_upload_request_checksum_context *context);
 
 /**
  * Get the size of the checksum output corresponding to the aws_s3_checksum_algorithm enum value.
@@ -166,14 +188,15 @@ AWS_S3_API
 int aws_checksum_finalize(struct aws_s3_checksum *checksum, struct aws_byte_buf *output);
 
 AWS_S3_API
-int aws_checksum_config_storage_init(
+int aws_s3_meta_request_checksum_config_storage_init(
     struct aws_allocator *allocator,
-    struct checksum_config_storage *internal_config,
+    struct aws_s3_meta_request_checksum_config_storage *internal_config,
     const struct aws_s3_checksum_config *config,
     const struct aws_http_message *message,
     const void *log_id);
 
 AWS_S3_API
-void aws_checksum_config_storage_cleanup(struct checksum_config_storage *internal_config);
+void aws_s3_meta_request_checksum_config_storage_cleanup(
+    struct aws_s3_meta_request_checksum_config_storage *internal_config);
 
 #endif /* AWS_S3_CHECKSUMS_H */

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 {

include/aws/s3/private/s3_meta_request_impl.h

@@ -25,7 +25,7 @@ struct aws_s3_request;
 struct aws_http_headers;
 struct aws_http_make_request_options;
 struct aws_retry_strategy;
-
+struct aws_s3_upload_request_checksum_context;
 enum aws_s3_meta_request_state {
     AWS_S3_META_REQUEST_STATE_ACTIVE,
     AWS_S3_META_REQUEST_STATE_FINISHED,
@@ -276,7 +276,7 @@ struct aws_s3_meta_request {
     const bool should_compute_content_md5;
 
     /* deep copy of the checksum config. */
-    struct checksum_config_storage checksum_config;
+    struct aws_s3_meta_request_checksum_config_storage checksum_config;
 
     /* checksum found in either a default get request, or in the initial head request of a multipart get */
     struct aws_byte_buf meta_request_level_response_header_checksum;
@@ -294,7 +294,7 @@ struct aws_s3_meta_request {
 struct aws_s3_mpu_part_info {
     uint64_t size;
     struct aws_string *etag;
-    struct aws_byte_buf checksum_base64;
+    struct aws_s3_upload_request_checksum_context *checksum_context;
     bool was_previously_uploaded;
 };
 

include/aws/s3/private/s3_request_messages.h

@@ -17,7 +17,8 @@ struct aws_byte_buf;
 struct aws_byte_cursor;
 struct aws_string;
 struct aws_array_list;
-struct checksum_config_storage;
+struct aws_s3_meta_request_checksum_config_storage;
+struct aws_s3_upload_request_checksum_context;
 
 AWS_EXTERN_C_BEGIN
 
@@ -52,8 +53,7 @@ struct aws_input_stream *aws_s3_message_util_assign_body(
     struct aws_allocator *allocator,
     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_s3_upload_request_checksum_context *checksum_context);
 
 /* Create an HTTP request for an S3 Ranged Get Object Request, using the given request as a basis */
 AWS_S3_API
@@ -76,7 +76,7 @@ AWS_S3_API
 struct aws_http_message *aws_s3_create_multipart_upload_message_new(
     struct aws_allocator *allocator,
     struct aws_http_message *base_message,
-    const struct checksum_config_storage *checksum_config);
+    const struct aws_s3_meta_request_checksum_config_storage *checksum_config);
 
 /* Create an HTTP request for an S3 Put Object request, using the original request as a basis.  Creates and assigns a
  * body stream using the passed in buffer.  If multipart is not needed, part number and upload_id can be 0 and NULL,
@@ -89,8 +89,7 @@ struct aws_http_message *aws_s3_upload_part_message_new(
     uint32_t part_number,
     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_s3_upload_request_checksum_context *checksum_context);
 
 /* 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,
@@ -116,7 +115,7 @@ struct aws_http_message *aws_s3_complete_multipart_message_new(
     struct aws_byte_buf *body_buffer,
     const struct aws_string *upload_id,
     const struct aws_array_list *parts,
-    const struct checksum_config_storage *checksum_config);
+    const struct aws_s3_meta_request_checksum_config_storage *checksum_config);
 
 AWS_S3_API
 struct aws_http_message *aws_s3_abort_multipart_upload_message_new(