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

99 lines
2.3 KiB
Go
Raw Permalink Normal View History

2024-07-23 10:23:43 +08:00
package connect
import (
"context"
"errors"
"fmt"
2024-07-31 16:49:14 +08:00
"gitea.bvbej.com/bvbej/base-golang/pkg/mux"
2024-07-23 10:23:43 +08:00
"github.com/gin-gonic/gin"
"go.uber.org/zap"
"net/http"
"net/url"
"time"
2024-07-31 16:49:14 +08:00
"gitea.bvbej.com/bvbej/base-golang/pkg/websocket/peer"
2024-07-23 10:23:43 +08:00
"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)
HandlerFunc() mux.HandlerFunc
}
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 && !errors.Is(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))
}
func (ws *wsAcceptor) HandlerFunc() mux.HandlerFunc {
return func(c mux.Context) {
ws.GinHandle(c.Context())
}
}