From e5cfd74bce685fbba438edc2c496a6abaef9e662 Mon Sep 17 00:00:00 2001
From: Pierre MULLER <pmuller@scaleway.com>
Date: Thu, 24 Feb 2022 16:28:24 +0100
Subject: [PATCH] add ExpectedSourceBucketOwner support for copy operation

---
 s3transfer/copies.py          |  4 +++-
 s3transfer/manager.py         |  1 +
 tests/functional/test_copy.py | 23 ++++++++++++++++++++++-
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/s3transfer/copies.py b/s3transfer/copies.py
index a1dfdc8b..f769b370 100644
--- a/s3transfer/copies.py
+++ b/s3transfer/copies.py
@@ -39,7 +39,7 @@ class CopySubmissionTask(SubmissionTask):
         'CopySourceSSECustomerAlgorithm': 'SSECustomerAlgorithm',
         'CopySourceSSECustomerKeyMD5': 'SSECustomerKeyMD5',
         'RequestPayer': 'RequestPayer',
-        'ExpectedBucketOwner': 'ExpectedBucketOwner',
+        'ExpectedSourceBucketOwner': 'ExpectedBucketOwner',
     }
 
     UPLOAD_PART_COPY_ARGS = [
@@ -55,6 +55,7 @@ class CopySubmissionTask(SubmissionTask):
         'SSECustomerKeyMD5',
         'RequestPayer',
         'ExpectedBucketOwner',
+        'ExpectedSourceBucketOwner',
     ]
 
     CREATE_MULTIPART_ARGS_BLACKLIST = [
@@ -67,6 +68,7 @@ class CopySubmissionTask(SubmissionTask):
         'CopySourceSSECustomerKeyMD5',
         'MetadataDirective',
         'TaggingDirective',
+        'ExpectedSourceBucketOwner',
     ]
 
     COMPLETE_MULTIPART_ARGS = ['RequestPayer', 'ExpectedBucketOwner']
diff --git a/s3transfer/manager.py b/s3transfer/manager.py
index ff6afa12..b9f676c6 100644
--- a/s3transfer/manager.py
+++ b/s3transfer/manager.py
@@ -192,6 +192,7 @@ class TransferManager:
         'CopySourceSSECustomerKeyMD5',
         'MetadataDirective',
         'TaggingDirective',
+        'ExpectedSourceBucketOwner',
     ]
 
     ALLOWED_DELETE_ARGS = [
diff --git a/tests/functional/test_copy.py b/tests/functional/test_copy.py
index 86f05cbf..9cf4299e 100644
--- a/tests/functional/test_copy.py
+++ b/tests/functional/test_copy.py
@@ -227,17 +227,24 @@ def test_copy_with_extra_args(self):
 
     def test_copy_maps_extra_args_to_head_object(self):
         self.extra_args['CopySourceSSECustomerAlgorithm'] = 'AES256'
+        self.extra_args[
+            'ExpectedSourceBucketOwner'
+        ] = 'expectedsourcebucketowner'
+        self.extra_args['ExpectedBucketOwner'] = 'expectedbucketowner'
 
         expected_head_params = {
             'Bucket': 'mysourcebucket',
             'Key': 'mysourcekey',
             'SSECustomerAlgorithm': 'AES256',
+            'ExpectedBucketOwner': 'expectedsourcebucketowner',
         }
         expected_copy_object = {
             'Bucket': self.bucket,
             'Key': self.key,
             'CopySource': self.copy_source,
             'CopySourceSSECustomerAlgorithm': 'AES256',
+            'ExpectedSourceBucketOwner': 'expectedsourcebucketowner',
+            'ExpectedBucketOwner': 'expectedbucketowner',
         }
 
         self.add_head_object_response(expected_params=expected_head_params)
@@ -423,9 +430,10 @@ def test_copy_with_extra_args(self):
         # This extra argument should be added to the head object,
         # the create multipart upload, and upload part copy.
         self.extra_args['RequestPayer'] = 'requester'
+        self.extra_args['ExpectedBucketOwner'] = 'expectedbucketowner'
 
         head_params, add_copy_kwargs = self._get_expected_params()
-        head_params.update(self.extra_args)
+        head_params['RequestPayer'] = 'requester'
         self.add_head_object_response(expected_params=head_params)
 
         self._add_params_to_expected_params(
@@ -444,8 +452,17 @@ def test_copy_with_extra_args(self):
     def test_copy_blacklists_args_to_create_multipart(self):
         # This argument can never be used for multipart uploads
         self.extra_args['MetadataDirective'] = 'COPY'
+        self.extra_args[
+            'ExpectedSourceBucketOwner'
+        ] = 'expectedsourcebucketowner'
 
         head_params, add_copy_kwargs = self._get_expected_params()
+        head_params['ExpectedBucketOwner'] = 'expectedsourcebucketowner'
+        self._add_params_to_expected_params(
+            add_copy_kwargs,
+            ['copy'],
+            {'ExpectedSourceBucketOwner': 'expectedsourcebucketowner'},
+        )
         self.add_head_object_response(expected_params=head_params)
         self.add_successful_copy_responses(**add_copy_kwargs)
 
@@ -493,12 +510,16 @@ def test_copy_passes_args_to_create_multipart_and_upload_part(self):
 
     def test_copy_maps_extra_args_to_head_object(self):
         self.extra_args['CopySourceSSECustomerAlgorithm'] = 'AES256'
+        self.extra_args[
+            'ExpectedSourceBucketOwner'
+        ] = 'expectedsourcebucketowner'
 
         head_params, add_copy_kwargs = self._get_expected_params()
 
         # The CopySourceSSECustomerAlgorithm needs to get mapped to
         # SSECustomerAlgorithm for HeadObject
         head_params['SSECustomerAlgorithm'] = 'AES256'
+        head_params['ExpectedBucketOwner'] = 'expectedsourcebucketowner'
         self.add_head_object_response(expected_params=head_params)
 
         # However, it needs to remain the same for UploadPartCopy.