99 lines
2.3 KiB
Go
99 lines
2.3 KiB
Go
package connect
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"git.bvbej.com/bvbej/base-golang/pkg/mux"
|
|
"github.com/gin-gonic/gin"
|
|
"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) 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())
|
|
}
|
|
}
|