base-golang/pkg/websocket/peer/connect/acceptor.go

package connect

import (
	"context"
	"fmt"
	"github.com/gin-gonic/gin"
	"go.uber.org/zap"
	"net/http"
	"net/url"
	"time"

	"git.bvbej.com/bvbej/base-golang/pkg/websocket/peer"
	"github.com/gorilla/websocket"
)

var _ WsAcceptor = (*wsAcceptor)(nil)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type WsAcceptor interface {
	Start(addr string) error
	Stop()
	GinHandle(ctx *gin.Context)
}

type wsAcceptor struct {
	server  *http.Server
	sessMgr *peer.SessionManager
	logger  *zap.Logger
}

func NewWsAcceptor(sessMgr *peer.SessionManager, loggers *zap.Logger) WsAcceptor {
	return &wsAcceptor{
		sessMgr: sessMgr,
		logger:  loggers,
	}
}

func (ws *wsAcceptor) Start(addr string) error {
	urlObj, err := url.Parse(addr)
	if err != nil {
		return fmt.Errorf("websocket urlparse failed. url(%s) %v", addr, err)
	}
	if urlObj.Path == "" {
		return fmt.Errorf("websocket start failed. expect path in url to listen addr:%s", addr)
	}

	http.HandleFunc(urlObj.Path, func(w http.ResponseWriter, r *http.Request) {
		c, upgradeErr := upgrader.Upgrade(w, r, nil)
		if upgradeErr != nil {
			ws.logger.Sugar().Errorf("upgrade http failed: %s", upgradeErr)
			return
		}
		ws.sessMgr.Register <- peer.NewSession(NewConnection(c, ws.sessMgr))
	})

	ws.server = &http.Server{Addr: urlObj.Host}
	err = ws.server.ListenAndServe()
	if err != nil && err != http.ErrServerClosed {
		return fmt.Errorf("websocket ListenAndServe addr:%s failed:%v", addr, err)
	}

	return nil
}

func (ws *wsAcceptor) Stop() {
	ws.sessMgr.CloseAllSession()

	if ws.server != nil {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
		defer cancel()
		if err := ws.server.Shutdown(ctx); err != nil {
			ws.logger.Sugar().Errorf("server shutdown err:[%s]", err)
		}
	}
}

func (ws *wsAcceptor) GinHandle(ctx *gin.Context) {
	c, upgradeErr := upgrader.Upgrade(ctx.Writer, ctx.Request, nil)
	if upgradeErr != nil {
		ws.logger.Sugar().Errorf("upgrade http failed: %s", upgradeErr)
		return
	}
	ws.sessMgr.Register <- peer.NewSession(NewConnection(c, ws.sessMgr))
}