Something went wrong. Try again.
A privacy-first, self-hosted, fully open source personal knowledge management software, written in typescript and golang. (PERSONAL FORK)
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499// SiYuan - Refactor your thinking// Copyright (c) 2020-present, b3log.org//// This program is free software: you can redistribute it and/or modify// it under the terms of the GNU Affero General Public License as published by// the Free Software Foundation, either version 3 of the License, or// (at your option) any later version.//// This program is distributed in the hope that it will be useful,// but WITHOUT ANY WARRANTY; without even the implied warranty of// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the// GNU Affero General Public License for more details.//// You should have received a copy of the GNU Affero General Public License// along with this program. If not, see <https://www.gnu.org/licenses/>.
package util
import ( "sync" "time"
"github.com/88250/gulu" "github.com/olahol/melody" "github.com/siyuan-note/eventbus")
var ( WebSocketServer *melody.Melody
// map[string]map[string]*melody.Session{} sessions = sync.Map{} // {appId, {sessionId, session}})
func BroadcastByTypeAndExcludeApp(excludeApp, typ, cmd string, code int, msg string, data interface{}) { sessions.Range(func(key, value interface{}) bool { appSessions := value.(*sync.Map) if key == excludeApp { return true }
appSessions.Range(func(key, value interface{}) bool { session := value.(*melody.Session) if t, ok := session.Get("type"); ok && typ == t { event := NewResult() event.Cmd = cmd event.Code = code event.Msg = msg event.Data = data session.Write(event.Bytes()) } return true }) return true })}
func BroadcastByTypeAndApp(typ, app, cmd string, code int, msg string, data interface{}) { appSessions, ok := sessions.Load(app) if !ok { return }
appSessions.(*sync.Map).Range(func(key, value interface{}) bool { session := value.(*melody.Session) if t, ok := session.Get("type"); ok && typ == t { event := NewResult() event.Cmd = cmd event.Code = code event.Msg = msg event.Data = data session.Write(event.Bytes()) } return true })}
// BroadcastByType 广播所有实例上 typ 类型的会话。func BroadcastByType(typ, cmd string, code int, msg string, data interface{}) { typeSessions := SessionsByType(typ) for _, sess := range typeSessions { event := NewResult() event.Cmd = cmd event.Code = code event.Msg = msg event.Data = data sess.Write(event.Bytes()) }}
func SessionsByType(typ string) (ret []*melody.Session) { ret = []*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 t, ok := session.Get("type"); ok && typ == t { ret = append(ret, session) } return true }) return true }) return}
func AddPushChan(session *melody.Session) { appID := session.Request.URL.Query().Get("app") session.Set("app", appID) id := session.Request.URL.Query().Get("id") session.Set("id", id) typ := session.Request.URL.Query().Get("type") session.Set("type", typ)
if appSessions, ok := sessions.Load(appID); !ok { appSess := &sync.Map{} appSess.Store(id, session) sessions.Store(appID, appSess) } else { (appSessions.(*sync.Map)).Store(id, session) }}
func RemovePushChan(session *melody.Session) { app, _ := session.Get("app") id, _ := session.Get("id")
if nil == app || nil == id { return }
appSess, _ := sessions.Load(app) if nil != appSess { appSessions := appSess.(*sync.Map) appSessions.Delete(id) if 1 > lenOfSyncMap(appSessions) { sessions.Delete(app) } }}
func lenOfSyncMap(m *sync.Map) (ret int) { m.Range(func(key, value interface{}) bool { ret++ return true }) return}
func ClosePushChan(id string) { sessions.Range(func(key, value interface{}) bool { appSessions := value.(*sync.Map) appSessions.Range(func(key, value interface{}) bool { session := value.(*melody.Session) if sid, _ := session.Get("id"); sid == id { session.CloseWithMsg([]byte(" close websocket")) RemovePushChan(session) } return true }) return true })}
func ReloadUIResetScroll() { BroadcastByType("main", "reloadui", 0, "", map[string]interface{}{"resetScroll": true})}
func ReloadUI() { BroadcastByType("main", "reloadui", 0, "", nil)}
func PushTxErr(msg string, code int, data interface{}) { BroadcastByType("main", "txerr", code, msg, data)}
func PushUpdateMsg(msgId string, msg string, timeout int) { BroadcastByType("main", "msg", 0, msg, map[string]interface{}{"id": msgId, "closeTimeout": timeout})}
func PushMsg(msg string, timeout int) (msgId string) { msgId = gulu.Rand.String(7) BroadcastByType("main", "msg", 0, msg, map[string]interface{}{"id": msgId, "closeTimeout": timeout}) return}
func PushMsgWithApp(app, msg string, timeout int) (msgId string) { msgId = gulu.Rand.String(7) if "" == app { BroadcastByType("main", "msg", 0, msg, map[string]interface{}{"id": msgId, "closeTimeout": timeout}) return } BroadcastByTypeAndApp("main", app, "msg", 0, msg, map[string]interface{}{"id": msgId, "closeTimeout": timeout}) return}
func PushErrMsg(msg string, timeout int) (msgId string) { msgId = gulu.Rand.String(7) BroadcastByType("main", "msg", -1, msg, map[string]interface{}{"id": msgId, "closeTimeout": timeout}) return}
func PushStatusBar(msg string) { msg += " (" + time.Now().Format("2006-01-02 15:04:05") + ")" BroadcastByType("main", "statusbar", 0, msg, nil)}
func PushBackgroundTask(data map[string]interface{}) { BroadcastByType("main", "backgroundtask", 0, "", data)}
func PushReloadFiletree() { BroadcastByType("filetree", "reloadFiletree", 0, "", nil)}
func PushReloadTag() { BroadcastByType("main", "reloadTag", 0, "", nil)}
type BlockStatResult struct { RuneCount int `json:"runeCount"` WordCount int `json:"wordCount"` LinkCount int `json:"linkCount"` ImageCount int `json:"imageCount"` RefCount int `json:"refCount"` BlockCount int `json:"blockCount"`}
func ContextPushMsg(context map[string]interface{}, msg string) { switch context[eventbus.CtxPushMsg].(int) { case eventbus.CtxPushMsgToNone: break case eventbus.CtxPushMsgToProgress: PushEndlessProgress(msg) case eventbus.CtxPushMsgToStatusBar: PushStatusBar(msg) case eventbus.CtxPushMsgToStatusBarAndProgress: PushStatusBar(msg) PushEndlessProgress(msg) }}
const ( PushProgressCodeProgressed = 0 // 有进度 PushProgressCodeEndless = 1 // 无进度 PushProgressCodeEnd = 2 // 关闭进度)
func PushClearAllMsg() { ClearPushProgress(100) PushClearMsg("")}
func ClearPushProgress(total int) { PushProgress(PushProgressCodeEnd, total, total, "")}
func PushEndlessProgress(msg string) { PushProgress(PushProgressCodeEndless, 1, 1, msg)}
func PushProgress(code, current, total int, msg string) { BroadcastByType("main", "progress", code, msg, map[string]interface{}{ "current": current, "total": total, })}
// PushClearMsg 会清空指定消息。func PushClearMsg(msgId string) { BroadcastByType("main", "cmsg", 0, "", map[string]interface{}{"id": msgId})}
// PushClearProgress 取消进度遮罩。func PushClearProgress() { BroadcastByType("main", "cprogress", 0, "", nil)}
func PushUpdateIDs(ids map[string]string) { BroadcastByType("main", "updateids", 0, "", ids)}
func PushReloadDoc(rootID string) { BroadcastByType("main", "reloaddoc", 0, "", rootID)}
func PushSaveDoc(rootID, typ string, sources interface{}) { evt := NewCmdResult("savedoc", 0, PushModeBroadcast) evt.Data = map[string]interface{}{ "rootID": rootID, "type": typ, "sources": sources, } PushEvent(evt)}
func PushReloadDocInfo(docInfo map[string]any) { BroadcastByType("filetree", "reloadDocInfo", 0, "", docInfo)}
func PushReloadProtyle(rootID string) { BroadcastByType("protyle", "reload", 0, "", rootID)}
func PushSetRefDynamicText(rootID, blockID, defBlockID, refText string) { BroadcastByType("main", "setRefDynamicText", 0, "", map[string]interface{}{"rootID": rootID, "blockID": blockID, "defBlockID": defBlockID, "refText": refText})}
func PushSetDefRefCount(rootID, blockID string, defIDs []string, refCount, rootRefCount int) { BroadcastByType("main", "setDefRefCount", 0, "", map[string]interface{}{"rootID": rootID, "blockID": blockID, "refCount": refCount, "rootRefCount": rootRefCount, "defIDs": defIDs})}
func PushLocalShorthandCount(count int) { BroadcastByType("main", "setLocalShorthandCount", 0, "", map[string]interface{}{"count": count})}
func PushProtyleLoading(rootID, msg string) { BroadcastByType("protyle", "addLoading", 0, msg, rootID)}
func PushReloadEmojiConf() { BroadcastByType("main", "reloadEmojiConf", 0, "", nil)}
func PushDownloadProgress(id string, percent float32) { evt := NewCmdResult("downloadProgress", 0, PushModeBroadcast) evt.Data = map[string]interface{}{ "id": id, "percent": percent, } PushEvent(evt)}
func PushEvent(event *Result) { msg := event.Bytes() mode := event.PushMode switch mode { case PushModeBroadcast: Broadcast(msg) case PushModeSingleSelf: single(msg, event.AppId, event.SessionId) case PushModeBroadcastExcludeSelf: broadcastOthers(msg, event.SessionId) case PushModeBroadcastExcludeSelfApp: broadcastOtherApps(msg, event.AppId) case PushModeBroadcastApp: broadcastApp(msg, event.AppId) case PushModeBroadcastMainExcludeSelfApp: broadcastOtherAppMains(msg, event.AppId) }}
func single(msg []byte, appId, sid string) { sessions.Range(func(key, value interface{}) bool { appSessions := value.(*sync.Map) if key != appId { return true }
appSessions.Range(func(key, value interface{}) bool { session := value.(*melody.Session) if id, _ := session.Get("id"); id == sid { session.Write(msg) } return true }) return true })}
func Broadcast(msg []byte) { sessions.Range(func(key, value interface{}) bool { appSessions := value.(*sync.Map) appSessions.Range(func(key, value interface{}) bool { session := value.(*melody.Session) session.Write(msg) return true }) return true })}
func broadcastOtherApps(msg []byte, excludeApp string) { sessions.Range(func(key, value interface{}) bool { appSessions := value.(*sync.Map) appSessions.Range(func(key, value interface{}) bool { session := value.(*melody.Session) if app, _ := session.Get("app"); app == excludeApp { return true } session.Write(msg) return true }) return true })}
func broadcastOtherAppMains(msg []byte, excludeApp string) { sessions.Range(func(key, value interface{}) bool { appSessions := value.(*sync.Map) appSessions.Range(func(key, value interface{}) bool { session := value.(*melody.Session) if app, _ := session.Get("app"); app == excludeApp { return true }
if t, ok := session.Get("type"); ok && "main" != t { return true }
session.Write(msg) return true }) return true })}
func broadcastApp(msg []byte, app string) { sessions.Range(func(key, value interface{}) bool { appSessions := value.(*sync.Map) appSessions.Range(func(key, value interface{}) bool { session := value.(*melody.Session) if sessionApp, _ := session.Get("app"); sessionApp != app { return true } session.Write(msg) return true }) return true })}
func broadcastOthers(msg []byte, excludeSID string) { sessions.Range(func(key, value interface{}) bool { appSessions := value.(*sync.Map) appSessions.Range(func(key, value interface{}) bool { session := value.(*melody.Session) if id, _ := session.Get("id"); id == excludeSID { return true } session.Write(msg) return true }) return true })}
func CountSessions() (ret int) { sessions.Range(func(key, value interface{}) bool { ret++ return true }) 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) }}