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

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