Skip to content

Commit 72b6826

Browse files
TingDaoKTingDaoK
committed
refactor it
1 parent e29d64a commit 72b6826

File tree

4 files changed

+61
-69
lines changed

4 files changed

+61
-69
lines changed

‎include/aws/s3/private/s3_checksums.h‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -62,8 +62,7 @@ struct aws_s3_meta_request_checksum_config_storage {
6262
};
6363

6464
/**
65-
* a stream that takes in a stream, computes a running checksum as it is read, and outputs the checksum when the stream
66-
* is destroyed.
65+
* a stream that takes in a stream
6766
* Note: seek this stream will immediately fail, as it would prevent an accurate calculation of the
6867
* checksum.
6968
*
@@ -72,15 +71,19 @@ struct aws_s3_meta_request_checksum_config_storage {
7271
* outputs the checksum of existing stream to checksum_output upon destruction. Will be kept
7372
* alive by the checksum stream
7473
* @param algorithm Checksum algorithm to use.
75-
* @param checksum_output Checksum of the `existing_stream`, owned by caller, which will be calculated when this stream
76-
* is destroyed.
7774
*/
7875
AWS_S3_API
7976
struct aws_input_stream *aws_checksum_stream_new(
8077
struct aws_allocator *allocator,
8178
struct aws_input_stream *existing_stream,
82-
enum aws_s3_checksum_algorithm algorithm,
83-
struct aws_byte_buf *checksum_output);
79+
enum aws_s3_checksum_algorithm algorithm);
80+
81+
/**
82+
* Finalize the checksum has read so far to the output checksum buf with base64 encoding.
83+
* Not thread safe.
84+
*/
85+
AWS_S3_API
86+
int aws_checksum_stream_finalize_checksum(struct aws_input_stream *checksum_stream, struct aws_byte_buf *checksum_buf);
8487

8588
/**
8689
* TODO: properly support chunked encoding.

‎source/s3_checksum_stream.c‎

Lines changed: 32 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -14,44 +14,8 @@ struct aws_checksum_stream {
1414
struct aws_input_stream *old_stream;
1515
struct aws_s3_checksum *checksum;
1616
struct aws_byte_buf checksum_result;
17-
/* base64 encoded checksum of the stream, updated at end of stream */
18-
struct aws_byte_buf *encoded_checksum_output;
19-
bool checksum_finalized;
2017
};
2118

22-
static int s_finalize_checksum(struct aws_checksum_stream *impl) {
23-
if (impl->checksum_finalized) {
24-
return AWS_OP_SUCCESS;
25-
}
26-
27-
if (aws_checksum_finalize(impl->checksum, &impl->checksum_result) != AWS_OP_SUCCESS) {
28-
AWS_LOGF_ERROR(
29-
AWS_LS_S3_CLIENT,
30-
"Failed to calculate checksum with error code %d (%s).",
31-
aws_last_error(),
32-
aws_error_str(aws_last_error()));
33-
aws_byte_buf_reset(&impl->checksum_result, true);
34-
impl->checksum_finalized = true;
35-
return aws_raise_error(AWS_ERROR_S3_CHECKSUM_CALCULATION_FAILED);
36-
}
37-
struct aws_byte_cursor checksum_result_cursor = aws_byte_cursor_from_buf(&impl->checksum_result);
38-
if (aws_base64_encode(&checksum_result_cursor, impl->encoded_checksum_output) != AWS_OP_SUCCESS) {
39-
AWS_LOGF_ERROR(
40-
AWS_LS_S3_CLIENT,
41-
"Failed to base64 encode checksum with error code %d (%s). Output capacity: %zu length %zu",
42-
aws_last_error(),
43-
aws_error_str(aws_last_error()),
44-
impl->encoded_checksum_output->capacity,
45-
impl->encoded_checksum_output->len);
46-
aws_byte_buf_reset(&impl->checksum_result, true);
47-
impl->checksum_finalized = true;
48-
return aws_raise_error(AWS_ERROR_S3_CHECKSUM_CALCULATION_FAILED);
49-
}
50-
51-
impl->checksum_finalized = true;
52-
return AWS_OP_SUCCESS;
53-
}
54-
5519
static int s_aws_input_checksum_stream_seek(
5620
struct aws_input_stream *stream,
5721
int64_t offset,
@@ -82,14 +46,6 @@ static int s_aws_input_checksum_stream_read(struct aws_input_stream *stream, str
8246
if (aws_checksum_update(impl->checksum, &to_sum)) {
8347
return AWS_OP_ERR;
8448
}
85-
/* If we're at the end of the stream, compute and store the final checksum */
86-
struct aws_stream_status status;
87-
if (aws_input_stream_get_status(impl->old_stream, &status)) {
88-
return AWS_OP_ERR;
89-
}
90-
if (status.is_end_of_stream) {
91-
return s_finalize_checksum(impl);
92-
}
9349
return AWS_OP_SUCCESS;
9450
}
9551

@@ -127,11 +83,8 @@ static struct aws_input_stream_vtable s_aws_input_checksum_stream_vtable = {
12783
struct aws_input_stream *aws_checksum_stream_new(
12884
struct aws_allocator *allocator,
12985
struct aws_input_stream *existing_stream,
130-
enum aws_s3_checksum_algorithm algorithm,
131-
struct aws_byte_buf *checksum_output) {
86+
enum aws_s3_checksum_algorithm algorithm) {
13287
AWS_PRECONDITION(existing_stream);
133-
AWS_PRECONDITION(checksum_output);
134-
AWS_PRECONDITION(checksum_output->len == 0 && "Checksum output buffer is not empty");
13588

13689
struct aws_checksum_stream *impl = aws_mem_calloc(allocator, 1, sizeof(struct aws_checksum_stream));
13790
impl->allocator = allocator;
@@ -143,7 +96,6 @@ struct aws_input_stream *aws_checksum_stream_new(
14396
}
14497
aws_byte_buf_init(&impl->checksum_result, allocator, impl->checksum->digest_size);
14598
impl->old_stream = aws_input_stream_acquire(existing_stream);
146-
impl->encoded_checksum_output = checksum_output;
14799
aws_ref_count_init(
148100
&impl->base.ref_count, impl, (aws_simple_completion_callback *)s_aws_input_checksum_stream_destroy);
149101

@@ -152,3 +104,34 @@ struct aws_input_stream *aws_checksum_stream_new(
152104
aws_mem_release(impl->allocator, impl);
153105
return NULL;
154106
}
107+
108+
int aws_checksum_stream_finalize_checksum(struct aws_input_stream *checksum_stream, struct aws_byte_buf *checksum_buf) {
109+
AWS_PRECONDITION(checksum_buf);
110+
AWS_PRECONDITION(checksum_buf->len == 0 && "Checksum output buffer is not empty");
111+
112+
struct aws_checksum_stream *impl = AWS_CONTAINER_OF(checksum_stream, struct aws_checksum_stream, base);
113+
114+
if (aws_checksum_finalize(impl->checksum, &impl->checksum_result) != AWS_OP_SUCCESS) {
115+
AWS_LOGF_ERROR(
116+
AWS_LS_S3_CLIENT,
117+
"Failed to calculate checksum with error code %d (%s).",
118+
aws_last_error(),
119+
aws_error_str(aws_last_error()));
120+
aws_byte_buf_reset(&impl->checksum_result, true);
121+
return aws_raise_error(AWS_ERROR_S3_CHECKSUM_CALCULATION_FAILED);
122+
}
123+
struct aws_byte_cursor checksum_result_cursor = aws_byte_cursor_from_buf(&impl->checksum_result);
124+
if (aws_base64_encode(&checksum_result_cursor, checksum_buf) != AWS_OP_SUCCESS) {
125+
AWS_LOGF_ERROR(
126+
AWS_LS_S3_CLIENT,
127+
"Failed to base64 encode checksum with error code %d (%s). Output capacity: %zu length %zu",
128+
aws_last_error(),
129+
aws_error_str(aws_last_error()),
130+
checksum_buf->capacity,
131+
checksum_buf->len);
132+
aws_byte_buf_reset(&impl->checksum_result, true);
133+
return aws_raise_error(AWS_ERROR_S3_CHECKSUM_CALCULATION_FAILED);
134+
}
135+
136+
return AWS_OP_SUCCESS;
137+
}

‎source/s3_chunk_stream.c‎

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ struct aws_chunk_stream {
2727
/* Pointing to the stream we read from */
2828
struct aws_input_stream *current_stream;
2929
struct aws_input_stream *chunk_body_stream;
30+
struct aws_input_stream *checksum_stream;
3031

3132
struct aws_s3_upload_request_checksum_context *checksum_context;
3233
struct aws_byte_buf pre_chunk_buffer;
@@ -62,7 +63,16 @@ static int s_set_post_chunk_stream(struct aws_chunk_stream *parent_stream) {
6263
struct aws_byte_cursor post_trailer_cursor = aws_byte_cursor_from_string(s_post_trailer);
6364
struct aws_byte_cursor colon_cursor = aws_byte_cursor_from_string(s_colon);
6465
/* After the checksum stream released, the checksum will be calculated. */
65-
parent_stream->checksum_context->checksum_calculated = true;
66+
if (parent_stream->checksum_context->checksum_calculated == false) {
67+
AWS_ASSERT(parent_stream->checksum_context->base64_checksum.len == 0);
68+
AWS_ASSERT(parent_stream->checksum_stream != NULL);
69+
/* Get the checksum from the checksum stream of the data read from it. */
70+
if (aws_checksum_stream_finalize_checksum(
71+
parent_stream->checksum_stream, &parent_stream->checksum_context->base64_checksum)) {
72+
return AWS_OP_ERR;
73+
}
74+
parent_stream->checksum_context->checksum_calculated = true;
75+
}
6676
struct aws_byte_cursor checksum_result_cursor =
6777
aws_byte_cursor_from_buf(&parent_stream->checksum_context->base64_checksum);
6878
AWS_ASSERT(parent_stream->checksum_context->encoded_checksum_size == checksum_result_cursor.len);
@@ -159,12 +169,9 @@ static int s_aws_input_chunk_stream_get_length(struct aws_input_stream *stream,
159169

160170
static void s_aws_input_chunk_stream_destroy(struct aws_chunk_stream *impl) {
161171
if (impl) {
162-
if (impl->current_stream) {
163-
aws_input_stream_release(impl->current_stream);
164-
}
165-
if (impl->chunk_body_stream) {
166-
aws_input_stream_release(impl->chunk_body_stream);
167-
}
172+
aws_input_stream_release(impl->current_stream);
173+
aws_input_stream_release(impl->chunk_body_stream);
174+
aws_input_stream_release(impl->checksum_stream);
168175
aws_byte_buf_clean_up(&impl->pre_chunk_buffer);
169176
aws_byte_buf_clean_up(&impl->post_chunk_buffer);
170177
/* Either we calculated the checksum, or we the checksum is empty. Otherwise, something was wrong. */
@@ -196,12 +203,10 @@ struct aws_input_stream *aws_chunk_stream_new(
196203

197204
/* Extract algorithm and buffer from context */
198205
enum aws_s3_checksum_algorithm algorithm = AWS_SCA_NONE;
199-
struct aws_byte_buf *checksum_buffer = NULL;
200206

201207
impl->checksum_context = aws_s3_upload_request_checksum_context_acquire(checksum_context);
202208

203209
algorithm = checksum_context->algorithm;
204-
checksum_buffer = &checksum_context->base64_checksum;
205210
bool checksum_calculated = checksum_context->checksum_calculated;
206211

207212
int64_t stream_length = 0;
@@ -227,10 +232,11 @@ struct aws_input_stream *aws_chunk_stream_new(
227232
}
228233
if (!checksum_calculated) {
229234
/* Wrap the existing stream with checksum stream to calculate the checksum when reading from it. */
230-
impl->chunk_body_stream = aws_checksum_stream_new(allocator, existing_stream, algorithm, checksum_buffer);
231-
if (impl->chunk_body_stream == NULL) {
235+
impl->checksum_stream = aws_checksum_stream_new(allocator, existing_stream, algorithm);
236+
if (impl->checksum_stream == NULL) {
232237
goto error;
233238
}
239+
impl->chunk_body_stream = aws_input_stream_acquire(impl->checksum_stream);
234240
} else {
235241
/* No need to calculate the checksum during read, use the existing stream directly. */
236242
impl->chunk_body_stream = aws_input_stream_acquire(existing_stream);

‎tests/s3_checksum_stream_test.c‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,7 @@ static int compare_checksum_stream(struct aws_allocator *allocator, struct aws_b
2828
struct aws_byte_cursor checksum_result_cursor = aws_byte_cursor_from_buf(&compute_checksum_output);
2929
aws_base64_encode(&checksum_result_cursor, &compute_encoded_checksum_output);
3030
struct aws_input_stream *cursor_stream = aws_input_stream_new_from_cursor(allocator, input);
31-
struct aws_input_stream *stream =
32-
aws_checksum_stream_new(allocator, cursor_stream, algorithm, &stream_checksum_output);
31+
struct aws_input_stream *stream = aws_checksum_stream_new(allocator, cursor_stream, algorithm);
3332
aws_input_stream_release(cursor_stream);
3433
struct aws_stream_status status;
3534
AWS_ZERO_STRUCT(status);
@@ -38,11 +37,12 @@ static int compare_checksum_stream(struct aws_allocator *allocator, struct aws_b
3837
read_buf.len = 0;
3938
ASSERT_TRUE(aws_input_stream_get_status(stream, &status) == 0);
4039
}
41-
aws_input_stream_release(stream);
40+
aws_checksum_stream_finalize_checksum(stream, &stream_checksum_output);
4241
ASSERT_TRUE(aws_byte_buf_eq(&compute_encoded_checksum_output, &stream_checksum_output));
4342
aws_byte_buf_clean_up(&compute_checksum_output);
4443
aws_byte_buf_clean_up(&stream_checksum_output);
4544
aws_byte_buf_clean_up(&compute_encoded_checksum_output);
45+
aws_input_stream_release(stream);
4646
}
4747
aws_byte_buf_clean_up(&read_buf);
4848
return AWS_OP_SUCCESS;

0 commit comments

Comments
 (0)