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