refactor: replace localStorage with HTTP-only cookies for auth tokens (#24365)

Signed-off-by: NeatGuyCoding <15627489+NeatGuyCoding@users.noreply.github.com>
Signed-off-by: lyzno1 <yuanyouhuilyz@gmail.com>
Signed-off-by: kenwoodjw <blackxin55+@gmail.com>
Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Yunlu Wen <wylswz@163.com>
Co-authored-by: Joel <iamjoel007@gmail.com>
Co-authored-by: GareArc <chen4851@purdue.edu>
Co-authored-by: NFish <douxc512@gmail.com>
Co-authored-by: Davide Delbianco <davide.delbianco@outlook.com>
Co-authored-by: minglu7 <1347866672@qq.com>
Co-authored-by: Ponder <ruan.lj@foxmail.com>
Co-authored-by: crazywoola <100913391+crazywoola@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: heyszt <270985384@qq.com>
Co-authored-by: Asuka Minato <i@asukaminato.eu.org>
Co-authored-by: Guangdong Liu <liugddx@gmail.com>
Co-authored-by: Eric Guo <eric.guocz@gmail.com>
Co-authored-by: NeatGuyCoding <15627489+NeatGuyCoding@users.noreply.github.com>
Co-authored-by: XlKsyt <caixuesen@outlook.com>
Co-authored-by: Dhruv Gorasiya <80987415+DhruvGorasiya@users.noreply.github.com>
Co-authored-by: crazywoola <427733928@qq.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: lyzno1 <92089059+lyzno1@users.noreply.github.com>
Co-authored-by: hj24 <mambahj24@gmail.com>
Co-authored-by: GuanMu <ballmanjq@gmail.com>
Co-authored-by: 非法操作 <hjlarry@163.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Tonlo <123lzs123@gmail.com>
Co-authored-by: Yusuke Yamada <yamachu.dev@gmail.com>
Co-authored-by: Novice <novice12185727@gmail.com>
Co-authored-by: kenwoodjw <blackxin55+@gmail.com>
Co-authored-by: Ademílson Tonato <ademilsonft@outlook.com>
Co-authored-by: znn <jubinkumarsoni@gmail.com>
Co-authored-by: yangzheli <43645580+yangzheli@users.noreply.github.com>

api/libs/external_api.py

@@ -9,7 +9,9 @@ from werkzeug.exceptions import HTTPException
 from werkzeug.http import HTTP_STATUS_CODES
 
 from configs import dify_config
+from constants import COOKIE_NAME_ACCESS_TOKEN, COOKIE_NAME_CSRF_TOKEN, COOKIE_NAME_REFRESH_TOKEN
 from core.errors.error import AppInvokeQuotaExceededError
+from libs.token import is_secure
 
 
 def http_status_message(code):
@@ -67,6 +69,19 @@ def register_external_error_handlers(api: Api):
             # If you need WWW-Authenticate for 401, add it to headers
             if status_code == 401:
                 headers["WWW-Authenticate"] = 'Bearer realm="api"'
+                # Check if this is a forced logout error - clear cookies
+                error_code = getattr(e, "error_code", None)
+                if error_code == "unauthorized_and_force_logout":
+                    # Add Set-Cookie headers to clear auth cookies
+
+                    secure = is_secure()
+                    # response is not accessible, so we need to do it ugly
+                    common_part = "Path=/; Expires=Thu, 01 Jan 1970 00:00:00 GMT; HttpOnly"
+                    headers["Set-Cookie"] = [
+                        f'{COOKIE_NAME_ACCESS_TOKEN}=""; {common_part}{"; Secure" if secure else ""}; SameSite=Lax',
+                        f'{COOKIE_NAME_CSRF_TOKEN}=""; {common_part}{"; Secure" if secure else ""}; SameSite=Lax',
+                        f'{COOKIE_NAME_REFRESH_TOKEN}=""; {common_part}{"; Secure" if secure else ""}; SameSite=Lax',
+                    ]
             return data, status_code, headers
 
     _ = handle_http_exception

api/libs/login.py

@@ -7,6 +7,7 @@ from flask_login.config import EXEMPT_METHODS  # type: ignore
 from werkzeug.local import LocalProxy
 
 from configs import dify_config
+from libs.token import check_csrf_token
 from models import Account
 from models.model import EndUser
 
@@ -73,6 +74,9 @@ def login_required(func: Callable[P, R]):
             pass
         elif current_user is not None and not current_user.is_authenticated:
             return current_app.login_manager.unauthorized()  # type: ignore
+        # we put csrf validation here for less conflicts
+        # TODO: maybe find a better place for it.
+        check_csrf_token(request, current_user.id)
         return current_app.ensure_sync(func)(*args, **kwargs)
 
     return decorated_view

api/libs/token.py

@@ -0,0 +1,208 @@
+import logging
+import re
+from datetime import UTC, datetime, timedelta
+
+from flask import Request
+from werkzeug.exceptions import Unauthorized
+from werkzeug.wrappers import Response
+
+from configs import dify_config
+from constants import (
+    COOKIE_NAME_ACCESS_TOKEN,
+    COOKIE_NAME_CSRF_TOKEN,
+    COOKIE_NAME_PASSPORT,
+    COOKIE_NAME_REFRESH_TOKEN,
+    HEADER_NAME_CSRF_TOKEN,
+    HEADER_NAME_PASSPORT,
+)
+from libs.passport import PassportService
+
+logger = logging.getLogger(__name__)
+
+CSRF_WHITE_LIST = [
+    re.compile(r"/console/api/apps/[a-f0-9-]+/workflows/draft"),
+]
+
+
+# server is behind a reverse proxy, so we need to check the url
+def is_secure() -> bool:
+    return dify_config.CONSOLE_WEB_URL.startswith("https") and dify_config.CONSOLE_API_URL.startswith("https")
+
+
+def _real_cookie_name(cookie_name: str) -> str:
+    if is_secure():
+        return "__Host-" + cookie_name
+    else:
+        return cookie_name
+
+
+def _try_extract_from_header(request: Request) -> str | None:
+    """
+    Try to extract access token from header
+    """
+    auth_header = request.headers.get("Authorization")
+    if auth_header:
+        if " " not in auth_header:
+            return None
+        else:
+            auth_scheme, auth_token = auth_header.split(None, 1)
+            auth_scheme = auth_scheme.lower()
+            if auth_scheme != "bearer":
+                return None
+            else:
+                return auth_token
+    return None
+
+
+def extract_csrf_token(request: Request) -> str | None:
+    """
+    Try to extract CSRF token from header or cookie.
+    """
+    return request.headers.get(HEADER_NAME_CSRF_TOKEN)
+
+
+def extract_csrf_token_from_cookie(request: Request) -> str | None:
+    """
+    Try to extract CSRF token from cookie.
+    """
+    return request.cookies.get(_real_cookie_name(COOKIE_NAME_CSRF_TOKEN))
+
+
+def extract_access_token(request: Request) -> str | None:
+    """
+    Try to extract access token from cookie, header or params.
+
+    Access token is either for console session or webapp passport exchange.
+    """
+
+    def _try_extract_from_cookie(request: Request) -> str | None:
+        return request.cookies.get(_real_cookie_name(COOKIE_NAME_ACCESS_TOKEN))
+
+    return _try_extract_from_cookie(request) or _try_extract_from_header(request)
+
+
+def extract_webapp_passport(app_code: str, request: Request) -> str | None:
+    """
+    Try to extract app token from header or params.
+
+    Webapp access token (part of passport) is only used for webapp session.
+    """
+
+    def _try_extract_passport_token_from_cookie(request: Request) -> str | None:
+        return request.cookies.get(_real_cookie_name(COOKIE_NAME_PASSPORT + "-" + app_code))
+
+    def _try_extract_passport_token_from_header(request: Request) -> str | None:
+        return request.headers.get(HEADER_NAME_PASSPORT)
+
+    ret = _try_extract_passport_token_from_cookie(request) or _try_extract_passport_token_from_header(request)
+    return ret
+
+
+def set_access_token_to_cookie(request: Request, response: Response, token: str, samesite: str = "Lax"):
+    response.set_cookie(
+        _real_cookie_name(COOKIE_NAME_ACCESS_TOKEN),
+        value=token,
+        httponly=True,
+        secure=is_secure(),
+        samesite=samesite,
+        max_age=int(dify_config.ACCESS_TOKEN_EXPIRE_MINUTES * 60),
+        path="/",
+    )
+
+
+def set_refresh_token_to_cookie(request: Request, response: Response, token: str):
+    response.set_cookie(
+        _real_cookie_name(COOKIE_NAME_REFRESH_TOKEN),
+        value=token,
+        httponly=True,
+        secure=is_secure(),
+        samesite="Lax",
+        max_age=int(60 * 60 * 24 * dify_config.REFRESH_TOKEN_EXPIRE_DAYS),
+        path="/",
+    )
+
+
+def set_csrf_token_to_cookie(request: Request, response: Response, token: str):
+    response.set_cookie(
+        _real_cookie_name(COOKIE_NAME_CSRF_TOKEN),
+        value=token,
+        httponly=False,
+        secure=is_secure(),
+        samesite="Lax",
+        max_age=int(60 * dify_config.ACCESS_TOKEN_EXPIRE_MINUTES),
+        path="/",
+    )
+
+
+def _clear_cookie(
+    response: Response,
+    cookie_name: str,
+    samesite: str = "Lax",
+    http_only: bool = True,
+):
+    response.set_cookie(
+        _real_cookie_name(cookie_name),
+        "",
+        expires=0,
+        path="/",
+        secure=is_secure(),
+        httponly=http_only,
+        samesite=samesite,
+    )
+
+
+def clear_access_token_from_cookie(response: Response, samesite: str = "Lax"):
+    _clear_cookie(response, COOKIE_NAME_ACCESS_TOKEN, samesite)
+
+
+def clear_refresh_token_from_cookie(response: Response):
+    _clear_cookie(response, COOKIE_NAME_REFRESH_TOKEN)
+
+
+def clear_csrf_token_from_cookie(response: Response):
+    _clear_cookie(response, COOKIE_NAME_CSRF_TOKEN, http_only=False)
+
+
+def check_csrf_token(request: Request, user_id: str):
+    # some apis are sent by beacon, so we need to bypass csrf token check
+    # since these APIs are post, they are already protected by SameSite: Lax, so csrf is not required.
+    def _unauthorized():
+        raise Unauthorized("CSRF token is missing or invalid.")
+
+    for pattern in CSRF_WHITE_LIST:
+        if pattern.match(request.path):
+            return
+
+    csrf_token = extract_csrf_token(request)
+    csrf_token_from_cookie = extract_csrf_token_from_cookie(request)
+
+    if csrf_token != csrf_token_from_cookie:
+        _unauthorized()
+
+    if not csrf_token:
+        _unauthorized()
+    verified = {}
+    try:
+        verified = PassportService().verify(csrf_token)
+    except:
+        _unauthorized()
+
+    if verified.get("sub") != user_id:
+        _unauthorized()
+
+    exp: int | None = verified.get("exp")
+    if not exp:
+        _unauthorized()
+    else:
+        time_now = int(datetime.now().timestamp())
+        if exp < time_now:
+            _unauthorized()
+
+
+def generate_csrf_token(user_id: str) -> str:
+    exp_dt = datetime.now(UTC) + timedelta(minutes=dify_config.ACCESS_TOKEN_EXPIRE_MINUTES)
+    payload = {
+        "exp": int(exp_dt.timestamp()),
+        "sub": user_id,
+    }
+    return PassportService().issue(payload)