🎨 Auto-close browser page when publish service is closed https://github.com/siyuan-note/siyuan/issues/16587#issuecomment-3698421929 (#16804)

kernel/model/auth.go

@@ -159,3 +159,15 @@ func GetClaimRole(claims jwt.MapClaims) Role {
 	}
 	return RoleVisitor
 }
+
+// IsPublishServiceToken 检查 token 是否来自发布服务
+func IsPublishServiceToken(token *jwt.Token) bool {
+	if token == nil || !token.Valid {
+		return false
+	}
+	claims := GetTokenClaims(token)
+	if tokenIssuer, ok := claims["iss"].(string); ok {
+		return tokenIssuer == iss
+	}
+	return false
+}

kernel/server/proxy/publish.go

@@ -70,7 +70,6 @@ func initPublishService() {
 	if err := initPublishListener(); err == nil {
 		go startPublishReverseProxyService()
 	}
-	return
 }
 
 func initPublishListener() (err error) {
@@ -93,6 +92,9 @@ func closePublishListener() {
 		return
 	}
 
+	// 关闭所有发布服务的 WebSocket 连接
+	util.ClosePublishServiceSessions()
+
 	if err := server.Shutdown(context.Background()); err != nil {
 		logging.LogErrorf("shutdown server failed: %s", err)
 	}
@@ -101,7 +103,6 @@ func closePublishListener() {
 		logging.LogErrorf("close server failed: %s", err)
 	}
 	server, listener = nil, nil
-	return
 }
 
 func startPublishReverseProxyService() {

kernel/server/serve.go

@@ -671,6 +671,13 @@ func serveWebSocket(ginServer *gin.Engine) {
 			return
 		}
 
+		// 标记发布服务的连接
+		if token := model.ParseXAuthToken(s.Request); token != nil {
+			if model.IsPublishServiceToken(token) {
+				s.Set("isPublish", true)
+			}
+		}
+
 		util.AddPushChan(s)
 		//sessionId, _ := s.Get("id")
 		//logging.LogInfof("ws [%s] connected", sessionId)

kernel/util/websocket.go

@@ -177,7 +177,6 @@ func PushTxErr(msg string, code int, data interface{}) {
 
 func PushUpdateMsg(msgId string, msg string, timeout int) {
 	BroadcastByType("main", "msg", 0, msg, map[string]interface{}{"id": msgId, "closeTimeout": timeout})
-	return
 }
 
 func PushMsg(msg string, timeout int) (msgId string) {
@@ -454,3 +453,46 @@ func CountSessions() (ret int) {
 	})
 	return
 }
+
+// ClosePublishServiceSessions 关闭所有发布服务的 WebSocket 连接
+func ClosePublishServiceSessions() {
+	if WebSocketServer == nil {
+		return
+	}
+
+	// 收集所有发布服务的会话
+	var publishSessions []*melody.Session
+	sessions.Range(func(key, value interface{}) bool {
+		appSessions := value.(*sync.Map)
+		appSessions.Range(func(key, value interface{}) bool {
+			session := value.(*melody.Session)
+			if isPublish, ok := session.Get("isPublish"); ok && isPublish == true {
+				publishSessions = append(publishSessions, session)
+			}
+			return true
+		})
+		return true
+	})
+
+	// 发送消息通知客户端关闭页面
+	for _, session := range publishSessions {
+		event := NewResult()
+		event.Cmd = "closepublishpage"
+		event.Code = 0
+		event.Msg = "SiYuan publish service closed"
+		event.Data = map[string]interface{}{
+			"reason": "publish service closed",
+		}
+		session.Write(event.Bytes())
+	}
+
+	// 等待一小段时间让消息发送完成、客户端刷新页面之后显示消息
+	time.Sleep(500 * time.Millisecond)
+
+	// 关闭所有发布服务的 WebSocket 连接
+	for _, session := range publishSessions {
+		// 使用 "close websocket" 作为关闭消息，客户端检测到后会停止重连
+		session.CloseWithMsg([]byte("  close websocket: publish service closed"))
+		RemovePushChan(session)
+	}
+}