From ff4d215f78a8d6e6db6052be18b5f9f4db9c7d68 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Yingyi=20/=20=E9=A2=96=E9=80=B8?=
 <49649786+Zuoqiu-Yingyi@users.noreply.github.com>
Date: Thu, 28 Aug 2025 16:20:12 +0800
Subject: [PATCH] :art: Add cookie-based auth in publish proxy (#15692)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* chore(publish-auth): Add TODO for cookie-based auth in publish proxy

A TODO comment was added to indicate future implementation of authentication using cookies in the PublishServiceTransport RoundTrip method.

* :art: Add session-based authentication for publish proxy

Introduces session management using cookies for the publish reverse proxy server. Adds session ID generation, storage, and validation in kernel/model/auth.go, and updates the proxy transport to check for valid sessions before falling back to basic authentication. Sets a session cookie upon successful basic auth login.

* :bug: Fixed the issue of repeatedly setting cookies

* :art: Dynamically remove invalid session IDs

* ♻️ Revert changes in pnpm-lock.yaml
---
 kernel/go.mod                  |  2 +-
 kernel/model/auth.go           | 33 ++++++++++++++++++++++++++--
 kernel/server/proxy/publish.go | 39 ++++++++++++++++++++++++++++++----
 3 files changed, 67 insertions(+), 7 deletions(-)

diff --git a/kernel/go.mod b/kernel/go.mod
index 6981bccd8..db76b4e11 100644
--- a/kernel/go.mod
+++ b/kernel/go.mod
@@ -37,6 +37,7 @@ require (
 	github.com/go-ole/go-ole v1.3.0
 	github.com/gofrs/flock v0.12.1
 	github.com/golang-jwt/jwt/v5 v5.2.2
+	github.com/google/uuid v1.6.0
 	github.com/gorilla/css v1.0.1
 	github.com/gorilla/websocket v1.5.3
 	github.com/imroc/req/v3 v3.54.2
@@ -129,7 +130,6 @@ require (
 	github.com/goccy/go-json v0.10.5 // indirect
 	github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 // indirect
 	github.com/google/pprof v0.0.0-20250607225305-033d6d78b36a // indirect
-	github.com/google/uuid v1.6.0 // indirect
 	github.com/gopherjs/gopherjs v1.17.2 // indirect
 	github.com/gorilla/context v1.1.2 // indirect
 	github.com/gorilla/securecookie v1.1.2 // indirect
diff --git a/kernel/model/auth.go b/kernel/model/auth.go
index e1551ff70..90f62c82e 100644
--- a/kernel/model/auth.go
+++ b/kernel/model/auth.go
@@ -19,8 +19,10 @@ package model
 import (
 	"crypto/rand"
 	"net/http"
+	"sync"
 
 	"github.com/golang-jwt/jwt/v5"
+	"github.com/google/uuid"
 	"github.com/siyuan-note/logging"
 )
 
@@ -29,12 +31,15 @@ type Account struct {
 	Password string
 	Token    string
 }
-type AccountsMap map[string]*Account
+type AccountsMap map[string]*Account // username -> account
+type SessionsMap map[string]string   // sessionID -> username
 type ClaimsKeyType string
 
 const (
 	XAuthTokenKey = "X-Auth-Token"
 
+	SessionIdCookieName = "publish-visitor-session-id"
+
 	ClaimsContextKey = "claims"
 
 	iss = "siyuan-publish-reverse-proxy-server"
@@ -46,13 +51,37 @@ const (
 
 var (
 	accountsMap = AccountsMap{}
-	jwtKey      = make([]byte, 32)
+	sessionsMap = SessionsMap{}
+	sessionLock = sync.Mutex{}
+
+	jwtKey = make([]byte, 32)
 )
 
 func GetBasicAuthAccount(username string) *Account {
 	return accountsMap[username]
 }
 
+func GetBasicAuthUsernameBySessionID(sessionID string) string {
+	return sessionsMap[sessionID]
+}
+
+func GetNewSessionID() string {
+	sessionID := uuid.New().String()
+	return sessionID
+}
+
+func AddSession(sessionID, username string) {
+	sessionLock.Lock()
+	defer sessionLock.Unlock()
+	sessionsMap[sessionID] = username
+}
+
+func DeleteSession(sessionID string) {
+	sessionLock.Lock()
+	defer sessionLock.Unlock()
+	delete(sessionsMap, sessionID)
+}
+
 func InitAccounts() {
 	accountsMap = AccountsMap{
 		"": &Account{}, // 匿名用户
diff --git a/kernel/server/proxy/publish.go b/kernel/server/proxy/publish.go
index edac7eeb6..b350bdccf 100644
--- a/kernel/server/proxy/publish.go
+++ b/kernel/server/proxy/publish.go
@@ -125,10 +125,28 @@ func rewrite(r *httputil.ProxyRequest) {
 
 func (PublishServiceTransport) RoundTrip(request *http.Request) (response *http.Response, err error) {
 	if model.Conf.Publish.Auth.Enable {
+		// Session Auth
+		sessionIdCookie, cookieErr := request.Cookie(model.SessionIdCookieName)
+		if cookieErr == nil {
+			// Check session ID
+			sessionID := sessionIdCookie.Value
+			if username := model.GetBasicAuthUsernameBySessionID(sessionID); username != "" {
+				// Valid session
+				if account := model.GetBasicAuthAccount(username); account != nil {
+					// Valid account
+					request.Header.Set(model.XAuthTokenKey, account.Token)
+					response, err = http.DefaultTransport.RoundTrip(request)
+					return
+				} else {
+					// Invalid account, remove session
+					model.DeleteSession(sessionID)
+				}
+			}
+		}
+
 		// Basic Auth
 		username, password, ok := request.BasicAuth()
 		account := model.GetBasicAuthAccount(username)
-
 		if !ok ||
 			account == nil ||
 			account.Username == "" || // 匿名用户
@@ -149,13 +167,26 @@ func (PublishServiceTransport) RoundTrip(request *http.Request) (response *http.
 				ContentLength: -1,
 			}, nil
 		} else {
+			// set session cookie
+			sessionID := model.GetNewSessionID()
+			cookie := &http.Cookie{
+				Name:     model.SessionIdCookieName,
+				Value:    sessionID,
+				Path:     "/",
+				HttpOnly: true,
+			}
+			model.AddSession(sessionID, username)
+
 			// set JWT
 			request.Header.Set(model.XAuthTokenKey, account.Token)
+			response, err = http.DefaultTransport.RoundTrip(request)
+
+			response.Header.Add("Set-Cookie", cookie.String())
+			return
 		}
 	} else {
 		request.Header.Set(model.XAuthTokenKey, model.GetBasicAuthAccount("").Token)
+		response, err = http.DefaultTransport.RoundTrip(request)
+		return
 	}
-
-	response, err = http.DefaultTransport.RoundTrip(request)
-	return
 }