first commit

pkg/websocket/peer/connect/acceptor.go

@@ -0,0 +1,98 @@
+package connect
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"git.bvbej.com/bvbej/base-golang/pkg/mux"
+	"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)
+	HandlerFunc() mux.HandlerFunc
+}
+
+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 && !errors.Is(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))
+}
+
+func (ws *wsAcceptor) HandlerFunc() mux.HandlerFunc {
+	return func(c mux.Context) {
+		ws.GinHandle(c.Context())
+	}
+}