feat: add authentication headers when proxying pipeline websocket requests

changes/1457.feature.md

@@ -0,0 +1 @@
+Add `X-BackendAI-SessionID` and `X-BackendAI-SSO` headers to the pipeline websocket proxy handler.

src/ai/backend/manager/registry.py

@@ -2560,7 +2560,7 @@ async def _restart_kernel(kernel: KernelRow) -> None:
                     keepalive_timeout=self.rpc_keepalive_timeout,
                 ) as rpc:
                     updated_config: Dict[str, Any] = {
-                        # TODO: support resacling of sub-containers
+                        # TODO: support rescaling of sub-containers
                     }
                     kernel_info = await rpc.call.restart_kernel(
                         str(kernel.session_id),

src/ai/backend/web/config.py

@@ -86,7 +86,12 @@
         t.Key("pipeline"): t.Dict(
             {
                 t.Key("endpoint", default=None): t.Null | tx.URL,
-            }
+                t.Key("jwt"): t.Dict(
+                    {
+                        t.Key("secret"): t.String,
+                    },
+                ),
+            },
         ).allow_extra("*"),
         t.Key("ui"): t.Dict(
             {

src/ai/backend/web/proxy.py

@@ -5,7 +5,7 @@
 import json
 import logging
 import random
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from typing import Optional, Tuple, Union, cast
 
 import aiohttp
@@ -207,18 +207,18 @@ async def web_handler(request: web.Request, *, is_anonymous=False) -> web.Stream
                 aiohttp_session = request.cookies.get("AIOHTTP_SESSION")
                 if not (sso_token := request.headers.get("X-BackendAI-SSO")):
                     jwt_secret = request.app["config"]["pipeline"]["jwt"]["secret"]
-                    now = datetime.now(tz=timezone(timedelta(hours=9)))
-                    payload = {
+                    now = datetime.now().astimezone()
+                    claims = {
                         # Registered claims
-                        "exp": now + timedelta(hours=1),
+                        "exp": now + timedelta(days=1),
                         "iss": "Backend.AI Webserver",
                         "iat": now,
                         # Private claims
                         "aiohttp_session": aiohttp_session,
                         "access_key": api_session.config.access_key,
                         # "secret_key": api_session.config.secret_key,
                     }
-                    sso_token = jwt.encode(payload, key=jwt_secret, algorithm="HS256")
+                    sso_token = jwt.encode(claims, key=jwt_secret, algorithm="HS256")
                 api_rqst.headers["X-BackendAI-SSO"] = sso_token
                 if session_id := (request_headers.get("X-BackendAI-SessionID") or aiohttp_session):
                     api_rqst.headers["X-BackendAI-SessionID"] = session_id
@@ -382,13 +382,14 @@ async def websocket_handler(request, *, is_anonymous=False) -> web.StreamRespons
         else:
             endpoint = endpoint.with_scheme("ws")
             log.info(f"WEBSOCKET_HANDLER {request.path} -> {endpoint}/{real_path}")
-        api_session = await asyncio.shield(get_anonymous_session(request, endpoint))
+        api_session = await asyncio.shield(get_api_session(request, endpoint))
     elif is_anonymous:
         api_session = await asyncio.shield(get_anonymous_session(request, api_endpoint))
     else:
         api_session = await asyncio.shield(get_api_session(request, api_endpoint))
     try:
         async with api_session:
+            request_headers = extra_config_headers.check(request.headers)
             request_api_version = request.headers.get("X-BackendAI-Version", None)
             fill_forwarding_hdrs_to_api_session(request, api_session)
             api_rqst = Request(
@@ -399,6 +400,25 @@ async def websocket_handler(request, *, is_anonymous=False) -> web.StreamRespons
                 content_type=request.content_type,
                 override_api_version=request_api_version,
             )
+            if proxy_path == "pipeline":
+                aiohttp_session = request.cookies.get("AIOHTTP_SESSION")
+                if not (sso_token := request.headers.get("X-BackendAI-SSO")):
+                    jwt_secret = request.app["config"]["pipeline"]["jwt"]["secret"]
+                    now = datetime.now().astimezone()
+                    claims = {
+                        # Registered claims
+                        "exp": now + timedelta(days=1),
+                        "iss": "Backend.AI Webserver",
+                        "iat": now,
+                        # Private claims
+                        "aiohttp_session": aiohttp_session,
+                        "access_key": api_session.config.access_key,
+                        # "secret_key": api_session.config.secret_key,
+                    }
+                    sso_token = jwt.encode(claims, key=jwt_secret, algorithm="HS256")
+                api_rqst.headers["X-BackendAI-SSO"] = sso_token
+                if session_id := (request_headers.get("X-BackendAI-SessionID") or aiohttp_session):
+                    api_rqst.headers["X-BackendAI-SessionID"] = session_id
             async with api_rqst.connect_websocket() as up_conn:
                 down_conn = web.WebSocketResponse()
                 await down_conn.prepare(request)