package wshandler import ( "bytes" "context" "encoding/gob" "encoding/hex" "encoding/json" "fmt" "io" "net/http" "strings" "sync" "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 joinedRooms []*room } func (conn *wsconn) popRoom(r *room) int { for i, jr := range conn.joinedRooms { if jr == r { conn.joinedRooms = append(conn.joinedRooms[:i], conn.joinedRooms[i+1:]...) break } } return len(conn.joinedRooms) } func (conn *wsconn) pushRoom(r *room) { conn.joinedRooms = append(conn.joinedRooms, r) } type UpstreamMessage struct { Alias string Accid primitive.ObjectID Target string Body any Tag []string } type DownstreamMessage struct { Alias string `json:",omitempty"` Body any `json:",omitempty"` Tag []string `json:",omitempty"` } type commandType string const ( commandType_JoinRoom = commandType("join_room") commandType_LeaveRoom = commandType("leave_room") ) 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 EventReceiver interface { OnClientMessageReceived(sender *Sender, messageType WebSocketMessageType, body io.Reader) OnRoomCreated(region, name string) OnRoomDestroyed(region, name string) } type subhandler struct { redisMsgChanName string redisCmdChanName string redisSync *redis.Client connInOutChan chan *wsconn deliveryChan chan any localDeliveryChan chan any callReceiver EventReceiver connWaitGroup sync.WaitGroup region string receiverChain []EventReceiver } // WebsocketHandler : type WebsocketHandler struct { subhandlers map[string]*subhandler } type wsConfig struct { gocommon.RegionStorageConfig Maingate string `json:"maingate_service_url"` } var config wsConfig func init() { gocommon.LoadConfig(&config) gob.Register(UpstreamMessage{}) gob.Register(commandMessage{}) gob.Register(map[string]any{}) gob.Register(primitive.A{}) gob.Register(primitive.M{}) gob.Register(primitive.ObjectID{}) } func NewWebsocketHandler() (*WebsocketHandler, error) { // 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 // } subhandlers := make(map[string]*subhandler) for region, cfg := range config.RegionStorage { redisSync, err := gocommon.NewRedisClient(cfg.Redis["wshandler"]) if err != nil { return nil, err } sh := &subhandler{ 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, } subhandlers[region] = sh } return &WebsocketHandler{ subhandlers: subhandlers, }, nil } func (ws *WebsocketHandler) RegisterReceiver(region string, receiver EventReceiver) { if sh := ws.subhandlers[region]; sh != nil { sh.receiverChain = append(sh.receiverChain, receiver) } } type nilReceiver struct{} func (r *nilReceiver) OnClientMessageReceived(sender *Sender, messageType WebSocketMessageType, body io.Reader) { } func (r *nilReceiver) OnRoomCreated(region, name string) {} func (r *nilReceiver) OnRoomDestroyed(region, name string) {} type chainReceiver struct { chain []EventReceiver } func (r *chainReceiver) OnClientMessageReceived(sender *Sender, messageType WebSocketMessageType, body io.Reader) { for _, cr := range r.chain { cr.OnClientMessageReceived(sender, messageType, body) } } func (r *chainReceiver) OnRoomCreated(region, name string) { for _, cr := range r.chain { cr.OnRoomCreated(region, name) } } func (r *chainReceiver) OnRoomDestroyed(region, name string) { for _, cr := range r.chain { cr.OnRoomDestroyed(region, name) } } func (ws *WebsocketHandler) Start(ctx context.Context) { for _, sh := range ws.subhandlers { chain := sh.receiverChain if len(chain) == 0 { sh.callReceiver = &nilReceiver{} } else if len(chain) == 1 { sh.callReceiver = chain[0] } else { sh.callReceiver = &chainReceiver{chain: sh.receiverChain} } go sh.mainLoop(ctx) } } func (ws *WebsocketHandler) Cleanup() { for _, sh := range ws.subhandlers { sh.connWaitGroup.Wait() } } func (ws *WebsocketHandler) RegisterHandlers(serveMux *http.ServeMux, prefix string) error { for region, sh := range ws.subhandlers { 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 { if sh := ws.subhandlers[region]; sh != nil { state, _ := sh.redisSync.Get(context.Background(), accid.Hex()).Result() return state } return "" } func (ws *WebsocketHandler) SetState(region string, accid primitive.ObjectID, state string) { if sh := ws.subhandlers[region]; sh != nil { sh.redisSync.SetArgs(context.Background(), accid.Hex(), state, redis.SetArgs{Mode: "XX"}).Result() } } func (ws *WebsocketHandler) SendUpstreamMessage(region string, msg *UpstreamMessage) { sh := ws.subhandlers[region] if sh != nil { sh.localDeliveryChan <- msg } } func (ws *WebsocketHandler) EnterRoom(region string, room string, accid primitive.ObjectID) { sh := ws.subhandlers[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.subhandlers[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 { buffer := bytes.NewBuffer([]byte(raw.Payload)) dec := gob.NewDecoder(buffer) if raw.Channel == sh.redisMsgChanName { var msg UpstreamMessage if err := dec.Decode(&msg); err == nil { sh.deliveryChan <- &msg } else { logger.Println("decode UpstreamMessage failed :", err) } } else if raw.Channel == sh.redisCmdChanName { var cmd commandMessage if err := dec.Decode(&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) roomDestroyChan := make(chan string, 1000) findRoom := func(name string, create bool) *room { room := rooms[name] if room == nil && create { room = makeRoom(name, roomDestroyChan) rooms[name] = room room.start(ctx) go sh.callReceiver.OnRoomCreated(sh.region, name) } return room } // 유저에게서 온 메세지, 소켓 연결/해체 처리 for { buffer := bytes.NewBuffer(make([]byte, 0, 1024)) buffer.Reset() select { case destroyedRoom := <-roomDestroyChan: delete(rooms, destroyedRoom) go sh.callReceiver.OnRoomDestroyed(sh.region, destroyedRoom) 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 해야 함 ds, _ := json.Marshal(DownstreamMessage{ Alias: usermsg.Alias, Body: usermsg.Body, Tag: usermsg.Tag, }) conn.WriteMessage(websocket.TextMessage, ds) break } } var err error enc := gob.NewEncoder(buffer) if err = enc.Encode(usermsg); err == nil { _, err = sh.redisSync.Publish(context.Background(), sh.redisMsgChanName, buffer.Bytes()).Result() } if err != nil { logger.Println("gob.Encode or publish failed :", err) } 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 { conn.pushRoom(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 { if conn.popRoom(room.out(conn)) == 0 { conn.Close() } break } } } // 위에서 break 안걸리면 나한테 없으므로 publish를 해야 함. 그러면 다른 호스트가 deliveryChan으로 받는다 var err error enc := gob.NewEncoder(buffer) if err = enc.Encode(usermsg); err == nil { _, err = sh.redisSync.Publish(context.Background(), sh.redisCmdChanName, buffer.Bytes()).Result() } if err != nil { logger.Println("gob.Encode or Publish failed :", err) } } 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 { ds, _ := json.Marshal(DownstreamMessage{ Alias: usermsg.Alias, Body: usermsg.Body, Tag: usermsg.Tag, }) conn.WriteMessage(websocket.TextMessage, ds) } } 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 { conn.pushRoom(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 { if conn.popRoom(room.out(conn)) == 0 { conn.Close() } } } } default: logger.Println("usermsg is unknown type") } case c := <-sh.connInOutChan: if c.Conn == nil { delete(entireConns, c.sender.Accid.Hex()) for _, room := range c.joinedRooms { room.out(c) } c.joinedRooms = nil go sh.callReceiver.OnClientMessageReceived(c.sender, Disconnected, nil) } else { entireConns[c.sender.Accid.Hex()] = c go sh.callReceiver.OnClientMessageReceived(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.OnClientMessageReceived(c.sender, CloseMessage, r) break } if messageType == websocket.TextMessage { // 유저가 직접 보낸 메시지 sh.callReceiver.OnClientMessageReceived(c.sender, TextMessage, r) } else if messageType == websocket.BinaryMessage { sh.callReceiver.OnClientMessageReceived(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 := r.Header.Get("Authorization") req, _ := http.NewRequest("GET", fmt.Sprintf("%s/query?sk=%s", config.Maingate, sk), nil) req.Header.Add("Authorization", auth) client := http.Client{} resp, err := client.Do(req) if err != nil { w.WriteHeader(http.StatusInternalServerError) logger.Error("authorize query failed :", err) return } defer resp.Body.Close() var authinfo gocommon.Authinfo dec := json.NewDecoder(resp.Body) if err = dec.Decode(&authinfo); err != nil { w.WriteHeader(http.StatusInternalServerError) logger.Error("authorize query failed :", err) 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 = authinfo.Accid.Hex() } upgrade_core(sh, conn, authinfo.Accid, alias) }