Option for logging request and response bodies

Author: simonw

asgi_proxy/__init__.py

@@ -2,7 +2,9 @@
 from urllib.parse import urlparse
 
 
-def asgi_proxy(backend, log=None, timeout=None):
+def asgi_proxy(
+    backend, log=None, timeout=None, log_request_body=None, log_response_body=None
+):
     backend_host = urlparse(backend).netloc
 
     async def asgi_proxy(scope, receive, send):
@@ -32,6 +34,10 @@ async def asgi_proxy(scope, receive, send):
             body += message.get("body", b"")
             more_body = message.get("more_body", False)
 
+        # Log request body if requested
+        if log_request_body:
+            log_request_body(body)
+
         async with httpx.AsyncClient(timeout=timeout) as client:
             try:
                 # Stream it, in case of long streaming responses
@@ -41,6 +47,10 @@ async def asgi_proxy(scope, receive, send):
                     if log:
                         log.info(f"Request: {method} {url}")
                         log.info(f"Response: {resp.status_code} {resp.reason_phrase}")
+
+                    # Initialize response body accumulator if needed
+                    response_body = b"" if log_response_body else None
+
                     # Start the response
                     await send(
                         {
@@ -57,6 +67,8 @@ async def asgi_proxy(scope, receive, send):
                         # aiter_raw not aiter_bytes because we don't want
                         # content decoding to have been applied
                         async for chunk in resp.aiter_raw():
+                            if log_response_body:
+                                response_body += chunk
                             await send(
                                 {
                                     "type": "http.response.body",
@@ -71,9 +83,18 @@ async def asgi_proxy(scope, receive, send):
                                 f"Client disconnected: {e.__class__.__name__}: {e}"
                             )
                         await send({"type": "http.response.body", "more_body": False})
+
+                        # Log response body even if client disconnected
+                        if log_response_body and response_body is not None:
+                            log_response_body(response_body)
                         return
 
                     await send({"type": "http.response.body", "more_body": False})
+
+                    # Log response body if requested
+                    if log_response_body and response_body is not None:
+                        log_response_body(response_body)
+
             except httpx.TimeoutException as ex:
                 if log:
                     log.error(f"Timeout error occurred: {ex.__class__.__name__}: {ex}")

asgi_proxy/__main__.py

@@ -22,14 +22,31 @@
         "--timeout", type=float, default=None, help="Timeout in seconds"
     )
     parser.add_argument("-v", "--verbose", action="store_true", help="Verbose logging")
+    parser.add_argument(
+        "-b",
+        "--bodies",
+        action="store_true",
+        help="Log full request and response bodies",
+    )
 
     args = parser.parse_args()
 
-    if args.verbose:
+    if args.verbose or args.bodies:
         import logging
 
         logging.basicConfig(level=logging.INFO)
-        app = asgi_proxy(args.url, log=logging, timeout=args.timeout)
+        log_request_body = None
+        log_response_body = None
+        if args.bodies:
+            log_request_body = logging.info
+            log_response_body = logging.info
+        app = asgi_proxy(
+            args.url,
+            log=logging,
+            timeout=args.timeout,
+            log_request_body=log_request_body,
+            log_response_body=log_response_body,
+        )
     else:
         app = asgi_proxy(args.url, timeout=args.timeout)