still allow it to be null

Author: TingDaoK

include/aws/s3/private/s3_checksums.h

@@ -93,7 +93,9 @@ 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_buffer   Required.
+ * @param checksum_buffer   Optional.
+ *                          - If the checksum_buffer is NULL, the checksum will still be calculated to append as
+ *                          trailer.
  *                          - 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.

source/s3_chunk_stream.c

@@ -28,7 +28,9 @@ struct aws_chunk_stream {
 
     /* 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_byte_buf *passed_in_checksum_buffer;
+    /* If there is no passed in checksum buffer, we still calculate the checksum. This stores the checksum. */
+    struct aws_byte_buf checksum_buffer;
 
     struct aws_input_stream *chunk_body_stream;
     struct aws_byte_buf pre_chunk_buffer;
@@ -63,12 +65,22 @@ 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_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;
+    if (!parent_stream->passed_in_checksum_buffer || parent_stream->passed_in_checksum_buffer->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);
+        }
+        checksum_result_cursor = aws_byte_cursor_from_buf(&parent_stream->checksum_buffer);
+        if (parent_stream->passed_in_checksum_buffer) {
+            /* the passed in checksum buffer is empty, initialize it with the calculated checksum */
+            aws_byte_buf_init_copy(
+                parent_stream->passed_in_checksum_buffer, parent_stream->allocator, &parent_stream->checksum_buffer);
+        }
+    } else {
+        /* The passed-in checksum buffer already have the checksum. */
+        checksum_result_cursor = aws_byte_cursor_from_buf(parent_stream->passed_in_checksum_buffer);
     }
-    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,
@@ -170,6 +182,7 @@ 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);
+        aws_byte_buf_clean_up(&impl->checksum_buffer);
         aws_mem_release(impl->allocator, impl);
     }
 }
@@ -188,13 +201,12 @@ struct aws_input_stream *aws_chunk_stream_new(
     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_buffer = checksum_buffer;
+    impl->passed_in_checksum_buffer = checksum_buffer;
 
     int64_t stream_length = 0;
     int64_t final_chunk_len = 0;
@@ -224,13 +236,14 @@ struct aws_input_stream *aws_chunk_stream_new(
     if (aws_base64_compute_encoded_len(checksum_len, &encoded_checksum_len)) {
         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)) {
+    if (!checksum_buffer || checksum_buffer->len == 0) {
+        /* Empty passed-in buffer or no passed-in, 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);
+        impl->chunk_body_stream =
+            aws_checksum_stream_new(allocator, existing_stream, algorithm, &impl->checksum_buffer);
         if (impl->chunk_body_stream == NULL) {
             goto error;
         }
@@ -284,6 +297,7 @@ struct aws_input_stream *aws_chunk_stream_new(
     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_buffer);
     aws_mem_release(impl->allocator, impl);
     return NULL;
 }