package wshandler import ( "context" "encoding/hex" "encoding/json" "fmt" "io" "net/http" "strings" "sync" "time" "go.mongodb.org/mongo-driver/bson/primitive" "repositories.action2quare.com/ayo/gocommon" "repositories.action2quare.com/ayo/gocommon/flagx" "repositories.action2quare.com/ayo/gocommon/logger" "github.com/go-redis/redis/v8" "github.com/gorilla/websocket" ) var noSessionFlag = flagx.Bool("nosession", false, "nosession=[true|false]") type wsconn struct { *websocket.Conn sender *Sender } type UpstreamMessage struct { Alias string Accid primitive.ObjectID Target string Body []byte Tag []string } type DownstreamMessage struct { Alias string `json:",omitempty"` Body string `json:",omitempty"` Tag []string `json:",omitempty"` } type commandType string const ( commandType_JoinRoom = commandType("join_room") commandType_LeaveRoom = commandType("leave_room") commandType_WriteControl = commandType("write_control") ) type commandMessage struct { Cmd commandType Args []any } type WebSocketMessageType int const ( TextMessage = WebSocketMessageType(websocket.TextMessage) BinaryMessage = WebSocketMessageType(websocket.BinaryMessage) CloseMessage = WebSocketMessageType(websocket.CloseMessage) PingMessage = WebSocketMessageType(websocket.PingMessage) PongMessage = WebSocketMessageType(websocket.PongMessage) Connected = WebSocketMessageType(100) Disconnected = WebSocketMessageType(101) ) type Sender struct { Region string Accid primitive.ObjectID Alias string } type WebSocketMessageReceiver func(sender *Sender, messageType WebSocketMessageType, body io.Reader) type subhandler struct { authCache *gocommon.AuthCollection redisMsgChanName string redisCmdChanName string redisSync *redis.Client connInOutChan chan *wsconn deliveryChan chan any localDeliveryChan chan any callReceiver WebSocketMessageReceiver connWaitGroup sync.WaitGroup region string } // WebsocketHandler : type WebsocketHandler struct { authCaches map[string]*subhandler RedisSync *redis.Client receiverChain map[string][]WebSocketMessageReceiver } type wsConfig struct { SyncPipeline string `json:"ws_sync_pipeline"` } func NewWebsocketHandler(authglobal gocommon.AuthCollectionGlobal) (wsh *WebsocketHandler) { var config wsConfig gocommon.LoadConfig(&config) redisSync, err := gocommon.NewRedisClient(config.SyncPipeline, 0) if err != nil { panic(err) } // decoder := func(r io.Reader) *T { // if r == nil { // // 접속이 끊겼을 때. // return nil // } // var m T // dec := json.NewDecoder(r) // if err := dec.Decode(&m); err != nil { // logger.Println(err) // } // // decoding 실패하더라도 빈 *T를 내보냄 // return &m // } authCaches := make(map[string]*subhandler) for _, region := range authglobal.Regions() { sh := &subhandler{ authCache: authglobal.Get(region), redisMsgChanName: fmt.Sprintf("_wsh_msg_%s", region), redisCmdChanName: fmt.Sprintf("_wsh_cmd_%s", region), redisSync: redisSync, connInOutChan: make(chan *wsconn), deliveryChan: make(chan any, 1000), localDeliveryChan: make(chan any, 100), region: region, } authCaches[region] = sh } return &WebsocketHandler{ authCaches: authCaches, RedisSync: redisSync, receiverChain: make(map[string][]WebSocketMessageReceiver), } } func (ws *WebsocketHandler) RegisterReceiver(region string, receiver WebSocketMessageReceiver) { ws.receiverChain[region] = append(ws.receiverChain[region], receiver) } func (ws *WebsocketHandler) Start(ctx context.Context) { for region, sh := range ws.authCaches { chain := ws.receiverChain[region] if len(chain) == 0 { sh.callReceiver = func(sender *Sender, messageType WebSocketMessageType, body io.Reader) {} } else if len(chain) == 1 { sh.callReceiver = chain[0] } else { sh.callReceiver = func(sender *Sender, messageType WebSocketMessageType, body io.Reader) { for _, r := range chain { r(sender, messageType, body) } } } go sh.mainLoop(ctx) } } func (ws *WebsocketHandler) Cleanup() { for _, sh := range ws.authCaches { sh.connWaitGroup.Wait() } } func (ws *WebsocketHandler) RegisterHandlers(serveMux *http.ServeMux, prefix string) error { for region, sh := range ws.authCaches { if region == "default" { region = "" } url := gocommon.MakeHttpHandlerPattern(prefix, region, "ws") if *noSessionFlag { serveMux.HandleFunc(url, sh.upgrade_nosession) } else { serveMux.HandleFunc(url, sh.upgrade) } } return nil } func (ws *WebsocketHandler) GetState(region string, accid primitive.ObjectID) string { state, err := ws.RedisSync.Get(context.Background(), accid.Hex()).Result() if err == redis.Nil { return "" } return state } func (ws *WebsocketHandler) SetState(region string, accid primitive.ObjectID, state string) { ws.RedisSync.SetArgs(context.Background(), accid.Hex(), state, redis.SetArgs{Mode: "XX"}).Result() } func (ws *WebsocketHandler) SendUpstreamMessage(region string, msg *UpstreamMessage) { sh := ws.authCaches[region] if sh != nil { sh.localDeliveryChan <- msg } } func (ws *WebsocketHandler) SendCloseMessage(region string, target string, text string) { sh := ws.authCaches[region] if sh != nil { sh.localDeliveryChan <- &commandMessage{ Cmd: commandType_WriteControl, Args: []any{ target, int(websocket.CloseMessage), websocket.FormatCloseMessage(websocket.CloseNormalClosure, text), }, } } } func (ws *WebsocketHandler) EnterRoom(region string, room string, accid primitive.ObjectID) { sh := ws.authCaches[region] if sh != nil { sh.localDeliveryChan <- &commandMessage{ Cmd: commandType_JoinRoom, Args: []any{room, accid}, } } } func (ws *WebsocketHandler) LeaveRoom(region string, room string, accid primitive.ObjectID) { sh := ws.authCaches[region] if sh != nil { sh.localDeliveryChan <- &commandMessage{ Cmd: commandType_LeaveRoom, Args: []any{room, accid}, } } } func (sh *subhandler) mainLoop(ctx context.Context) { defer func() { s := recover() if s != nil { logger.Error(s) } }() // redis channel에서 유저가 보낸 메시지를 읽는 go rountine go func() { var pubsub *redis.PubSub for { if pubsub == nil { pubsub = sh.redisSync.Subscribe(ctx, sh.redisMsgChanName, sh.redisCmdChanName) } raw, err := pubsub.ReceiveMessage(ctx) if err == nil { if raw.Channel == sh.redisMsgChanName { var msg UpstreamMessage if err := json.Unmarshal([]byte(raw.Payload), &msg); err == nil { sh.deliveryChan <- &msg } else { logger.Println("decode UpstreamMessage failed :", err) } } else if raw.Channel == sh.redisCmdChanName { var cmd commandMessage if err := json.Unmarshal([]byte(raw.Payload), &cmd); err == nil { sh.deliveryChan <- &cmd } else { logger.Println("decode UpstreamMessage failed :", err) } } } else { logger.Println("pubsub.ReceiveMessage failed :", err) pubsub.Close() pubsub = nil if ctx.Err() != nil { break } } } }() entireConns := make(map[string]*wsconn) rooms := make(map[string]*room) findRoom := func(name string, create bool) *room { room := rooms[name] if room == nil && create { room = makeRoom(name) rooms[name] = room room.start(ctx) } return room } // 유저에게서 온 메세지, 소켓 연결/해체 처리 for { select { case usermsg := <-sh.localDeliveryChan: // 로컬에 connection이 있는지 먼저 확인해 보기 위한 채널 // 없으면 publish한다. switch usermsg := usermsg.(type) { case *UpstreamMessage: target := usermsg.Target if target[0] == '@' { accid := target[1:] conn := entireConns[accid] if conn != nil { // 이 경우 아니면 publish 해야 함 conn.WriteMessage(websocket.TextMessage, usermsg.Body) break } } if bt, err := json.Marshal(usermsg); err == nil { sh.redisSync.Publish(context.Background(), sh.redisMsgChanName, bt).Result() } case *commandMessage: if usermsg.Cmd == commandType_JoinRoom && len(usermsg.Args) == 2 { roomName := usermsg.Args[0].(string) accid := usermsg.Args[1].(primitive.ObjectID) conn := entireConns[accid.Hex()] if conn != nil { findRoom(roomName, true).in(conn) break } } else if usermsg.Cmd == commandType_LeaveRoom && len(usermsg.Args) == 2 { roomName := usermsg.Args[0].(string) accid := usermsg.Args[1].(primitive.ObjectID) conn := entireConns[accid.Hex()] if conn != nil { if room := findRoom(roomName, false); room != nil { room.out(conn) break } } } else if usermsg.Cmd == commandType_WriteControl && len(usermsg.Args) == 2 { accid := usermsg.Args[0].(string) conn := entireConns[accid] if conn != nil { conn.WriteControl(usermsg.Args[1].(int), usermsg.Args[2].([]byte), time.Time{}) break } } // 위에서 break 안걸리면 나한테 없으므로 publish를 해야 함. 그러면 다른 호스트가 deliveryChan으로 받는다 if bt, err := json.Marshal(usermsg); err == nil { sh.redisSync.Publish(context.Background(), sh.redisCmdChanName, bt).Result() } } case usermsg := <-sh.deliveryChan: switch usermsg := usermsg.(type) { case *UpstreamMessage: target := usermsg.Target if target[0] == '#' { // 룸에 브로드캐스팅 roomName := target[1:] if room := findRoom(roomName, false); room != nil { room.broadcast(usermsg) } } else if target[0] == '@' { accid := target[1:] conn := entireConns[accid] if conn != nil { conn.WriteMessage(websocket.TextMessage, usermsg.Body) } } case *commandMessage: if usermsg.Cmd == commandType_JoinRoom && len(usermsg.Args) == 2 { roomName := usermsg.Args[0].(string) accid := usermsg.Args[1].(primitive.ObjectID) conn := entireConns[accid.Hex()] if conn != nil { findRoom(roomName, true).in(conn) } } else if usermsg.Cmd == commandType_LeaveRoom && len(usermsg.Args) == 2 { roomName := usermsg.Args[0].(string) accid := usermsg.Args[1].(primitive.ObjectID) conn := entireConns[accid.Hex()] if conn != nil { if room := findRoom(roomName, false); room != nil { room.out(conn) } } } default: logger.Println("usermsg is unknown type") } case c := <-sh.connInOutChan: if c.Conn == nil { delete(entireConns, c.sender.Accid.Hex()) for _, room := range rooms { room.out(c) } sh.callReceiver(c.sender, Disconnected, nil) } else { entireConns[c.sender.Accid.Hex()] = c sh.callReceiver(c.sender, Connected, nil) } } } } func upgrade_core(sh *subhandler, conn *websocket.Conn, accid primitive.ObjectID, alias string) { newconn := &wsconn{ Conn: conn, sender: &Sender{ Region: sh.region, Alias: alias, Accid: accid, }, } sh.connInOutChan <- newconn sh.connWaitGroup.Add(1) go func(c *wsconn, accid primitive.ObjectID, deliveryChan chan<- any) { sh.redisSync.Set(context.Background(), accid.Hex(), "online", 0) for { messageType, r, err := c.NextReader() if err != nil { c.Close() break } if messageType == websocket.CloseMessage { sh.callReceiver(c.sender, CloseMessage, r) break } if messageType == websocket.TextMessage { // 유저가 직접 보낸 메시지 sh.callReceiver(c.sender, TextMessage, r) } else if messageType == websocket.BinaryMessage { sh.callReceiver(c.sender, BinaryMessage, r) } } sh.redisSync.Del(context.Background(), accid.Hex()) sh.connWaitGroup.Done() c.Conn = nil sh.connInOutChan <- c }(newconn, accid, sh.deliveryChan) } func (sh *subhandler) upgrade_nosession(w http.ResponseWriter, r *http.Request) { // 클라이언트 접속 defer func() { s := recover() if s != nil { logger.Error(s) } io.Copy(io.Discard, r.Body) r.Body.Close() }() auth := strings.Split(r.Header.Get("Authorization"), " ") if len(auth) != 2 { w.WriteHeader(http.StatusBadRequest) return } if auth[0] != "Editor" { w.WriteHeader(http.StatusBadRequest) return } temp, err := hex.DecodeString(auth[1]) if err != nil { w.WriteHeader(http.StatusBadRequest) return } if len(temp) != len(primitive.NilObjectID) { w.WriteHeader(http.StatusBadRequest) return } raw := (*[12]byte)(temp) accid := primitive.ObjectID(*raw) var upgrader = websocket.Upgrader{} // use default options conn, err := upgrader.Upgrade(w, r, nil) if err != nil { w.WriteHeader(http.StatusInternalServerError) return } var alias string if v := r.Header.Get("AS-X-ALIAS"); len(v) > 0 { alias = v } else { alias = accid.Hex() } upgrade_core(sh, conn, accid, alias) } func (sh *subhandler) upgrade(w http.ResponseWriter, r *http.Request) { // 클라이언트 접속 defer func() { s := recover() if s != nil { logger.Error(s) } io.Copy(io.Discard, r.Body) r.Body.Close() }() sk := r.Header.Get("AS-X-SESSION") auth := strings.Split(r.Header.Get("Authorization"), " ") if len(auth) != 2 { //TODO : 클라이언트는 BadRequest를 받으면 로그인 화면으로 돌아가야 한다. w.WriteHeader(http.StatusBadRequest) return } authtoken := auth[1] accid, success := sh.authCache.IsValid(sk, authtoken) if !success { w.WriteHeader(http.StatusUnauthorized) return } var upgrader = websocket.Upgrader{} // use default options conn, err := upgrader.Upgrade(w, r, nil) if err != nil { w.WriteHeader(http.StatusInternalServerError) return } var alias string if v := r.Header.Get("AS-X-ALIAS"); len(v) > 0 { alias = v } else { alias = accid.Hex() } upgrade_core(sh, conn, accid, alias) }