Skip to content

Commit

Permalink
Honor "request_checksum_calculation = WHEN_REQUIRED" in config (#327)
Browse files Browse the repository at this point in the history
  • Loading branch information
metadaddy committed Jan 18, 2025
1 parent cb67175 commit cee3ab5
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 2 deletions.
3 changes: 2 additions & 1 deletion s3transfer/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -519,7 +519,8 @@ def _validate_all_known_args(self, actual, allowed):
)

def _add_operation_defaults(self, extra_args):
set_default_checksum_algorithm(extra_args)
if self.client.meta.config.request_checksum_calculation == "when_supported":
set_default_checksum_algorithm(extra_args)

def _submit_transfer(
self, call_args, submission_task_cls, extra_main_kwargs=None
Expand Down
73 changes: 72 additions & 1 deletion tests/functional/test_upload.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from botocore.awsrequest import AWSRequest
from botocore.client import Config
from botocore.exceptions import ClientError
from botocore.httpchecksum import DEFAULT_CHECKSUM_ALGORITHM
from botocore.stub import ANY

from s3transfer.manager import TransferConfig, TransferManager
Expand Down Expand Up @@ -74,6 +75,16 @@ def setUp(self):
'before-parameter-build.s3.*', self.collect_body
)

# A list to keep track of all of the headers sent over the wire
# and their order.
self.sent_headers = []
# collect_headers needs to be called before the stubber, or it
# won't be called at all. Note that we must specify the service
# as '*', otherwise the stubber short-circuits the event.
self.client.meta.events.register_first(
'before-call.*.*', self.collect_headers
)

def tearDown(self):
super().tearDown()
shutil.rmtree(self.tempdir)
Expand All @@ -99,6 +110,9 @@ def collect_body(self, params, model, **kwargs):
)
self.sent_bodies.append(self._stream_body(params['Body']))

def collect_headers(self, model, params, request_signer, context, **kwargs):
self.sent_headers.append(params['headers'])

def _stream_body(self, body):
read_amt = 8 * 1024
data = body.read(read_amt)
Expand Down Expand Up @@ -163,6 +177,9 @@ def add_put_object_response_with_default_expected_params(
def assert_put_object_body_was_correct(self):
self.assertEqual(self.sent_bodies, [self.content])

def assert_put_object_request_checksum_algorithm_was_correct(self):
self.assertEqual(self.sent_headers[0]['x-amz-sdk-checksum-algorithm'], DEFAULT_CHECKSUM_ALGORITHM)

def test_upload(self):
self.extra_args['RequestPayer'] = 'requester'
self.add_put_object_response_with_default_expected_params(
Expand All @@ -174,6 +191,7 @@ def test_upload(self):
future.result()
self.assert_expected_client_calls_were_correct()
self.assert_put_object_body_was_correct()
self.assert_put_object_request_checksum_algorithm_was_correct()

def test_upload_with_checksum(self):
self.extra_args['ChecksumAlgorithm'] = 'sha256'
Expand All @@ -186,6 +204,50 @@ def test_upload_with_checksum(self):
future.result()
self.assert_expected_client_calls_were_correct()
self.assert_put_object_body_was_correct()
self.assertEqual(self.sent_headers[0]['x-amz-sdk-checksum-algorithm'], 'sha256')

def test_upload_with_request_checksum_calculation_when_supported(self):
self.add_put_object_response_with_default_expected_params(
extra_expected_params={'ChecksumAlgorithm': DEFAULT_CHECKSUM_ALGORITHM}
)
future = self.manager.upload(
self.filename, self.bucket, self.key, self.extra_args
)
future.result()
self.assert_expected_client_calls_were_correct()
self.assert_put_object_body_was_correct()
self.assert_put_object_request_checksum_algorithm_was_correct()

def test_upload_with_request_checksum_calculation_when_required(self):
self.reset_stubber_with_new_client(
{'config': Config(request_checksum_calculation='when_required')}
)
self.client.meta.events.register(
'before-parameter-build.s3.*', self.collect_body
)
self.client.meta.events.register_first(
'before-call.*.*', self.collect_headers
)
self._manager = TransferManager(self.client, self.config)

# We need a custom response with no 'ChecksumAlgorithm' in the expected params
self.stubber.add_response(
method='put_object',
service_response={},
expected_params={
'Body': ANY,
'Bucket': self.bucket,
'Key': self.key,
},
)

future = self.manager.upload(
self.filename, self.bucket, self.key, self.extra_args
)
future.result()
self.assert_expected_client_calls_were_correct()
self.assert_put_object_body_was_correct()
self.assertNotIn('x-amz-sdk-checksum-algorithm', self.sent_headers[0])

def test_upload_with_s3express_default_checksum(self):
s3express_bucket = "mytestbucket--usw2-az6--x-s3"
Expand All @@ -200,6 +262,7 @@ def test_upload_with_s3express_default_checksum(self):
future.result()
self.assert_expected_client_calls_were_correct()
self.assert_put_object_body_was_correct()
self.assert_put_object_request_checksum_algorithm_was_correct()

def test_upload_for_fileobj(self):
self.add_put_object_response_with_default_expected_params()
Expand All @@ -210,6 +273,7 @@ def test_upload_for_fileobj(self):
future.result()
self.assert_expected_client_calls_were_correct()
self.assert_put_object_body_was_correct()
self.assert_put_object_request_checksum_algorithm_was_correct()

def test_upload_for_seekable_filelike_obj(self):
self.add_put_object_response_with_default_expected_params()
Expand All @@ -220,6 +284,7 @@ def test_upload_for_seekable_filelike_obj(self):
future.result()
self.assert_expected_client_calls_were_correct()
self.assert_put_object_body_was_correct()
self.assert_put_object_request_checksum_algorithm_was_correct()

def test_upload_for_seekable_filelike_obj_that_has_been_seeked(self):
self.add_put_object_response_with_default_expected_params()
Expand All @@ -232,6 +297,7 @@ def test_upload_for_seekable_filelike_obj_that_has_been_seeked(self):
future.result()
self.assert_expected_client_calls_were_correct()
self.assertEqual(b''.join(self.sent_bodies), self.content[seek_pos:])
self.assert_put_object_request_checksum_algorithm_was_correct()

def test_upload_for_non_seekable_filelike_obj(self):
self.add_put_object_response_with_default_expected_params()
Expand All @@ -242,6 +308,7 @@ def test_upload_for_non_seekable_filelike_obj(self):
future.result()
self.assert_expected_client_calls_were_correct()
self.assert_put_object_body_was_correct()
self.assert_put_object_request_checksum_algorithm_was_correct()

def test_sigv4_progress_callbacks_invoked_once(self):
# Reset the client and manager to use sigv4
Expand All @@ -251,6 +318,9 @@ def test_sigv4_progress_callbacks_invoked_once(self):
self.client.meta.events.register(
'before-parameter-build.s3.*', self.collect_body
)
self.client.meta.events.register_first(
'before-call.*.*', self.collect_headers
)
self._manager = TransferManager(self.client, self.config)

# Add the stubbed response.
Expand All @@ -265,6 +335,7 @@ def test_sigv4_progress_callbacks_invoked_once(self):

# The amount of bytes seen should be the same as the file size
self.assertEqual(subscriber.calculate_bytes_seen(), len(self.content))
self.assert_put_object_request_checksum_algorithm_was_correct()

def test_uses_provided_osutil(self):
osutil = RecordingOSUtils()
Expand Down Expand Up @@ -316,6 +387,7 @@ def test_upload_with_bandwidth_limiter(self):

self.assert_expected_client_calls_were_correct()
self.assert_put_object_body_was_correct()
self.assert_put_object_request_checksum_algorithm_was_correct()

def test_raise_exception_on_s3_object_lambda_resource(self):
s3_object_lambda_arn = (
Expand All @@ -325,7 +397,6 @@ def test_raise_exception_on_s3_object_lambda_resource(self):
with self.assertRaisesRegex(ValueError, 'methods do not support'):
self.manager.upload(self.filename, s3_object_lambda_arn, self.key)


class TestMultipartUpload(BaseUploadTest):
__test__ = True

Expand Down

0 comments on commit cee3ab5

Please sign in to comment.