package connect import ( "context" "fmt" "go.uber.org/zap" "net/http" "net/url" "time" "github.com/gorilla/websocket" "git.bvbej.com/bvbej/base-golang/pkg/websocket/peer" ) var upgrader = websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true }, } type wsAcceptor struct { server *http.Server logger *zap.Logger } func (ws *wsAcceptor) Start(addr string, sessMgr *peer.SessionManager, loggers *zap.Logger) 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, err := upgrader.Upgrade(w, r, nil) if err != nil { fmt.Println(err) return } sessMgr.Register <- peer.NewSession(newConnection(c, sessMgr)) }) ws.logger = loggers ws.server = &http.Server{ Addr: urlObj.Host, Handler: nil, } 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.Error("server shutdown err", zap.Error(err)) } } func init() { peer.RegisterAcceptor("ws", new(wsAcceptor)) }