sonet/internal/gateway_ws/server/http_server.go

package server

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"net/http"
	"sonet/pkg/utils/conver"
	"sonet/pkg/utils/resp"
	"time"
)

type HttpServer struct {
	connHandler *ConnHandler
}

func NewHttpServer(connHandler *ConnHandler) *HttpServer {
	return &HttpServer{
		connHandler: connHandler,
	}
}

func (s *HttpServer) Run(port int) error {
	server := gin.Default()
	server.GET("/ws", s.upgrade)
	return server.Run(fmt.Sprintf(":%d", port))
}

var (
	HandshakeTimeout = 3 * time.Second
	ReadDeadline     = 5 * time.Second
	WriteDeadline    = 5 * time.Second
	// PongWait Time allowed to read the next pong message from the peer.
	PongWait = 60 * time.Second

	MaxMessageSize  = conver.MustParseDataUnitInt("4M")
	ReadBufferSize  = conver.MustParseDataUnitInt("4Ki")
	WriteBufferSize = conver.MustParseDataUnitInt("4Ki")
)

func (s *HttpServer) upgrade(ctx *gin.Context) {
	upgrader := websocket.Upgrader{
		ReadBufferSize:   ReadBufferSize,
		WriteBufferSize:  WriteBufferSize,
		HandshakeTimeout: HandshakeTimeout,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
	conn, err := upgrader.Upgrade(ctx.Writer, ctx.Request, nil)
	if err != nil {
		ctx.JSON(http.StatusOK, resp.Error(err.Error()))
		return
	}

	// https://github.com/gorilla/websocket/blob/a68708917c6a4f06314ab4e52493cc61359c9d42/examples/chat/conn.go#L50
	conn.SetReadLimit(int64(MaxMessageSize))
	conn.SetPongHandler(func(string) error {
		return conn.SetReadDeadline(time.Now().Add(PongWait))
	})

	go s.connHandler.handleConn(conn)
}