Merge pull request #5212 from IBM/fix/session-affinity-auth-context

fix(session-affinity): carry the edge-validated auth context across cross-worker forwards

AUTH_ENCRYPTION_SECRET is the critical security here, before it goes to merge we need to add a way to validate if this was changed in production. It should not use default value.

Author: cafalchio

mcpgateway/auth_context.py

@@ -221,6 +221,25 @@ def get_internal_mcp_auth_context(request: Request) -> Optional[Dict[str, Any]]:
     return None
 
 
+def encode_internal_mcp_auth_context(auth_context: Dict[str, Any]) -> str:
+    """Encode an edge-validated auth context for forwarding to a trusted internal dispatcher.
+
+    Mirror of :func:`decode_internal_mcp_auth_context`. Packages the auth context
+    so a trusted internal dispatcher (e.g. ``/_internal/mcp/rpc``) can use it
+    without re-running authentication. The unpadded base64url shape keeps the
+    value header-safe.
+
+    Args:
+        auth_context: Auth-context dict to encode.
+
+    Returns:
+        Unpadded base64url-encoded JSON payload for the
+        ``x-contextforge-auth-context`` header.
+    """
+    encoded = base64.urlsafe_b64encode(orjson.dumps(auth_context)).decode("ascii")
+    return encoded.rstrip("=")
+
+
 def decode_internal_mcp_auth_context(header_value: str) -> Dict[str, Any]:
     """Decode the trusted internal MCP auth header payload.
 

mcpgateway/config.py

@@ -458,6 +458,7 @@ def validate_auth_header_name(cls, v: str) -> str:
             "/sse",  # Exempt: SSE is a server-sent event stream, not vulnerable to CSRF
             "/message",  # Exempt: MCP SSE message endpoint
             "/rpc",  # Exempt: JSON-RPC is a programmatic protocol, not browser-based
+            "/_internal/mcp/",  # Exempt: loopback-only, HMAC-gated internal dispatch (affinity/Rust forwards); not browser-reachable
         ],
         description="Paths exempt from CSRF protection",
     )

mcpgateway/main.py

@@ -277,18 +277,28 @@
 
 
 def _is_trusted_internal_mcp_runtime_request(request: Request) -> bool:
-    """Return whether the request came from the local Rust runtime sidecar.
+    """Return whether the request came from a trusted local internal source.
+
+    Two callers are trusted today:
+
+    - ``"rust"`` — the local Rust runtime sidecar (over loopback).
+    - ``"affinity"`` — the in-process dispatch used by session-affinity
+      forwarding to reach the owner worker, carrying the identity the edge
+      already validated.
+
+    Both share the same gates: a shared-secret HMAC header AND a loopback client
+    address. Only the ``x-contextforge-mcp-runtime`` marker value differs.
 
     Args:
         request: Incoming request to inspect.
 
     Returns:
-        ``True`` when the request carries the trusted Rust runtime marker from
-        loopback, otherwise ``False``.
+        ``True`` when the request carries a trusted internal-runtime marker
+        from loopback, otherwise ``False``.
     """
     runtime_marker = request.headers.get("x-contextforge-mcp-runtime")
     client_host = getattr(getattr(request, "client", None), "host", None)
-    if runtime_marker != "rust" or not has_valid_internal_mcp_runtime_auth_header(request) or client_host not in ("127.0.0.1", "::1"):
+    if runtime_marker not in ("rust", "affinity") or not has_valid_internal_mcp_runtime_auth_header(request) or client_host not in ("127.0.0.1", "::1"):
         return False
     # Defense-in-depth: /_internal/a2a/* endpoints must refuse requests when
     # A2A support is disabled, even from an otherwise-trusted local sidecar.

mcpgateway/services/session_affinity.py

@@ -49,6 +49,7 @@
     MessageHandlerFactory,
 )
 from mcpgateway.utils.internal_http import (
+    internal_loopback_base_url,
     post_rpc_in_process,
 )
 
@@ -1021,27 +1022,32 @@ async def _execute_forwarded_request(self, request: Dict[str, Any]) -> Dict[str,
     async def _execute_forwarded_http_request(self, request: Dict[str, Any], redis: Any) -> None:
         """Execute a forwarded Streamable HTTP request in-process and reply via Redis.
 
-        A forwarded request always lands here on the worker that OWNS the
-        downstream session (its ``WORKER_ID`` matched the Redis owner key), so the
-        bound upstream session in this process's ``UpstreamSessionRegistry`` can
-        serve it. We therefore dispatch the JSON-RPC call to the local ``/rpc``
-        route via an **in-process ASGI transport** instead of a network loopback
-        to ``127.0.0.1``: a real loopback hits the shared gunicorn socket and the
-        kernel routes it to an arbitrary worker that does not hold this session,
-        which breaks upstream-session reuse and fails the request. The
-        ``x-forwarded-internally`` header stops the re-entered handler from
-        forwarding again.
+        The request lands on the worker that owns the downstream session, so the
+        bound upstream session in this process can serve it. It is dispatched
+        in-process to the trusted-internal ``/_internal/mcp/rpc`` endpoint rather
+        than the public ``/rpc``.
+
+        Public ``/rpc`` only understands ContextForge JWTs and cookies, so an
+        OAuth bearer or an ``MCP_REQUIRE_AUTH=false`` public-only request would
+        401 if re-authenticated there. Instead, the originating worker's
+        already-validated identity rides in ``auth_context`` (the encoded
+        ``x-contextforge-auth-context`` header); with the
+        ``x-contextforge-mcp-runtime: affinity`` marker, the shared-secret HMAC,
+        and the loopback client address, the trusted endpoint accepts the request
+        and reconstructs the same user the edge established.
 
         Args:
             request: Serialized HTTP request data from Redis Pub/Sub containing:
                 - type: "http_forward"
                 - response_channel: Redis channel to publish response to
                 - mcp_session_id: Session identifier
-                - method: HTTP method (GET, POST, DELETE)
-                - path: Request path (e.g., /servers/{id}/mcp)
-                - headers: Request headers dict (lowercased keys)
-                - body: Hex-encoded request body
-            redis: Redis client for publishing response
+                - method: HTTP method (POST primarily)
+                - path: Original request path (e.g., /servers/{id}/mcp)
+                - headers: Original request headers (lowercased keys)
+                - body: Hex-encoded JSON-RPC request body
+                - auth_context: base64url-encoded auth context from the
+                  originating worker's ``streamable_http_auth()`` result
+            redis: Redis client for publishing the response
         """
         response_channel = request.get("response_channel")
 
@@ -1057,6 +1063,7 @@ async def _publish(status: int, body: bytes, headers: Optional[Dict[str, str]] =
             headers = request.get("headers", {})
             body_hex = request.get("body", "")
             mcp_session_id = request.get("mcp_session_id")
+            auth_context_header = request.get("auth_context") or ""
             body = bytes.fromhex(body_hex) if body_hex else b""
 
             session_short = mcp_session_id[:8] if mcp_session_id and len(mcp_session_id) >= 8 else "unknown"
@@ -1079,7 +1086,8 @@ async def _publish(status: int, body: bytes, headers: Optional[Dict[str, str]] =
                 await _publish(202, b"")
                 return
 
-            # Inject the virtual server id (from the path) into params so /rpc routes correctly.
+            # Inject the virtual server id (from the path) into params so the
+            # dispatcher routes to the right virtual server.
             server_match = _SERVER_ID_RE.search(path)
             if server_match and isinstance(json_body, dict):
                 if not isinstance(json_body.get("params"), dict):
@@ -1088,31 +1096,53 @@ async def _publish(status: int, body: bytes, headers: Optional[Dict[str, str]] =
                 body = orjson.dumps(json_body)
 
             # First-Party - lazy imports avoid a circular dependency with main/transport.
-            # First-Party
+            from mcpgateway.auth_context import _expected_internal_mcp_runtime_auth_header, encode_internal_mcp_auth_context  # pylint: disable=import-outside-toplevel,protected-access
+            from mcpgateway.main import app  # pylint: disable=import-outside-toplevel
             from mcpgateway.utils.passthrough_headers import safe_extract_and_filter_for_loopback  # pylint: disable=import-outside-toplevel
-            from mcpgateway.utils.verify_credentials import _resolve_auth_header_name  # pylint: disable=import-outside-toplevel
-
+            from mcpgateway.utils.verify_credentials import _resolve_auth_header_name  # pylint: disable=import-outside-toplevel,protected-access
+
+            # A forwarded payload may omit the auth context; fall back to an
+            # encoded public-only ({}) context so the endpoint applies default
+            # visibility instead of rejecting the dispatch with 400.
+            if not auth_context_header:
+                auth_context_header = encode_internal_mcp_auth_context({})
+
+            # Trust headers for the internal /_internal/mcp/rpc endpoint:
+            # - x-contextforge-mcp-runtime: "affinity" caller marker
+            # - x-contextforge-mcp-runtime-auth: shared-secret HMAC
+            # - x-contextforge-auth-context: the encoded edge auth context, so the
+            #   endpoint reconstructs the same user without re-authenticating.
             rpc_headers = {
                 "content-type": "application/json",
                 "x-mcp-session-id": mcp_session_id or "",
-                "x-forwarded-internally": "true",
+                "x-contextforge-mcp-runtime": "affinity",
+                "x-contextforge-mcp-runtime-auth": _expected_internal_mcp_runtime_auth_header(),
+                "x-contextforge-auth-context": auth_context_header,
             }
-            auth_header = _resolve_auth_header_name(settings).lower()
-            if auth_header in headers:
-                rpc_headers[auth_header] = headers[auth_header]
+            # Preserve the bearer under the configured auth header (AUTH_HEADER_NAME),
+            # not a hardcoded "authorization": the CSRF bearer short-circuit keys on
+            # the configured header, so a custom header would otherwise be dropped.
+            # This is defense-in-depth; the endpoint trusts the auth-context above.
+            auth_header_name = _resolve_auth_header_name(settings).lower()
+            original_auth = headers.get(auth_header_name) or headers.get(_resolve_auth_header_name(settings))
+            if original_auth:
+                rpc_headers[auth_header_name] = original_auth
             # Preserve passthrough headers destined for upstream MCP servers (#3640).
             rpc_headers.update(safe_extract_and_filter_for_loopback(headers))
 
-            # Dispatch IN-PROCESS so /rpc resolves the bound upstream session from
-            # this worker's registry instead of bouncing back through the shared
-            # socket to a random worker (see post_rpc_in_process).
-            response = await post_rpc_in_process(
-                content=body,
-                headers=rpc_headers,
-                timeout=settings.mcpgateway_pool_rpc_forward_timeout,
-            )
+            # Dispatch IN-PROCESS to the trusted internal endpoint. The explicit
+            # client=("127.0.0.1", 0) tells ASGITransport to set scope["client"]
+            # to a loopback address so the trust check accepts the request.
+            transport = httpx.ASGITransport(app=app, client=("127.0.0.1", 0))
+            async with httpx.AsyncClient(transport=transport, base_url=internal_loopback_base_url()) as client:
+                response = await client.post(
+                    "/_internal/mcp/rpc",
+                    content=body,
+                    headers=rpc_headers,
+                    timeout=settings.mcpgateway_pool_rpc_forward_timeout,
+                )
 
-            logger.debug(f"[HTTP_AFFINITY] Worker {WORKER_ID} | Session {session_short}... | Executed in-process: {response.status_code}")
+            logger.debug(f"[HTTP_AFFINITY] Worker {WORKER_ID} | Session {session_short}... | Executed in-process via /_internal/mcp/rpc: {response.status_code}")
 
             resp_headers = {"content-type": "application/json"}
             if mcp_session_id:
@@ -1149,6 +1179,7 @@ async def forward_to_owner(
         headers: Dict[str, str],
         body: bytes,
         query_string: str = "",
+        auth_context: Optional[str] = None,
     ) -> Optional[Dict[str, Any]]:
         """Forward a Streamable HTTP request to the worker that owns the session via Redis Pub/Sub.
 
@@ -1165,6 +1196,10 @@ async def forward_to_owner(
             headers: Request headers.
             body: Request body bytes.
             query_string: Query string if any.
+            auth_context: Encoded ``x-contextforge-auth-context`` value carrying the
+                originating worker's already-validated identity, so the owner can
+                dispatch to the trusted internal endpoint without re-authenticating
+                (OAuth bearers and ``MCP_REQUIRE_AUTH=false`` public-only survive).
 
         Returns:
             Dict with 'status', 'headers', and 'body' from the owner worker's response,
@@ -1206,6 +1241,8 @@ async def forward_to_owner(
                 "body": body.hex() if body else "",  # Hex encode binary body
                 "original_worker": WORKER_ID,
                 "timestamp": time.time(),
+                # Encoded edge identity; lets the owner dispatch without re-authenticating.
+                "auth_context": auth_context,
             }
 
             # Subscribe to response channel BEFORE publishing request (prevent race)

mcpgateway/transports/rust_mcp_runtime_proxy.py

@@ -15,18 +15,17 @@
 
 # Standard
 import asyncio
-import base64
 import logging
 import re
 from urllib.parse import urlsplit, urlunsplit
 
 # Third-Party
 import httpx
-import orjson
 from sqlalchemy import exists as sa_exists
 from starlette.types import Receive, Scope, Send
 
 # First-Party
+from mcpgateway.auth_context import encode_internal_mcp_auth_context
 from mcpgateway.config import settings
 from mcpgateway.db import fresh_db_session
 from mcpgateway.db import Server as DbServer
@@ -338,5 +337,4 @@ def _build_forwarded_auth_context_header() -> str | None:
     auth_context = get_streamable_http_auth_context()
     if not auth_context:
         return None
-    encoded = base64.urlsafe_b64encode(orjson.dumps(auth_context)).decode("ascii")
-    return encoded.rstrip("=")
+    return encode_internal_mcp_auth_context(auth_context)

mcpgateway/transports/streamablehttp_transport.py

@@ -4209,6 +4209,14 @@ async def handle_streamable_http(  # noqa: PLR0911,PLR0912,PLR0915 — pylint: d
                     # Session owned by another worker - forward the entire HTTP request
                     logger.info("[HTTP_AFFINITY] Worker %s | Session %s... | Owner: %s | Forwarding HTTP request", WORKER_ID, mcp_session_id[:8], owner)
 
+                    # Package the edge-validated identity so the owner can dispatch via
+                    # the trusted internal /_internal/mcp/rpc endpoint without
+                    # re-authenticating, letting OAuth and public-only sessions survive.
+                    # First-Party
+                    from mcpgateway.auth_context import encode_internal_mcp_auth_context  # pylint: disable=import-outside-toplevel
+
+                    encoded_auth_context = encode_internal_mcp_auth_context(get_streamable_http_auth_context())
+
                     # Read request body
                     body_parts = []
                     while True:
@@ -4230,6 +4238,7 @@ async def handle_streamable_http(  # noqa: PLR0911,PLR0912,PLR0915 — pylint: d
                         headers=headers,
                         body=body,
                         query_string=query_string,
+                        auth_context=encoded_auth_context,
                     )
 
                     if response:
@@ -4313,24 +4322,45 @@ async def handle_streamable_http(  # noqa: PLR0911,PLR0912,PLR0915 — pylint: d
                             body = orjson.dumps(json_body)
                             logger.debug("[HTTP_AFFINITY_LOCAL] Injected server_id %s into /rpc params", server_id)
 
+                        # Owner-direct path: dispatch to the trusted internal
+                        # /_internal/mcp/rpc endpoint carrying the edge-validated auth
+                        # context, so OAuth and public-only sessions are honored without
+                        # re-authenticating against public /rpc (JWTs/cookies only).
+                        # First-Party
+                        from mcpgateway.auth_context import _expected_internal_mcp_runtime_auth_header, encode_internal_mcp_auth_context  # pylint: disable=import-outside-toplevel,protected-access
+                        from mcpgateway.main import app  # pylint: disable=import-outside-toplevel
+                        from mcpgateway.utils.internal_http import internal_loopback_base_url  # pylint: disable=import-outside-toplevel
+                        from mcpgateway.utils.passthrough_headers import safe_extract_and_filter_for_loopback  # pylint: disable=import-outside-toplevel
+
                         rpc_headers = {
                             "content-type": "application/json",
                             "x-mcp-session-id": mcp_session_id,
-                            "x-forwarded-internally": "true",
+                            "x-contextforge-mcp-runtime": "affinity",
+                            "x-contextforge-mcp-runtime-auth": _expected_internal_mcp_runtime_auth_header(),
+                            "x-contextforge-auth-context": encode_internal_mcp_auth_context(get_streamable_http_auth_context()),
                         }
-                        _gw_auth_lower = _resolve_auth_header_name(settings).lower()
-                        if _gw_auth_lower in headers:
-                            rpc_headers[_gw_auth_lower] = headers[_gw_auth_lower]
+                        # Preserve the bearer under the configured auth header (AUTH_HEADER_NAME),
+                        # not a hardcoded "authorization": the CSRF bearer short-circuit keys on
+                        # the configured header, so a custom header would otherwise be dropped.
+                        # This is defense-in-depth; the endpoint trusts the auth-context above.
+                        _auth_header_name = _resolve_auth_header_name(settings).lower()
+                        _original_auth = headers.get(_auth_header_name) or headers.get(_resolve_auth_header_name(settings))
+                        if _original_auth:
+                            rpc_headers[_auth_header_name] = _original_auth
                         # Forward passthrough headers for upstream MCP servers (see #3640).
-                        # First-Party
-                        from mcpgateway.utils.passthrough_headers import safe_extract_and_filter_for_loopback  # pylint: disable=import-outside-toplevel
-
                         rpc_headers.update(safe_extract_and_filter_for_loopback(headers))
 
-                        # Dispatch IN-PROCESS so /rpc runs on this worker — the session
-                        # owner that holds the bound upstream session — instead of looping
-                        # back over the shared socket to a random worker (#4205 isolation).
-                        response = await post_rpc_in_process(content=body, headers=rpc_headers, timeout=30.0)
+                        # Dispatch in-process so the request runs on this worker, the
+                        # session owner that holds the bound upstream session, instead of
+                        # looping back over the shared socket to a random worker.
+                        transport = httpx.ASGITransport(app=app, client=("127.0.0.1", 0))
+                        async with httpx.AsyncClient(transport=transport, base_url=internal_loopback_base_url()) as client:
+                            response = await client.post(
+                                "/_internal/mcp/rpc",
+                                content=body,
+                                headers=rpc_headers,
+                                timeout=settings.mcpgateway_pool_rpc_forward_timeout,
+                            )
 
                         response_headers = [
                             (b"content-type", b"application/json"),