Merge pull request #292 from jeffersonwhite/api-request-uuid-support

Accept A-api-request-uuid and add it to presigned when available
asfadmin · Apr 17, 2024 · a9a00d1 · a9a00d1
2 parents 318aac2 + 01b6acf
commit a9a00d1
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 1 deletion.
diff --git a/rain_api_core/egress_util.py b/rain_api_core/egress_util.py
@@ -23,7 +23,16 @@ def hmacsha256(key: bytes, string: str) -> hmac.HMAC:
     return hmac.new(key, string.encode(), sha256)
 
 
-def get_presigned_url(session, bucket_name, object_name, region_name, expire_seconds, user_id, method='GET') -> str:
+def get_presigned_url(
+    session,
+    bucket_name,
+    object_name,
+    region_name,
+    expire_seconds,
+    user_id,
+    method='GET',
+    api_request_uuid=None
+) -> str:
     timez = datetime.utcnow().strftime('%Y%m%dT%H%M%SZ')
     datez = timez[:8]
     region_id = "." + region_name if region_name != "us-east-1" else ""
@@ -47,6 +56,12 @@ def get_presigned_url(session, bucket_name, object_name, region_name, expire_sec
         "X-Amz-SignedHeaders=host"
     ])
 
+    if api_request_uuid is not None:
+        can_query_string = "&".join([
+            f"A-api-request-uuid={api_request_uuid}",
+            can_query_string,
+        ])
+
     can_request = (
         f"{method}\n"
         f"/{object_name}\n"

diff --git a/tests/test_egress_util.py b/tests/test_egress_util.py
@@ -116,3 +116,36 @@ def test_get_presigned_url_with_newlines(mock_datetime):
         "&X-Amz-SignedHeaders=host"
         "&X-Amz-Signature=b5f7dab0ba1af2a1a5c625f526de6ab8047a20f165675890cfa8caf70cfed730"
     )
+
+
+@mock.patch(f"{MODULE}.datetime", autospec=True)
+def test_get_presigned_url_with_api_request_uuid(mock_datetime):
+    mock_datetime.utcnow.return_value = datetime(2024, 1, 1)
+    session = {
+        "Credentials": {
+            "AccessKeyId": "access_key_id",
+            "SecretAccessKey": "secret_access_key",
+            "SessionToken": "session_token"
+        }
+    }
+    presigned_url = get_presigned_url(
+        session,
+        "bucket_name",
+        "object_name",
+        "region_name",
+        500,
+        "user_id",
+        api_request_uuid="uuid_value"
+    )
+    assert presigned_url == (
+        "https://bucket_name.s3.region_name.amazonaws.com/object_name"
+        "?A-api-request-uuid=uuid_value"
+        "&A-userid=user_id"
+        "&X-Amz-Algorithm=AWS4-HMAC-SHA256"
+        "&X-Amz-Credential=access_key_id%2F20240101%2Fregion_name%2Fs3%2Faws4_request"
+        "&X-Amz-Date=20240101T000000Z"
+        "&X-Amz-Expires=500"
+        "&X-Amz-Security-Token=session_token"
+        "&X-Amz-SignedHeaders=host"
+        "&X-Amz-Signature=3010924b3725f8b47a3486ea4005f658128ddc5dde5ca405dfe1bd69bd2ae7a8"
+    )