From 8cc7a6888473b3150f63434e5efae134cbb1c5b1 Mon Sep 17 00:00:00 2001
From: Jeongseok Kang <jskang@lablup.com>
Date: Tue, 25 Jul 2023 00:50:25 +0900
Subject: [PATCH] feat: Issue a signed token for authorizing the pipeline
 service (#1350)

---
 changes/1350.feature.md       |  1 +
 configs/webserver/sample.conf |  1 +
 src/ai/backend/web/proxy.py   | 23 ++++++++++++++++++++++-
 3 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 changes/1350.feature.md

diff --git a/changes/1350.feature.md b/changes/1350.feature.md
new file mode 100644
index 0000000000..6adae5b8ea
--- /dev/null
+++ b/changes/1350.feature.md
@@ -0,0 +1 @@
+Issue a signed token to X-BackendAI-SSO header to authorize an user from the pipeline service
diff --git a/configs/webserver/sample.conf b/configs/webserver/sample.conf
index ed473def1a..1b3bc8cb77 100644
--- a/configs/webserver/sample.conf
+++ b/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"
diff --git a/src/ai/backend/web/proxy.py b/src/ai/backend/web/proxy.py
index 6751a63ca0..079113d167 100644
--- a/src/ai/backend/web/proxy.py
+++ b/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: