gocommon/wshandler/wshandler.go

package wshandler

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/gob"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"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"
	"repositories.action2quare.com/ayo/gocommon/session"

	"github.com/go-redis/redis/v8"
	"github.com/gorilla/websocket"
)

var noAuthFlag = flagx.Bool("noauth", false, "")

type Conn struct {
	innerConn *websocket.Conn
	spinLock  int32
}

func makeConn(conn *websocket.Conn) *Conn {
	return &Conn{
		innerConn: conn,
		spinLock:  0,
	}
}

type wsconn struct {
	*Conn
	sender       *Sender
	closeMessage string
}

type noCopy struct{}
type websocketWriter struct {
	_           noCopy
	innerConn   *websocket.Conn
	spinLock    *int32
	fingerprint int32
}

func (c websocketWriter) writeImpl(vf func() error) error {
	defer atomic.StoreInt32(c.spinLock, 0)

	for i := int64(0); ; i++ {
		if atomic.CompareAndSwapInt32(c.spinLock, 0, c.fingerprint) {
			return vf()
		}

		time.Sleep(time.Microsecond)
		if i >= int64(time.Second/time.Microsecond) && i&int64(time.Second/time.Microsecond) == 0 {
			// 1초동안 락 실패
			logger.Println("websocket write lock failed : ", i/int64(time.Second/time.Microsecond))
		}
	}
}

func (c websocketWriter) WriteJSON(v interface{}) error {
	return c.writeImpl(func() error { return c.innerConn.WriteJSON(v) })
}

func (c websocketWriter) WriteMessage(messageType int, data []byte) error {
	return c.writeImpl(func() error { return c.innerConn.WriteMessage(messageType, data) })
}

func (c websocketWriter) WritePreparedMessage(pm *websocket.PreparedMessage) error {
	return c.writeImpl(func() error { return c.innerConn.WritePreparedMessage(pm) })
}

func (c websocketWriter) WriteControl(messageType int, data []byte, deadline time.Time) error {
	return c.writeImpl(func() error { return c.innerConn.WriteControl(messageType, data, deadline) })
}

func (c *Conn) SetReadDeadline(t time.Time) error {
	return c.innerConn.SetReadDeadline(t)
}

func (c *Conn) Close() error {
	return c.innerConn.Close()
}

func (c *Conn) ReadMessage() (messageType int, p []byte, err error) {
	return c.innerConn.ReadMessage()
}

func (c *Conn) RemoteAddr() net.Addr {
	return c.innerConn.RemoteAddr()
}

func (c *Conn) NextReader() (messageType int, r io.Reader, err error) {
	return c.innerConn.NextReader()
}

var websocketWriterSeq = int32(1)

func (c *Conn) MakeWriter() websocketWriter {
	return websocketWriter{
		innerConn:   c.innerConn,
		spinLock:    &c.spinLock,
		fingerprint: atomic.AddInt32(&websocketWriterSeq, 1),
	}
}

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_EnterRoom     = commandType("enter_room")
	commandType_LeaveRoom     = commandType("leave_room")
	ForceShutdownCloseMessage = "force_shutdown"
)

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

type Sender struct {
	Accid primitive.ObjectID
	Alias string
}

type EventReceiver interface {
	OnClientMessageReceived(sender *Sender, messageType WebSocketMessageType, body io.Reader)
	OnRoomCreated(name string)
	OnRoomDestroyed(name string)
}

type send_msg_queue_elem struct {
	to   *Conn
	pmsg *websocket.PreparedMessage
	//msg  []byte
}

type websocketHandlerBase struct {
	redisMsgChanName  string
	redisCmdChanName  string
	redisSync         *redis.Client
	connInOutChan     chan *wsconn
	deliveryChan      chan any
	localDeliveryChan chan any
	forceCloseChan    chan primitive.ObjectID
	sendMsgChan       chan send_msg_queue_elem

	connWaitGroup   sync.WaitGroup
	sessionConsumer session.Consumer
}

type WebsocketHandler struct {
	WebsocketApiBroker
	websocketHandlerBase
}

type wsConfig struct {
	gocommon.StorageAddr `json:"storage"`
}

func init() {
	gob.Register(UpstreamMessage{})
	gob.Register(commandMessage{})
	gob.Register(map[string]any{})
	gob.Register(primitive.A{})
	gob.Register(primitive.M{})
	gob.Register(primitive.D{})
	gob.Register(primitive.ObjectID{})
	gob.Register([]any{})
}

func NewWebsocketHandler(consumer session.Consumer, redisUrl string) (*WebsocketHandler, error) {
	var config wsConfig
	if err := gocommon.LoadConfig(&config); err != nil {
		return nil, err
	}

	redisSync, err := gocommon.NewRedisClient(redisUrl)
	if err != nil {
		return nil, logger.ErrorWithCallStack(err)
	}

	sendchan := make(chan send_msg_queue_elem, 1000)
	go func() {
		sender := func(elem *send_msg_queue_elem) {
			defer func() {
				r := recover()
				if r != nil {
					logger.Println("send_msg_queue_elem sender recover :", r)
				}
			}()
			if elem == nil {
				return
			}
			if elem.to == nil {
				return
			}

			elem.to.MakeWriter().WritePreparedMessage(elem.pmsg)
		}

		for elem := range sendchan {
			sender(&elem)
		}
	}()

	ws := &WebsocketHandler{
		websocketHandlerBase: websocketHandlerBase{
			redisMsgChanName:  fmt.Sprintf("_wsh_msg_%d", redisSync.Options().DB),
			redisCmdChanName:  fmt.Sprintf("_wsh_cmd_%d", redisSync.Options().DB),
			redisSync:         redisSync,
			connInOutChan:     make(chan *wsconn),
			deliveryChan:      make(chan any, 1000),
			localDeliveryChan: make(chan any, 100),
			forceCloseChan:    make(chan primitive.ObjectID),
			sendMsgChan:       sendchan,
			sessionConsumer:   consumer,
		},
	}
	consumer.RegisterOnSessionInvalidated(ws.onSessionInvalidated)
	return ws, nil
}

func (ws *WebsocketHandler) Start(ctx context.Context) {
	ws.connWaitGroup.Add(1)
	go ws.mainLoop(ctx)
}

func (ws *WebsocketHandler) Cleanup() {
	ws.connWaitGroup.Wait()
}

func (ws *WebsocketHandler) RegisterHandlers(serveMux *http.ServeMux, prefix string) error {
	url := gocommon.MakeHttpHandlerPattern(prefix, "ws")
	if *noAuthFlag {
		serveMux.HandleFunc(url, ws.upgrade_nosession)
	} else {
		serveMux.HandleFunc(url, ws.upgrade)
	}

	return nil
}

func (ws *WebsocketHandler) SendUpstreamMessage(msg *UpstreamMessage) {
	ws.localDeliveryChan <- msg
}

func (ws *WebsocketHandler) WriteDirectMessage(c *Conn, messageType int, data []byte) {
	pmsg, _ := websocket.NewPreparedMessage(messageType, data)
	ws.sendMsgChan <- send_msg_queue_elem{
		to:   c,
		pmsg: pmsg,
	}
}

func (ws *WebsocketHandler) EnterRoom(room string, accid primitive.ObjectID) {
	ws.localDeliveryChan <- &commandMessage{
		Cmd:  commandType_EnterRoom,
		Args: []any{room, accid},
	}
}

func (ws *WebsocketHandler) LeaveRoom(room string, accid primitive.ObjectID) {
	ws.localDeliveryChan <- &commandMessage{
		Cmd:  commandType_LeaveRoom,
		Args: []any{room, accid},
	}
}

func (ws *WebsocketHandler) onSessionInvalidated(accid primitive.ObjectID) {
	ws.forceCloseChan <- accid
}

func (ws *WebsocketHandler) mainLoop(ctx context.Context) {
	defer func() {
		ws.connWaitGroup.Done()
		s := recover()
		if s != nil {
			logger.Error(s)
		}
	}()

	// redis channel에서 유저가 보낸 메시지를 읽는 go rountine
	go func() {
		var pubsub *redis.PubSub
		waittimer := time.Second

		for {
			if pubsub == nil {
				pubsub = ws.redisSync.Subscribe(ctx, ws.redisMsgChanName, ws.redisCmdChanName)
			}

			raw, err := pubsub.ReceiveMessage(ctx)
			if err == nil {
				waittimer = time.Second

				buffer := bytes.NewBuffer([]byte(raw.Payload))
				dec := gob.NewDecoder(buffer)

				if raw.Channel == ws.redisMsgChanName {
					var msg UpstreamMessage
					if err := dec.Decode(&msg); err == nil {
						ws.deliveryChan <- &msg
					} else {
						logger.Println("decode UpstreamMessage failed :", err)
					}
				} else if raw.Channel == ws.redisCmdChanName {
					var cmd commandMessage
					if err := dec.Decode(&cmd); err == nil {
						ws.deliveryChan <- &cmd
					} else {
						logger.Println("decode UpstreamMessage failed :", err)
					}
				}
			} else {
				logger.Println("pubsub.ReceiveMessage failed :", err)
				pubsub.Close()
				pubsub = nil

				time.Sleep(waittimer)
				waittimer = waittimer * 2
				if waittimer > time.Minute {
					waittimer = time.Minute
				}

				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, ws.sendMsgChan)
			rooms[name] = room
			room.start(ctx)
			//go ws.callReceiver.OnRoomCreated(name)
		}
		return room
	}

	shutdownConn := func(c *wsconn) {
		defer func() {
			r := recover()
			if r != nil {
				logger.Println(r)
			}
		}()
		c.closeMessage = ForceShutdownCloseMessage
		ws.ClientDisconnected(c)
		c.Close()
	}

	defer func() {
		for _, conn := range entireConns {
			shutdownConn(conn)
		}
	}()

	errProcessFailed_NotInRoom := errors.New("not in room")
	processPrivateUpstreamMessage := func(usermsg *UpstreamMessage) bool {
		target := usermsg.Target
		if target[0] == '#' {
			return false
		}
		roomindex := strings.IndexByte(target, '@')
		var accid string
		var roomid string
		if roomindex > 0 {
			accid = target[:roomindex]
			roomid = target[roomindex+1:]
		} else {
			accid = target
		}

		conn := entireConns[accid]
		if conn == nil {
			return false
		}

		if len(roomid) > 0 {
			if room := findRoom(roomid, false); room == nil {
				return false
			}
		}

		// conn에게 메시지 보내면 처리 완료
		ds, _ := json.Marshal(DownstreamMessage{
			Alias: usermsg.Alias,
			Body:  usermsg.Body,
			Tag:   usermsg.Tag,
		})
		pmsg, err := websocket.NewPreparedMessage(websocket.TextMessage, ds)
		if err != nil {
			logger.Println("websocket.NewPreparedMessage failed :", err)
			return false
		}

		ws.sendMsgChan <- send_msg_queue_elem{
			to:   conn.Conn,
			pmsg: pmsg,
		}

		return true
	}
	errCommandArgumentNotCorrect := errors.New("command arguments are not correct")
	processCommandMessage := func(usermsg *commandMessage) (bool, error) {
		switch usermsg.Cmd {
		case commandType_EnterRoom:
			if len(usermsg.Args) != 2 {
				return false, errCommandArgumentNotCorrect
			}
			roomName := usermsg.Args[0].(string)
			accid := usermsg.Args[1].(primitive.ObjectID)
			conn := entireConns[accid.Hex()]
			if conn == nil {
				return false, nil
			}
			findRoom(roomName, true).in(conn)

		case commandType_LeaveRoom:
			if len(usermsg.Args) != 2 {
				return false, errCommandArgumentNotCorrect
			}
			roomName := usermsg.Args[0].(string)
			accid := usermsg.Args[1].(primitive.ObjectID)
			room := findRoom(roomName, false)
			if room == nil {
				return false, errProcessFailed_NotInRoom
			}
			room.out(accid)
		}
		return true, nil
	}

	// 유저에게서 온 메세지, 소켓 연결/해체 처리
	unauthdata := []byte{0x03, 0xec}
	unauthdata = append(unauthdata, []byte("unauthorized")...)

	for {
		buffer := bytes.NewBuffer(make([]byte, 0, 1024))
		buffer.Reset()

		select {
		case <-ctx.Done():
			return

		case destroyedRoom := <-roomDestroyChan:
			delete(rooms, destroyedRoom)
			//go ws.callReceiver.OnRoomDestroyed(destroyedRoom)

		case usermsg := <-ws.localDeliveryChan:
			// 로컬에 connection이 있는지 먼저 확인해 보기 위한 채널
			// 없으면 publish한다.
			switch usermsg := usermsg.(type) {
			case *UpstreamMessage:
				if !processPrivateUpstreamMessage(usermsg) {
					// private메시지가 처리가 안됐으므로 publish
					enc := gob.NewEncoder(buffer)
					var err error
					if err = enc.Encode(usermsg); err == nil {
						_, err = ws.redisSync.Publish(context.Background(), ws.redisMsgChanName, buffer.Bytes()).Result()
					}

					if err != nil {
						logger.Println("gob.Encode or publish failed :", err)
					}
				}

			case *commandMessage:
				processed, err := processCommandMessage(usermsg)
				if err != nil {
					logger.Println("processCommandMessage returns err :", err)
					break
				}

				if !processed {
					var err error
					enc := gob.NewEncoder(buffer)
					if err = enc.Encode(usermsg); err == nil {
						_, err = ws.redisSync.Publish(context.Background(), ws.redisCmdChanName, buffer.Bytes()).Result()
					}
					if err != nil {
						logger.Println("gob.Encode or Publish failed :", err)
					}
				}
			}

		case usermsg := <-ws.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 {
					processPrivateUpstreamMessage(usermsg)
				}

			case *commandMessage:
				_, err := processCommandMessage(usermsg)
				if err != nil {
					logger.Println("processCommandMessage returns err :", err)
					break
				}

			default:
				logger.Println("usermsg is unknown type")
			}

		case c := <-ws.connInOutChan:
			if c.Conn == nil {
				delete(entireConns, c.sender.Accid.Hex())
				go ws.ClientDisconnected(c)
			} else {
				entireConns[c.sender.Accid.Hex()] = c
				go ws.ClientConnected(c)
			}

		case accid := <-ws.forceCloseChan:
			if conn := entireConns[accid.Hex()]; conn != nil {
				conn.MakeWriter().WriteControl(websocket.CloseMessage, unauthdata, time.Time{})
			}
		}
	}
}

func upgrade_core(ws *WebsocketHandler, conn *Conn, accid primitive.ObjectID, alias string) {
	newconn := &wsconn{
		Conn: conn,
		sender: &Sender{
			Alias: alias,
			Accid: accid,
		},
	}
	ws.connInOutChan <- newconn

	ws.connWaitGroup.Add(1)
	go func(c *wsconn, accid primitive.ObjectID, deliveryChan chan<- any) {
		defer func() {
			recover()
			ws.connWaitGroup.Done()
		}()

		for {
			messageType, r, err := c.innerConn.NextReader()
			if err != nil {
				if ce, ok := err.(*websocket.CloseError); ok {
					c.closeMessage = ce.Text
				}
				c.Close()
				break
			}

			if messageType == websocket.CloseMessage {
				closeMsg, _ := io.ReadAll(r)
				logger.Println("close message :", string(closeMsg))
				break
			}

			if messageType == websocket.BinaryMessage {
				var size [1]byte
				r.Read(size[:])
				cmd := make([]byte, size[0])
				r.Read(cmd)
				ws.Call(newconn.sender, string(cmd), r)
			}
		}
		c.Conn = nil
		if c.closeMessage != ForceShutdownCloseMessage {
			ws.connInOutChan <- c
		}
	}(newconn, accid, ws.deliveryChan)
}

func (ws *WebsocketHandler) 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
	}

	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)

	sk := r.Header.Get("AS-X-SESSION")
	authinfo, err := ws.sessionConsumer.Query(sk)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	if authinfo.Account != accid {
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	if authinfo.Invalidated() {
		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 {
		vt, _ := base64.StdEncoding.DecodeString(v)
		alias = string(vt)
	} else {
		alias = accid.Hex()
	}

	upgrade_core(ws, makeConn(conn), accid, alias)
}

func (ws *WebsocketHandler) 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")
	authinfo, err := ws.sessionConsumer.Query(sk)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		logger.Error("authorize query failed :", err)
		return
	}

	if authinfo.Account.IsZero() {
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	if authinfo.Invalidated() {
		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 {
		vt, _ := base64.StdEncoding.DecodeString(v)
		alias = string(vt)
	} else {
		alias = authinfo.Account.Hex()
	}

	upgrade_core(ws, makeConn(conn), authinfo.Account, alias)
}