From 99ac7f0e97d50318cb6e05f44a2135b39fac5f78 Mon Sep 17 00:00:00 2001
From: GareArc <chen4851@purdue.edu>
Date: Wed, 28 May 2025 17:05:36 +0800
Subject: [PATCH] feat: introduce enterprise login token

---
 api/controllers/web/passport.py     | 75 +++++++++++++++++++++++++++++
 api/services/webapp_auth_service.py |  2 +-
 2 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/api/controllers/web/passport.py b/api/controllers/web/passport.py
index 154c6772b7..a9e70ca1ec 100644
--- a/api/controllers/web/passport.py
+++ b/api/controllers/web/passport.py
@@ -1,9 +1,11 @@
 import uuid
+from datetime import UTC, datetime, timedelta
 
 from flask import request
 from flask_restful import Resource
 from werkzeug.exceptions import NotFound, Unauthorized
 
+from configs import dify_config
 from controllers.web import api
 from controllers.web.error import WebAppAuthRequiredError
 from extensions.ext_database import db
@@ -20,10 +22,17 @@ class PassportResource(Resource):
         system_features = FeatureService.get_system_features()
         app_code = request.headers.get("X-App-Code")
         user_id = request.args.get("user_id")
+        enterprise_login_token = request.args.get("enterprise_login_token")
 
         if app_code is None:
             raise Unauthorized("X-App-Code header is missing.")
 
+        # logic for exchange token for enterprise logined web user
+        enterprise_user_id = decode_enterprise_webapp_user_id(enterprise_login_token)
+        if enterprise_user_id:
+            # a web user has already logged in, exchange a token for this app without redirecting to the login page
+            return exchange_token_for_existing_web_user(app_code=app_code, user_id=enterprise_user_id)
+
         if system_features.webapp_auth.enabled:
             app_settings = EnterpriseService.WebAppAuth.get_app_access_mode_by_code(app_code=app_code)
             if not app_settings or not app_settings.access_mode == "public":
@@ -84,6 +93,72 @@ class PassportResource(Resource):
 api.add_resource(PassportResource, "/passport")
 
 
+def decode_enterprise_webapp_user_id(auth_header: str | None):
+    """
+    Decode the enterprise user session from the Authorization header.
+    """
+    if not auth_header:
+        return None
+
+    if " " not in auth_header:
+        raise Unauthorized("Invalid Authorization header format. Expected 'Bearer <api-key>' format.")
+
+    auth_scheme, tk = auth_header.split(None, 1)
+    auth_scheme = auth_scheme.lower()
+
+    if auth_scheme != "bearer":
+        raise Unauthorized("Invalid Authorization header format. Expected 'Bearer <api-key>' format.")
+
+    decoded = PassportService().verify(tk)
+    source = decoded.get("token_source")
+    if not source or source != "enterprise_login":
+        return None
+    user_id: str | None = decoded.get("user_id")
+    return user_id
+
+
+def exchange_token_for_existing_web_user(app_code: str, user_id: str):
+    """
+    Exchange a token for an existing web user session.
+    """
+
+    site = db.session.query(Site).filter(Site.code == app_code, Site.status == "normal").first()
+    if not site:
+        raise NotFound()
+
+    app_model = db.session.query(App).filter(App.id == site.app_id).first()
+    if not app_model or app_model.status != "normal" or not app_model.enable_site:
+        raise NotFound()
+    end_user = db.session.query(EndUser).filter(EndUser.app_id == app_model.id, EndUser.session_id == user_id).first()
+    if not end_user:
+        end_user = EndUser(
+            tenant_id=app_model.tenant_id,
+            app_id=app_model.id,
+            type="browser",
+            is_anonymous=True,
+            session_id=generate_session_id(),
+        )
+        db.session.add(end_user)
+        db.session.commit()
+
+    exp_dt = datetime.now(UTC) + timedelta(hours=dify_config.ACCESS_TOKEN_EXPIRE_MINUTES * 24)
+    exp = int(exp_dt.timestamp())
+    payload = {
+        "iss": site.id,
+        "sub": "Web API Passport",
+        "app_id": site.app_id,
+        "app_code": site.code,
+        "user_id": user_id,
+        "end_user_id": end_user.id,
+        "token_source": "webapp",
+        "exp": exp,
+    }
+    token: str = PassportService().issue(payload)
+    return {
+        "access_token": token,
+    }
+
+
 def generate_session_id():
     """
     Generate a unique session ID.
diff --git a/api/services/webapp_auth_service.py b/api/services/webapp_auth_service.py
index 79d5217de7..f4186e7f6f 100644
--- a/api/services/webapp_auth_service.py
+++ b/api/services/webapp_auth_service.py
@@ -133,7 +133,7 @@ class WebAppAuthService:
             "app_code": site.code,
             "user_id": account.id,
             "end_user_id": end_user_id,
-            "token_source": "webapp",
+            "token_source": "enterprise_login",
             "exp": exp,
         }