refactor(siv4_helper.py): move signing logic from client creation to an event hook

Author: Kyon Caldera

mcp_proxy_for_aws/hooks.py

@@ -17,8 +17,6 @@
 import httpx
 import json
 import logging
-from botocore.auth import SigV4Auth
-from botocore.awsrequest import AWSRequest
 from httpx._content import ByteStream
 from typing import Any, Dict, Optional
 
@@ -76,73 +74,50 @@ async def _handle_error_response(response: httpx.Response) -> None:
                 )
 
 
-def _resign_request_with_sigv4(
-    request: httpx.Request,
+async def _sign_request_hook(
     region: str,
     service: str,
-    profile: Optional[str] = None,
+    profile: Optional[str],
+    request: httpx.Request,
 ) -> None:
-    """Re-sign an HTTP request with AWS SigV4 after content modification.
+    """Request hook to sign HTTP requests with AWS SigV4.
 
-    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.
+    This hook signs the request with AWS SigV4 credentials and adds signature headers.
+
+    This should be the last hook called to ensure the signature includes any modifications.
 
     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)
+        request: The HTTP request object to sign (modified in-place)
     """
-    # Import here to avoid circular dependency
-    from mcp_proxy_for_aws.sigv4_helper import create_aws_session
-
-    # 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)
+    # Import here to avoid circular dependency and for compatibility
+    from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session
 
-    # Set the new Content-Length
+    # Set Content-Length for signing
     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
+    logger.info('Signing request with credentials for access key: %s', credentials.access_key)
 
-    # 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,
-    )
+    # Create SigV4 auth and use its signing logic
+    auth = SigV4HTTPXAuth(credentials, service, region)
 
-    # 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)
+    # Call auth_flow to sign the request (it modifies request in-place)
+    auth_flow = auth.auth_flow(request)
+    next(auth_flow)  # Execute the generator to perform signing
 
-    # Update request headers with signed headers
-    request.headers.update(dict(aws_request.headers))
-    logger.info('Request headers after re-signing: %s', request.headers)
+    logger.debug('Request headers after signing: %s', request.headers)
 
 
-async def _inject_metadata_hook(
-    metadata: Dict[str, Any], region: str, service: str, request: httpx.Request
-) -> None:
+async def _inject_metadata_hook(metadata: Dict[str, Any], 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 ===')
@@ -189,9 +164,6 @@ async def _inject_metadata_hook(
                 request.stream = ByteStream(new_content)
                 request._content = new_content
 
-                # Re-sign the request with the new content
-                _resign_request_with_sigv4(request, region, service)
-
                 logger.info('Injected metadata into _meta: %s', body['params']['_meta'])
 
         except (json.JSONDecodeError, KeyError, TypeError) as e:

mcp_proxy_for_aws/server.py

@@ -64,10 +64,10 @@ async def setup_mcp_mode(local_mcp: FastMCP, args) -> None:
 
     # Log server configuration
     logger.info(
-        'Using service: %s, region: %s, forwarding region: %s, profile: %s',
+        'Using service: %s, region: %s, metadata: %s, profile: %s',
         service,
         region,
-        metadata.get('AWS_REGION'),
+        metadata,
         profile,
     )
     logger.info('Running in MCP mode')

mcp_proxy_for_aws/sigv4_helper.py

@@ -21,7 +21,11 @@
 from botocore.awsrequest import AWSRequest
 from botocore.credentials import Credentials
 from functools import partial
-from mcp_proxy_for_aws.hooks import _handle_error_response, _inject_metadata_hook
+from mcp_proxy_for_aws.hooks import (
+    _handle_error_response,
+    _inject_metadata_hook,
+    _sign_request_hook,
+)
 from typing import Any, Dict, Generator, Optional
 
 
@@ -102,42 +106,12 @@ def create_aws_session(profile: Optional[str] = None) -> boto3.Session:
     return session
 
 
-def create_sigv4_auth(service: str, region: str, profile: Optional[str] = None) -> SigV4HTTPXAuth:
-    """Create SigV4 authentication for AWS requests.
-
-    Args:
-        service: AWS service name for SigV4 signing
-        profile: AWS profile to use (optional)
-        region: AWS region (defaults to AWS_REGION env var or us-east-1)
-
-    Returns:
-        SigV4HTTPXAuth instance
-
-    Raises:
-        ValueError: If credentials cannot be obtained
-    """
-    # Create session and get credentials
-    session = create_aws_session(profile)
-    credentials = session.get_credentials()
-
-    # Create SigV4Auth with explicit credentials
-    sigv4_auth = SigV4HTTPXAuth(
-        credentials=credentials,
-        service=service,
-        region=region,
-    )
-
-    logger.info("Created SigV4 authentication for service '%s' in region '%s'", service, region)
-    return sigv4_auth
-
-
 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,
 ) -> httpx.AsyncClient:
@@ -149,7 +123,6 @@ def create_sigv4_client(
         region: AWS region (optional, defaults to AWS_REGION env var or us-east-1)
         timeout: Timeout configuration for the HTTP client
         headers: Headers to include in requests
-        auth: Auth parameter (ignored as we provide our own)
         metadata: Metadata to inject into MCP _meta field
         **kwargs: Additional arguments to pass to httpx.AsyncClient
 
@@ -174,17 +147,15 @@ def create_sigv4_client(
         'Creating httpx.AsyncClient with custom headers: %s', client_kwargs.get('headers', {})
     )
 
-    # Create SigV4 auth
-    sigv4_auth = create_sigv4_auth(service, region, profile)
-
-    # Create the client with SigV4 auth and error handling event hook
-    logger.info("Creating httpx.AsyncClient with SigV4 authentication for service '%s'", service)
+    logger.info("Creating httpx.AsyncClient with SigV4 request hooks for service '%s'", service)
 
     return httpx.AsyncClient(
-        auth=sigv4_auth,
         **client_kwargs,
         event_hooks={
             'response': [_handle_error_response],
-            'request': [partial(_inject_metadata_hook, metadata or {}, region, service)],
+            'request': [
+                partial(_inject_metadata_hook, metadata or {}),
+                partial(_sign_request_hook, region, service, profile),
+            ],
         },
     )