diff --git a/app/src/index.ts b/app/src/index.ts index 51eb8358f..fbb711f9f 100644 --- a/app/src/index.ts +++ b/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(); diff --git a/app/src/mobile/index.ts b/app/src/mobile/index.ts index d0c91ea77..d3c0d40b2 100644 --- a/app/src/mobile/index.ts +++ b/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; diff --git a/app/src/util/processMessage.ts b/app/src/util/processMessage.ts index 649d3793c..00bfd4c69 100644 --- a/app/src/util/processMessage.ts +++ b/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 = `
${publishServiceClosedMsg}
`; + return true; + } + } + return false; +}; diff --git a/kernel/model/auth.go b/kernel/model/auth.go index 90f62c82e..40172c305 100644 --- a/kernel/model/auth.go +++ b/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 +} diff --git a/kernel/server/proxy/publish.go b/kernel/server/proxy/publish.go index f7f32e6ab..c1f0a0e8d 100644 --- a/kernel/server/proxy/publish.go +++ b/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() { diff --git a/kernel/server/serve.go b/kernel/server/serve.go index 7da000b20..e6536ff54 100644 --- a/kernel/server/serve.go +++ b/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) diff --git a/kernel/util/websocket.go b/kernel/util/websocket.go index 29aaa2937..a94cbe730 100644 --- a/kernel/util/websocket.go +++ b/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) + } +}