gocommon/wshandler/wshandler.go

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)
}