gocommon/wshandler/wshandler_peer.go

package wshandler

import (
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"strings"
	"time"

	"go.mongodb.org/mongo-driver/bson/primitive"
	"repositories.action2quare.com/ayo/gocommon"
	"repositories.action2quare.com/ayo/gocommon/logger"
	"repositories.action2quare.com/ayo/gocommon/session"

	"github.com/gorilla/websocket"
)

type WebsocketPeerHandler[T PeerInterface] struct {
	WebsocketPeerApiBroker[T]
	sessionConsumer session.Consumer
}

func NewWebsocketPeerHandler[T PeerInterface](consumer session.Consumer) WebsocketPeerHandler[T] {
	return WebsocketPeerHandler[T]{
		sessionConsumer: consumer,
	}
}

func (ws *WebsocketPeerHandler[T]) 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 *WebsocketPeerHandler[T]) upgrade_core(conn *websocket.Conn, accid primitive.ObjectID, nonce uint32) {
	go func(c *websocket.Conn, accid primitive.ObjectID) {
		peer := ws.CreatePeer(accid)
		var closeReason string

		defer func() {
			peer.ClientDisconnected(closeReason)
		}()

		response := make([]byte, 255)
		for {
			response = response[:5]
			messageType, r, err := c.NextReader()
			if err != nil {
				if ce, ok := err.(*websocket.CloseError); ok {
					closeReason = ce.Text
				}
				c.Close()
				break
			}

			if messageType == websocket.CloseMessage {
				closeMsg, _ := io.ReadAll(r)
				closeReason = string(closeMsg)
				break
			}

			if messageType == websocket.BinaryMessage {
				var flag [1]byte
				r.Read(flag[:])
				if flag[0] == 0xff {
					// nonce
					r.Read(response[1:5])

					var size [1]byte
					r.Read(size[:])

					cmd := make([]byte, size[0])
					r.Read(cmd)

					result, err := ws.Call(peer, string(cmd), r)

					if err != nil {
						response[0] = 21 // 21 : Negative Ack
						response = append(response, []byte(err.Error())...)
					} else {
						response[0] = 6 // 6 : Acknowledgement

						switch result := result.(type) {
						case string:
							response = append(response, []byte(result)...)

						case int8, int16, int32, int64, uint8, uint16, uint32, uint64:
							response = append(response, []byte(fmt.Sprintf("%d", result))...)

						case float32, float64:
							response = append(response, []byte(fmt.Sprintf("%f", result))...)

						case []byte:
							response = append(response, result...)

						default:
							j, _ := json.Marshal(result)
							response = append(response, j...)
						}
					}
					c.WriteMessage(websocket.BinaryMessage, response)
				} else {
					cmd := make([]byte, flag[0])
					r.Read(cmd)
					ws.Call(peer, string(cmd), r)
				}
			}
		}
	}(conn, accid)
}

func (ws *WebsocketPeerHandler[T]) 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)

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

	nonce := rand.New(rand.NewSource(time.Now().UnixNano())).Uint32()
	ws.upgrade_core(conn, accid, nonce)
}

func (ws *WebsocketPeerHandler[T]) 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")
	logger.Println("WebsocketHandler.upgrade sk :", sk)
	authinfo, err := ws.sessionConsumer.Query(sk)
	if 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 {
	// 	vt, _ := base64.StdEncoding.DecodeString(v)
	// 	alias = string(vt)
	// } else {
	// 	alias = authinfo.Account.Hex()
	// }
	nonce := rand.New(rand.NewSource(time.Now().UnixNano())).Uint32()
	ws.upgrade_core(conn, authinfo.Account, nonce)
}