gocommon/wshandler/wshandler.go

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
}