package connect import ( "context" "fmt" "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, sessMgr *peer.SessionManager) error Acceptor(path string, sessMgr *peer.SessionManager) Stop(sessMgr *peer.SessionManager) } type wsAcceptor struct { server *http.Server logger *zap.Logger } func NewWsAcceptor(loggers *zap.Logger) WsAcceptor { return &wsAcceptor{logger: loggers} } func (ws *wsAcceptor) Acceptor(path string, sessMgr *peer.SessionManager) { http.HandleFunc(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 } sessMgr.Register <- peer.NewSession(newConnection(c, sessMgr)) }) } func (ws *wsAcceptor) Start(addr string, sessMgr *peer.SessionManager) 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) } ws.Acceptor(urlObj.Path, 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(sessMgr *peer.SessionManager) { sessMgr.CloseAllSession() 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) } }