diff --git a/wshandler/wshandler.go b/wshandler/wshandler.go index 23c9550..7170bbd 100644 --- a/wshandler/wshandler.go +++ b/wshandler/wshandler.go @@ -81,9 +81,9 @@ type subhandler struct { // WebsocketHandler : type WebsocketHandler struct { - authCaches map[string]*subhandler - RedisSync *redis.Client - ClientMessageReceiver WebSocketMessageReceiver + authCaches map[string]*subhandler + RedisSync *redis.Client + receiverChain []WebSocketMessageReceiver } type wsConfig struct { @@ -130,9 +130,28 @@ func NewWebsocketHandler(authglobal gocommon.AuthCollectionGlobal) (wsh *Websock } return &WebsocketHandler{ - authCaches: authCaches, - RedisSync: redisSync, - ClientMessageReceiver: func(primitive.ObjectID, string, WebSocketMessageType, io.Reader) {}, + authCaches: authCaches, + RedisSync: redisSync, + } +} + +func (ws *WebsocketHandler) RegisterReceiver(receiver WebSocketMessageReceiver) { + ws.receiverChain = append(ws.receiverChain, receiver) +} + +func (ws *WebsocketHandler) Start(ctx context.Context) { + for _, sh := range ws.authCaches { + if len(ws.receiverChain) == 1 { + sh.callReceiver = ws.receiverChain[0] + } else { + sh.callReceiver = func(accid primitive.ObjectID, alias string, messageType WebSocketMessageType, body io.Reader) { + for _, r := range ws.receiverChain { + r(accid, alias, messageType, body) + } + } + } + + go sh.mainLoop(ctx) } } @@ -142,7 +161,7 @@ func (ws *WebsocketHandler) Cleanup() { } } -func (ws *WebsocketHandler) RegisterHandlers(ctx context.Context, serveMux *http.ServeMux, prefix string) error { +func (ws *WebsocketHandler) RegisterHandlers(serveMux *http.ServeMux, prefix string) error { for region, sh := range ws.authCaches { if region == "default" { region = "" @@ -153,8 +172,6 @@ func (ws *WebsocketHandler) RegisterHandlers(ctx context.Context, serveMux *http } else { serveMux.HandleFunc(url, sh.upgrade) } - sh.callReceiver = ws.ClientMessageReceiver - go sh.mainLoop(ctx) } return nil