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

app/src/index.ts

@@ -27,6 +27,7 @@ import {initMessage} from "./dialog/message";
 import {getAllTabs} from "./layout/getAll";
 import {getLocalStorage} from "./protyle/util/compatibility";
 import {getSearch} from "./util/functions";
+import {checkPublishServiceClosed} from "./util/processMessage";
 import {hideAllElements} from "./protyle/ui/hideElements";
 import {loadPlugins, reloadPlugin} from "./plugin/loader";
 import "./assets/scss/base.scss";
@@ -46,6 +47,9 @@ export class App {
     public appId: string;
 
     constructor() {
+        if (checkPublishServiceClosed()) {
+            return;
+        }
         registerServiceWorker(`${Constants.SERVICE_WORKER_PATH}?v=${Constants.SIYUAN_VERSION}`);
         addBaseURL();
 

app/src/mobile/index.ts

@@ -18,6 +18,7 @@ import {hideKeyboardToolbar, showKeyboardToolbar} from "./util/keyboardToolbar";
 import {getLocalStorage, writeText} from "../protyle/util/compatibility";
 import {getCurrentEditor, openMobileFileById} from "./editor";
 import {getSearch} from "../util/functions";
+import {checkPublishServiceClosed} from "../util/processMessage";
 import {initRightMenu} from "./menu";
 import {openChangelog} from "../boot/openChangelog";
 import {registerServiceWorker} from "../util/serviceWorker";
@@ -37,6 +38,9 @@ class App {
     public appId: string;
 
     constructor() {
+        if (checkPublishServiceClosed()) {
+            return;
+        }
         registerServiceWorker(`${Constants.SERVICE_WORKER_PATH}?v=${Constants.SIYUAN_VERSION}`);
         addBaseURL();
         this.appId = Constants.SIYUAN_APPID;

app/src/util/processMessage.ts

@@ -5,6 +5,7 @@ import {hideMessage, showMessage} from "../dialog/message";
 import {setStorageVal} from "../protyle/util/compatibility";
 import {Constants} from "../constants";
 import {fetchPost} from "./fetch";
+import {isBrowser} from "./functions";
 
 export const processMessage = (response: IWebSocketData) => {
     if ("msg" === response.cmd) {
@@ -61,6 +62,10 @@ export const processMessage = (response: IWebSocketData) => {
         }
         return false;
     }
+    if ("closepublishpage" === response.cmd) {
+        handlePublishServiceClosed(response.msg);
+        return false;
+    }
 
     // 小于 0 为提示：-2 提示；-1 报错，大于 0 的错误需处理，等于 0 的为正常操作
     if (response.code < 0) {
@@ -70,3 +75,22 @@ export const processMessage = (response: IWebSocketData) => {
 
     return response;
 };
+
+export const handlePublishServiceClosed = (msg: string) => {
+    if (isBrowser()) {
+        sessionStorage.setItem("siyuanPublishServiceClosed", msg || "");
+        window.location.reload();
+    }
+};
+
+export const checkPublishServiceClosed = (): boolean => {
+    if (isBrowser()) {
+        const publishServiceClosedMsg = sessionStorage.getItem("siyuanPublishServiceClosed");
+        if (publishServiceClosedMsg) {
+            sessionStorage.removeItem("siyuanPublishServiceClosed");
+            document.body.innerHTML = `<div style="display:flex;align-items:center;justify-content:center;height:100vh">${publishServiceClosedMsg}</div>`;
+            return true;
+        }
+    }
+    return false;
+};

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)
+	}
+}