From 4e9e111ee7809c4a2cd17acb08da4e08bd0c7708 Mon Sep 17 00:00:00 2001
From: Liang Ding <d@b3log.org>
Date: Tue, 10 Jan 2023 22:25:02 +0800
Subject: [PATCH] =?UTF-8?q?:art:=20=E6=94=B9=E8=BF=9B=E5=A4=9A=E5=B7=A5?=
 =?UTF-8?q?=E4=BD=9C=E7=A9=BA=E9=97=B4=E9=89=B4=E6=9D=83?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 kernel/api/system.go    |  3 ++-
 kernel/model/session.go | 17 ++++++++++-------
 kernel/server/serve.go  |  7 ++++---
 kernel/util/session.go  | 17 +++++++++++++++++
 4 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/kernel/api/system.go b/kernel/api/system.go
index 0d8378a03..6fe77a9e7 100644
--- a/kernel/api/system.go
+++ b/kernel/api/system.go
@@ -208,7 +208,8 @@ func setAccessAuthCode(c *gin.Context) {
 	model.Conf.Save()
 
 	session := util.GetSession(c)
-	session.AccessAuthCode = aac
+	workspaceSession := util.GetWorkspaceSession(session)
+	workspaceSession.AccessAuthCode = aac
 	session.Save(c)
 	go func() {
 		time.Sleep(200 * time.Millisecond)
diff --git a/kernel/model/session.go b/kernel/model/session.go
index b4607454c..67c319d70 100644
--- a/kernel/model/session.go
+++ b/kernel/model/session.go
@@ -63,6 +63,7 @@ func LoginAuth(c *gin.Context) {
 
 	var inputCaptcha string
 	session := util.GetSession(c)
+	workspaceSession := util.GetWorkspaceSession(session)
 	if util.NeedCaptcha() {
 		captchaArg := arg["captcha"]
 		if nil == captchaArg {
@@ -77,7 +78,7 @@ func LoginAuth(c *gin.Context) {
 			return
 		}
 
-		if strings.ToLower(session.Captcha) != strings.ToLower(inputCaptcha) {
+		if strings.ToLower(workspaceSession.Captcha) != strings.ToLower(inputCaptcha) {
 			ret.Code = 1
 			ret.Msg = Conf.Language(22)
 			return
@@ -90,7 +91,7 @@ func LoginAuth(c *gin.Context) {
 		ret.Msg = Conf.Language(83)
 
 		util.WrongAuthCount++
-		session.Captcha = gulu.Rand.String(7)
+		workspaceSession.Captcha = gulu.Rand.String(7)
 		if util.NeedCaptcha() {
 			ret.Code = 1 // 需要渲染验证码
 		}
@@ -103,9 +104,9 @@ func LoginAuth(c *gin.Context) {
 		return
 	}
 
-	session.AccessAuthCode = authCode
+	workspaceSession.AccessAuthCode = authCode
 	util.WrongAuthCount = 0
-	session.Captcha = gulu.Rand.String(7)
+	workspaceSession.Captcha = gulu.Rand.String(7)
 	if err := session.Save(c); nil != err {
 		logging.LogErrorf("save session failed: " + err.Error())
 		c.Status(500)
@@ -126,7 +127,8 @@ func GetCaptcha(c *gin.Context) {
 	}
 
 	session := util.GetSession(c)
-	session.Captcha = img.Text
+	workspaceSession := util.GetWorkspaceSession(session)
+	workspaceSession.Captcha = img.Text
 	if err = session.Save(c); nil != err {
 		logging.LogErrorf("save session failed: " + err.Error())
 		c.Status(500)
@@ -186,7 +188,8 @@ func CheckAuth(c *gin.Context) {
 
 	// 通过 Cookie
 	session := util.GetSession(c)
-	if session.AccessAuthCode == Conf.AccessAuthCode {
+	workspaceSession := util.GetWorkspaceSession(session)
+	if workspaceSession.AccessAuthCode == Conf.AccessAuthCode {
 		c.Next()
 		return
 	}
@@ -211,7 +214,7 @@ func CheckAuth(c *gin.Context) {
 		return
 	}
 
-	if session.AccessAuthCode != Conf.AccessAuthCode {
+	if workspaceSession.AccessAuthCode != Conf.AccessAuthCode {
 		userAgentHeader := c.GetHeader("User-Agent")
 		if strings.HasPrefix(userAgentHeader, "SiYuan/") || strings.HasPrefix(userAgentHeader, "Mozilla/") {
 			if "GET" != c.Request.Method {
diff --git a/kernel/server/serve.go b/kernel/server/serve.go
index 90a3e55c7..8b67e1e63 100644
--- a/kernel/server/serve.go
+++ b/kernel/server/serve.go
@@ -364,13 +364,14 @@ func serveWebSocket(ginServer *gin.Engine) {
 				if nil == val {
 					authOk = false
 				} else {
-					sess := map[string]interface{}{}
-					err = gulu.JSON.UnmarshalJSON([]byte(val.(string)), &sess)
+					sess := &util.SessionData{}
+					err = gulu.JSON.UnmarshalJSON([]byte(val.(string)), sess)
 					if nil != err {
 						authOk = false
 						logging.LogErrorf("unmarshal cookie failed: %s", err)
 					} else {
-						authOk = sess["AccessAuthCode"].(string) == model.Conf.AccessAuthCode
+						workspaceSess := util.GetWorkspaceSession(sess)
+						authOk = workspaceSess.AccessAuthCode == model.Conf.AccessAuthCode
 					}
 				}
 			}
diff --git a/kernel/util/session.go b/kernel/util/session.go
index 889950284..71e85ea50 100644
--- a/kernel/util/session.go
+++ b/kernel/util/session.go
@@ -30,6 +30,10 @@ func NeedCaptcha() bool {
 
 // SessionData represents the session.
 type SessionData struct {
+	Workspaces map[string]*WorkspaceSession // <WorkspacePath, WorkspaceSession>
+}
+
+type WorkspaceSession struct {
 	AccessAuthCode string
 	Captcha        string
 }
@@ -63,3 +67,16 @@ func GetSession(c *gin.Context) *SessionData {
 	c.Set("session", ret)
 	return ret
 }
+
+func GetWorkspaceSession(session *SessionData) (ret *WorkspaceSession) {
+	ret = &WorkspaceSession{}
+	if nil == session.Workspaces {
+		session.Workspaces = map[string]*WorkspaceSession{}
+	}
+	ret = session.Workspaces[WorkspaceDir]
+	if nil == ret {
+		ret = &WorkspaceSession{}
+		session.Workspaces[WorkspaceDir] = ret
+	}
+	return
+}