feat: Issue a signed token for authorizing the pipeline service

changes/1350.feature.md

@@ -0,0 +1 @@
+Issue a signed token to X-BackendAI-SSO header to authorize an user from the pipeline service

configs/webserver/sample.conf

@@ -94,6 +94,7 @@ max_file_upload_size = 4294967296
 
 [pipeline]
 #endpoint = "http://mlops.com:9500"
+jwt.secret = "7<:~[X,^Z1XM!*,Pe:PHR!bv,H~Q#l177<7gf_XHD6.<*<.t<[o|V5W(=0x:jTh-"
 
 [ui]
 brand = "Lablup Cloud"

src/ai/backend/web/proxy.py

@@ -5,9 +5,11 @@
 import json
 import logging
 import random
+from datetime import datetime, timedelta, timezone
 from typing import Optional, Tuple, Union, cast
 
 import aiohttp
+import jwt
 from aiohttp import web
 from Crypto.Cipher import AES
 from Crypto.Util.Padding import unpad
@@ -158,7 +160,7 @@ async def web_handler(request: web.Request, *, is_anonymous=False) -> web.Stream
             log.error("WEB_HANDLER: 'pipeline.endpoint' has not been set.")
         else:
             log.info(f"WEB_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))
     else:
@@ -201,6 +203,25 @@ async def web_handler(request: web.Request, *, is_anonymous=False) -> web.Stream
             for hdr in HTTP_HEADERS_TO_FORWARD:
                 if request.headers.get(hdr) is not None:
                     api_rqst.headers[hdr] = request.headers[hdr]
+            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(tz=timezone(timedelta(hours=9)))
+                    payload = {
+                        # Registered claims
+                        "exp": now + timedelta(hours=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")
+                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
             # Uploading request body happens at the entering of the block,
             # and downloading response body happens in the read loop inside.
             async with api_rqst.fetch() as up_resp: