diff --git a/include/aws/s3/private/s3_meta_request_impl.h b/include/aws/s3/private/s3_meta_request_impl.h index 024680c8f..a5edf3623 100644 --- a/include/aws/s3/private/s3_meta_request_impl.h +++ b/include/aws/s3/private/s3_meta_request_impl.h @@ -169,6 +169,7 @@ struct aws_s3_meta_request { /* Customer specified callbacks. */ aws_s3_meta_request_headers_callback_fn *headers_callback; + aws_s3_meta_request_should_continue_fn *continue_callback; aws_s3_meta_request_receive_body_callback_fn *body_callback; aws_s3_meta_request_finish_fn *finish_callback; aws_s3_meta_request_shutdown_fn *shutdown_callback; diff --git a/include/aws/s3/s3_client.h b/include/aws/s3/s3_client.h index c765d79da..029a40dad 100644 --- a/include/aws/s3/s3_client.h +++ b/include/aws/s3/s3_client.h @@ -178,6 +178,11 @@ typedef int(aws_s3_meta_request_receive_body_callback_fn)( /* User data specified by aws_s3_meta_request_options.*/ void *user_data); +/** + * Invoked to poll whether the user has canceled the meta request. + */ +typedef bool(aws_s3_meta_request_should_continue_fn)(void *user_data); + /** * Invoked when the entire meta request execution is complete. */ @@ -767,6 +772,12 @@ struct aws_s3_meta_request_options { */ aws_s3_meta_request_receive_body_callback_fn *body_callback; + /** + * Invoked to check whether the user has set the cancellation flag for this request. + * See `aws_s3_meta_request_should_continue_fn`. + */ + aws_s3_meta_request_should_continue_fn *continue_callback; + /** * Invoked when the entire meta request execution is complete. * See `aws_s3_meta_request_finish_fn`. diff --git a/source/s3_meta_request.c b/source/s3_meta_request.c index 2982bb40d..afaeebecf 100644 --- a/source/s3_meta_request.c +++ b/source/s3_meta_request.c @@ -306,6 +306,7 @@ int aws_s3_meta_request_init_base( meta_request->meta_request_level_running_response_sum = NULL; meta_request->user_data = options->user_data; + meta_request->continue_callback = options->continue_callback; meta_request->progress_callback = options->progress_callback; meta_request->telemetry_callback = options->telemetry_callback; meta_request->upload_review_callback = options->upload_review_callback; @@ -452,10 +453,19 @@ bool aws_s3_meta_request_has_finish_result_synced(struct aws_s3_meta_request *me ASSERT_SYNCED_DATA_LOCK_HELD(meta_request); if (!meta_request->synced_data.finish_result_set) { - return false; + /* + * Perform the equivalent of calling aws_s3_meta_request_cancel() if continue_callback() indicates + * that the request should be canceled. This is a work-around to enable canceling ongoing requests + * via a user-provided function. The work-around is needed until the AWS C++ SDK supports Cancel(). + * This function was chosen since it is evaluated several times during the lifetime of a request. + */ + if (meta_request->continue_callback && !meta_request->continue_callback(meta_request->user_data)) { + aws_s3_meta_request_set_fail_synced(meta_request, NULL, AWS_ERROR_S3_CANCELED); + aws_s3_meta_request_cancel_cancellable_requests_synced(meta_request, AWS_ERROR_S3_CANCELED); + } } - return true; + return meta_request->synced_data.finish_result_set; } struct aws_s3_meta_request *aws_s3_meta_request_acquire(struct aws_s3_meta_request *meta_request) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 353b07318..86f26e6c3 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -64,6 +64,8 @@ add_net_test_case(test_s3_cancel_mpd_empty_object_get_with_part_number_1_sent) add_net_test_case(test_s3_cancel_mpd_empty_object_get_with_part_number_1_completed) add_net_test_case(test_s3_cancel_mpd_pending_streaming) add_net_test_case(test_s3_cancel_prepare) +add_net_test_case(test_s3_continue_callback_ok) +add_net_test_case(test_s3_continue_callback_cancels) add_net_test_case(test_s3_get_object_tls_disabled) add_net_test_case(test_s3_get_object_tls_enabled) diff --git a/tests/s3_cancel_tests.c b/tests/s3_cancel_tests.c index 2ffc6b7a8..41d6be43e 100644 --- a/tests/s3_cancel_tests.c +++ b/tests/s3_cancel_tests.c @@ -705,6 +705,11 @@ static void s_test_s3_cancel_prepare_meta_request_prepare_request_on_original_do ++test_user_data->request_prepare_counters[request->request_tag]; + if (results->continue_callback != NULL) { + /* Cancel via s_test_continue_callback_cancels_after_first_part(). */ + return AWS_OP_SUCCESS; + } + /* Cancel after the first part is prepared, preventing any additional parts from being prepared. */ if (request->request_tag == AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_PART && test_user_data->request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_PART] == 1) { @@ -787,3 +792,112 @@ static int s_test_s3_cancel_prepare(struct aws_allocator *allocator, void *ctx) return 0; } + +static bool s_test_continue_callback_cancels_after_first_part(void *user_data) { + struct aws_s3_meta_request_test_results *results = user_data; + AWS_ASSERT(results != NULL); + + struct aws_s3_tester *tester = results->tester; + AWS_ASSERT(tester != NULL); + + /* Similar to s_test_s3_cancel_prepare_meta_request_prepare_request: cancel after the first part. */ + struct test_s3_cancel_prepare_user_data *test_user_data = tester->user_data; + return test_user_data->request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_PART] == 0; +} + +AWS_TEST_CASE(test_s3_continue_callback_cancels, s_test_s3_continue_callback_cancels) +static int s_test_s3_continue_callback_cancels(struct aws_allocator *allocator, void *ctx) { + (void)ctx; + + struct aws_s3_tester tester; + ASSERT_SUCCESS(aws_s3_tester_init(allocator, &tester)); + + struct test_s3_cancel_prepare_user_data test_user_data; + AWS_ZERO_STRUCT(test_user_data); + tester.user_data = &test_user_data; + + struct aws_s3_client *client = NULL; + + struct aws_s3_tester_client_options client_options; + AWS_ZERO_STRUCT(client_options); + + ASSERT_SUCCESS(aws_s3_tester_client_new(&tester, &client_options, &client)); + + struct aws_s3_client_vtable *patched_client_vtable = aws_s3_tester_patch_client_vtable(&tester, client, NULL); + patched_client_vtable->meta_request_factory = s_test_s3_cancel_prepare_meta_request_factory; + + { + struct aws_s3_tester_meta_request_options options = { + .allocator = allocator, + .client = client, + .meta_request_type = AWS_S3_META_REQUEST_TYPE_PUT_OBJECT, + .validate_type = AWS_S3_TESTER_VALIDATE_TYPE_EXPECT_FAILURE, + .continue_callback = s_test_continue_callback_cancels_after_first_part, + .put_options = + { + .ensure_multipart = true, + }, + }; + struct aws_s3_meta_request_test_results meta_request_test_results; + AWS_ZERO_STRUCT(meta_request_test_results); + + ASSERT_SUCCESS(aws_s3_tester_send_meta_request_with_options(&tester, &options, &meta_request_test_results)); + ASSERT_TRUE(meta_request_test_results.finished_error_code == AWS_ERROR_S3_CANCELED); + + aws_s3_meta_request_test_results_clean_up(&meta_request_test_results); + } + + ASSERT_TRUE( + test_user_data.request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_CREATE_MULTIPART_UPLOAD] == 1); + ASSERT_TRUE(test_user_data.request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_PART] == 1); + ASSERT_TRUE( + test_user_data.request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_ABORT_MULTIPART_UPLOAD] == 1); + ASSERT_TRUE( + test_user_data.request_prepare_counters[AWS_S3_AUTO_RANGED_PUT_REQUEST_TAG_COMPLETE_MULTIPART_UPLOAD] == 0); + + aws_s3_client_release(client); + aws_s3_tester_clean_up(&tester); + + return 0; +} + +static bool s_test_continue_callback_ok(void *user_data) { + (void)user_data; + return true; +} + +AWS_TEST_CASE(test_s3_continue_callback_ok, s_test_s3_continue_callback_ok) +static int s_test_s3_continue_callback_ok(struct aws_allocator *allocator, void *ctx) { + (void)ctx; + + struct aws_s3_tester tester; + ASSERT_SUCCESS(aws_s3_tester_init(allocator, &tester)); + + struct aws_s3_client *client = NULL; + + struct aws_s3_tester_client_options client_options; + AWS_ZERO_STRUCT(client_options); + + ASSERT_SUCCESS(aws_s3_tester_client_new(&tester, &client_options, &client)); + + { + struct aws_s3_tester_meta_request_options options = { + .allocator = allocator, + .client = client, + .meta_request_type = AWS_S3_META_REQUEST_TYPE_PUT_OBJECT, + .validate_type = AWS_S3_TESTER_VALIDATE_TYPE_EXPECT_SUCCESS, + .continue_callback = s_test_continue_callback_ok, + .put_options = + { + .ensure_multipart = true, + }, + }; + + ASSERT_SUCCESS(aws_s3_tester_send_meta_request_with_options(&tester, &options, NULL)); + } + + aws_s3_client_release(client); + aws_s3_tester_clean_up(&tester); + + return 0; +} diff --git a/tests/s3_tester.c b/tests/s3_tester.c index d879f3857..3d3e47bcd 100644 --- a/tests/s3_tester.c +++ b/tests/s3_tester.c @@ -114,6 +114,16 @@ static int s_s3_test_meta_request_header_callback( return AWS_OP_SUCCESS; } +static bool s_s3_test_meta_request_continue_callback(void *user_data) { + struct aws_s3_meta_request_test_results *meta_request_test_results = + (struct aws_s3_meta_request_test_results *)user_data; + + if (meta_request_test_results->continue_callback != NULL) { + return meta_request_test_results->continue_callback(user_data); + } + return true; +} + static int s_s3_test_meta_request_body_callback( struct aws_s3_meta_request *meta_request, const struct aws_byte_cursor *body, @@ -506,6 +516,9 @@ int aws_s3_tester_bind_meta_request( ASSERT_TRUE(options->headers_callback == NULL); options->headers_callback = s_s3_test_meta_request_header_callback; + ASSERT_TRUE(options->continue_callback == NULL); + options->continue_callback = s_s3_test_meta_request_continue_callback; + ASSERT_TRUE(options->body_callback == NULL); options->body_callback = s_s3_test_meta_request_body_callback; @@ -1732,6 +1745,7 @@ int aws_s3_tester_send_meta_request_with_options( } out_results->headers_callback = options->headers_callback; + out_results->continue_callback = options->continue_callback; out_results->body_callback = options->body_callback; out_results->finish_callback = options->finish_callback; out_results->progress_callback = options->progress_callback; diff --git a/tests/s3_tester.h b/tests/s3_tester.h index 03b1f1f6c..5b41e86e2 100644 --- a/tests/s3_tester.h +++ b/tests/s3_tester.h @@ -172,6 +172,7 @@ struct aws_s3_tester_meta_request_options { bool use_s3express_signing; aws_s3_meta_request_headers_callback_fn *headers_callback; + aws_s3_meta_request_should_continue_fn *continue_callback; aws_s3_meta_request_receive_body_callback_fn *body_callback; aws_s3_meta_request_finish_fn *finish_callback; aws_s3_meta_request_progress_fn *progress_callback; @@ -238,6 +239,7 @@ struct aws_s3_meta_request_test_results { struct aws_s3_tester *tester; aws_s3_meta_request_headers_callback_fn *headers_callback; + aws_s3_meta_request_should_continue_fn *continue_callback; aws_s3_meta_request_receive_body_callback_fn *body_callback; aws_s3_meta_request_finish_fn *finish_callback; aws_s3_meta_request_progress_fn *progress_callback;