❤️ 完整开源界面和内核 https://github.com/siyuan-note/siyuan/issues/5013

kernel/util/websocket.go

@@ -0,0 +1,269 @@
+// SiYuan - Build Your Eternal Digital Garden
+// 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"
+
+	"github.com/88250/melody"
+)
+
+var (
+	WebSocketServer = melody.New()
+
+	// map[string]map[string]*melody.Session{}
+	sessions = sync.Map{} // {appId, {sessionId, session}}
+)
+
+// 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 ReloadUI() {
+	evt := NewCmdResult("reloadui", 0, PushModeBroadcast, 0)
+	PushEvent(evt)
+}
+
+func PushTxErr(msg string, code int, data interface{}) {
+	BroadcastByType("main", "txerr", code, msg, data)
+}
+
+func PushMsg(msg string, timeout int) {
+	evt := NewCmdResult("msg", 0, PushModeBroadcast, 0)
+	evt.Msg = msg
+	evt.Data = map[string]interface{}{"closeTimeout": timeout}
+	PushEvent(evt)
+}
+
+func PushErrMsg(msg string, timeout int) {
+	evt := NewCmdResult("msg", 0, PushModeBroadcast, 0)
+	evt.Code = -1
+	evt.Msg = msg
+	evt.Data = map[string]interface{}{"closeTimeout": timeout}
+	PushEvent(evt)
+}
+
+const (
+	PushProgressCodeProgressed = 0 // 有进度
+	PushProgressCodeEndless    = 1 // 无进度
+	PushProgressCodeEnd        = 2 // 关闭进度
+)
+
+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) {
+	evt := NewCmdResult("progress", 0, PushModeBroadcast, 0)
+	evt.Msg = msg
+	evt.Code = code
+	evt.Data = map[string]interface{}{
+		"current": current,
+		"total":   total,
+	}
+	PushEvent(evt)
+}
+
+// PushClearMsg 会清空消息提示以及进度遮罩。
+func PushClearMsg() {
+	evt := NewCmdResult("cmsg", 0, PushModeBroadcast, 0)
+	PushEvent(evt)
+}
+
+func PushDownloadProgress(id string, percent float32) {
+	evt := NewCmdResult("downloadProgress", 0, PushModeBroadcast, 0)
+	evt.Data = map[string]interface{}{
+		"id":      id,
+		"percent": percent,
+	}
+	PushEvent(evt)
+}
+
+func PushEvent(event *Result) {
+	msg := event.Bytes()
+	mode := event.PushMode
+	if "reload" == event.Cmd {
+		mode = event.ReloadPushMode
+	}
+	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 PushModeNone:
+	}
+}
+
+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 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
+}