package wshandler import ( "context" "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 Richconn struct { *websocket.Conn closeFuncLock sync.Mutex alias string onClose map[string]func() } type UpstreamMessage struct { Alias string Accid primitive.ObjectID Target string Body []byte } type DownstreamMessage struct { Alias string Body string } type CommandType string const ( CommandType_JoinChannel = CommandType("join_channel") CommandType_LeaveChannel = CommandType("leave_channel") ) type CommandMessage struct { Cmd CommandType Args []string } func (rc *Richconn) RegistOnCloseFunc(name string, f func()) { rc.closeFuncLock.Lock() defer rc.closeFuncLock.Unlock() if rc.onClose == nil { f() return } rc.onClose[name] = f } func (rc *Richconn) HasOnCloseFunc(name string) bool { rc.closeFuncLock.Lock() defer rc.closeFuncLock.Unlock() if rc.onClose == nil { return false } _, ok := rc.onClose[name] return ok } func (rc *Richconn) UnregistOnCloseFunc(name string) (out func()) { rc.closeFuncLock.Lock() defer rc.closeFuncLock.Unlock() if rc.onClose == nil { return } out = rc.onClose[name] delete(rc.onClose, name) return } func (rc *Richconn) Closed() { rc.closeFuncLock.Lock() defer rc.closeFuncLock.Unlock() for _, f := range rc.onClose { f() } } func (rc *Richconn) WriteBytes(data []byte) error { return rc.WriteMessage(websocket.TextMessage, data) } type subhandler struct { authCache *gocommon.AuthCollection redisMsgChanName string redisCmdChanName string redisSync *redis.Client connsLock sync.Mutex connectedAlias map[string]bool connInOutChan chan *Richconn deliveryChan chan any } // WebsocketHandler : type WebsocketHandler struct { authCaches map[string]*subhandler RedisSync *redis.Client } 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) } 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, connectedAlias: make(map[string]bool), connInOutChan: make(chan *Richconn), deliveryChan: make(chan any, 1000), } authCaches[region] = sh } return &WebsocketHandler{ authCaches: authCaches, RedisSync: redisSync, } } func (ws *WebsocketHandler) Destructor() { } func (ws *WebsocketHandler) IsConnected(region string, alias string) bool { if sh := ws.authCaches[region]; sh != nil { return sh.connected(alias) } return false } func (ws *WebsocketHandler) RegisterHandlers(ctx context.Context, 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) } go sh.mainLoop(ctx) } return nil } func (sh *subhandler) setConnected(alias string, connected bool) { sh.connsLock.Lock() defer sh.connsLock.Unlock() if connected { sh.connectedAlias[alias] = true } else { delete(sh.connectedAlias, alias) } } func (sh *subhandler) connected(alias string) bool { sh.connsLock.Lock() defer sh.connsLock.Unlock() _, ok := sh.connectedAlias[alias] return ok } func (sh *subhandler) mainLoop(ctx context.Context) { defer func() { s := recover() if s != nil { logger.Error(s) } }() 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]*Richconn) 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.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] == '@' { // TODO : 특정 유저에게만 } case *CommandMessage: if usermsg.Cmd == CommandType_JoinChannel && len(usermsg.Args) == 2 { alias := usermsg.Args[0] roomName := usermsg.Args[1] conn := entireConns[alias] if conn != nil { findRoom(roomName, true).in(conn) } } else if usermsg.Cmd == CommandType_JoinChannel && len(usermsg.Args) == 2 { alias := usermsg.Args[0] roomName := usermsg.Args[1] conn := entireConns[alias] 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.alias) sh.setConnected(c.alias, false) for _, room := range rooms { room.out(c) } c.Closed() } else { sh.setConnected(c.alias, true) entireConns[c.alias] = c } } } } func upgrade_core(sh *subhandler, conn *websocket.Conn, accid primitive.ObjectID, alias string) { newconn := sh.makeRichConn(alias, conn) sh.connInOutChan <- newconn go func(c *Richconn, accid primitive.ObjectID, deliveryChan chan<- any) { for { messageType, r, err := c.NextReader() // 웹소켓에서 직접 메시지를 받지 않는다. if r != nil { io.Copy(io.Discard, r) } if err != nil { c.Close() break } if messageType == websocket.CloseMessage { break } } 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) } func (sh *subhandler) makeRichConn(alias string, conn *websocket.Conn) *Richconn { rc := Richconn{ Conn: conn, alias: alias, onClose: make(map[string]func()), } return &rc }