feat(sigv4_helper): add region and service argument to _inject_metadata_hook to allow for proper resigning of sigv4 to work

Author: Kyon Caldera

mcp_proxy_for_aws/sigv4_helper.py

@@ -33,10 +33,10 @@ class SigV4HTTPXAuth(httpx.Auth):
     """HTTPX Auth class that signs requests with AWS SigV4."""
 
     def __init__(
-            self,
-            credentials: Credentials,
-            service: str,
-            region: str,
+        self,
+        credentials: Credentials,
+        service: str,
+        region: str,
     ):
         """Initialize SigV4HTTPXAuth.
 
@@ -123,11 +123,70 @@ async def _handle_error_response(response: httpx.Response) -> None:
             raise e
 
 
-async def _inject_metadata_hook(metadata: Dict[str, Any], request: httpx.Request) -> None:
+def _resign_request_with_sigv4(
+    request: httpx.Request,
+    region: str,
+    service: str,
+    profile: Optional[str] = None,
+) -> None:
+    """Re-sign an HTTP request with AWS SigV4 after content modification.
+
+    This function removes old signature headers, creates a new signature based on
+    the current request content, and updates the request headers with the new signature.
+
+    Args:
+        request: The HTTP request object to re-sign (modified in-place)
+        region: AWS region for SigV4 signing
+        service: AWS service name for SigV4 signing
+        profile: AWS profile to use (optional)
+    """
+    # Remove old signature headers before re-signing
+    headers_to_remove = ['Content-Length', 'x-amz-date', 'x-amz-security-token', 'authorization']
+    for header in headers_to_remove:
+        request.headers.pop(header, None)
+
+    # Set the new Content-Length
+    request.headers['Content-Length'] = str(len(request.content))
+
+    logger.info('Headers after cleanup: %s', request.headers)
+
+    # Get AWS credentials
+    session = create_aws_session(profile)
+    credentials = session.get_credentials()
+    logger.info('Re-signing request with credentials for access key: %s', credentials.access_key)
+
+    # Create headers dict for signing, removing connection header like in auth_flow
+    headers_for_signing = dict(request.headers)
+    headers_for_signing.pop('connection', None)  # Remove connection header for signing
+
+    # Create SigV4 signer and AWS request
+    signer = SigV4Auth(credentials, service, region)
+    aws_request = AWSRequest(
+        method=request.method,
+        url=str(request.url),
+        data=request.content,
+        headers=headers_for_signing,
+    )
+
+    # Sign the request
+    logger.info('AWS request before signing: %s', aws_request.headers)
+    signer.add_auth(aws_request)
+    logger.info('AWS request after signing: %s', aws_request.headers)
+
+    # Update request headers with signed headers
+    request.headers.update(dict(aws_request.headers))
+    logger.info('Request headers after re-signing: %s', request.headers)
+
+
+async def _inject_metadata_hook(
+    metadata: Dict[str, Any], region: str, service: str, request: httpx.Request
+) -> None:
     """Request hook to inject metadata into MCP calls.
 
     Args:
         metadata: Dictionary of metadata to inject into _meta field
+        region: AWS region for SigV4 re-signing after metadata injection
+        service: AWS service name for SigV4 re-signing after metadata injection
         request: The HTTP request object
     """
     logger.info('=== Outgoing Request ===')
@@ -174,47 +233,9 @@ async def _inject_metadata_hook(metadata: Dict[str, Any], request: httpx.Request
                 request.stream = ByteStream(new_content)
                 request._content = new_content
 
-                # Remove old signature headers before re-signing
-                del request.headers['Content-Length']
-                del request.headers['x-amz-date']
-                del request.headers['x-amz-security-token']
-                del request.headers['authorization']
-
-                # Set the new Content-Length AFTER deleting the old one
-                request.headers['Content-Length'] = str(len(new_content))
-
-                logger.info('Headers are: %s', request.headers)
-                service = 'eks-mcp'
-                region = 'us-west-2'
-                session = create_aws_session()
-                credentials = session.get_credentials()
-                logger.info('DEBUG about to sign again: %s', credentials)
-
-                # Create headers dict for signing, removing connection header like in auth_flow
-                headers_for_signing = dict(request.headers)
-                headers_for_signing.pop('connection', None)  # Remove connection header for signing
-
-                signer = SigV4Auth(credentials, service, region)
-                aws_request = AWSRequest(
-                    method=request.method,
-                    url=str(request.url),
-                    data=request.content,
-                    headers=headers_for_signing,
-                )
-
-                # Sign the request with SigV4
-                logger.info('aws_request before signing: %s', aws_request.headers)
-                signer.add_auth(aws_request)
-                logger.info('aws_request after signing: %s', aws_request.headers)
-
-                # Update request headers with signed headers (this includes Authorization, x-amz-date, etc.)
-                request.headers['x-amz-date'] = aws_request.headers['x-amz-date']
-                request.headers['x-amz-security-token'] = aws_request.headers[
-                    'x-amz-security-token'
-                ]
-                request.headers['Authorization'] = aws_request.headers['Authorization']
+                # Re-sign the request with the new content
+                _resign_request_with_sigv4(request, region, service)
 
-                logger.info('New headers: %s', request.headers)
                 logger.info('Injected metadata into _meta: %s', body['params']['_meta'])
 
         except (json.JSONDecodeError, KeyError, TypeError) as e:
@@ -281,14 +302,14 @@ def create_sigv4_auth(service: str, region: str, profile: Optional[str] = None)
 
 
 def create_sigv4_client(
-        service: str,
-        region: str,
-        timeout: Optional[httpx.Timeout] = None,
-        profile: Optional[str] = None,
-        headers: Optional[Dict[str, str]] = None,
-        auth: Optional[httpx.Auth] = None,
-        metadata: Optional[Dict[str, Any]] = None,
-        **kwargs: Any,
+    service: str,
+    region: str,
+    timeout: Optional[httpx.Timeout] = None,
+    profile: Optional[str] = None,
+    headers: Optional[Dict[str, str]] = None,
+    auth: Optional[httpx.Auth] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    **kwargs: Any,
 ) -> httpx.AsyncClient:
     """Create an httpx.AsyncClient with SigV4 authentication.
 
@@ -334,6 +355,6 @@ def create_sigv4_client(
         **client_kwargs,
         event_hooks={
             'response': [_handle_error_response],
-            'request': [partial(_inject_metadata_hook, metadata or {})],
+            'request': [partial(_inject_metadata_hook, metadata or {}, region, service)],
         },
     )