From 1f668586f2dce94048e276401c4f136d2d278140 Mon Sep 17 00:00:00 2001
From: mountain <mountain@action2quare.com>
Date: Mon, 10 Jul 2023 12:50:27 +0900
Subject: [PATCH] =?UTF-8?q?wshandler=20reciever=20=ED=95=A8=EC=88=98=20?=
 =?UTF-8?q?=EC=B6=94=EA=B0=80=20=EC=88=98=EC=A0=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 wshandler/wshandler.go | 35 ++++++++++++++++++++++++++---------
 1 file changed, 26 insertions(+), 9 deletions(-)

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