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

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
	Stop(sessMgr *peer.SessionManager)
}

type wsAcceptor struct {
	server *http.Server
	logger *zap.Logger
	addr   string
}

func NewWsAcceptor(loggers *zap.Logger) WsAcceptor {
	return &wsAcceptor{logger: loggers}
}

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

	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)
			w.WriteHeader(http.StatusNotFound)
			_, _ = w.Write([]byte(http.StatusText(http.StatusNotFound)))
			return
		}
		sessMgr.Register <- peer.NewSession(newConnection(c, sessMgr))
	})

	ws.server = &http.Server{
		Addr:    urlObj.Host,
		Handler: nil,
	}
	ws.addr = addr
	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)
	}
}