package app import ( "context" "crypto/rand" "net/http" "strings" "time" "encoders.orly/hex" "github.com/coder/websocket" "lol.mleku.dev/chk" "lol.mleku.dev/log" "utils.orly/units" ) const ( DefaultWriteWait = 10 * time.Second DefaultPongWait = 60 * time.Second DefaultPingWait = DefaultPongWait / 2 DefaultMaxMessageSize = 1 * units.Mb // CloseMessage denotes a close control message. The optional message // payload contains a numeric code and text. Use the FormatCloseMessage // function to format a close message payload. CloseMessage = 8 // PingMessage denotes a ping control message. The optional message payload // is UTF-8 encoded text. PingMessage = 9 // PongMessage denotes a pong control message. The optional message payload // is UTF-8 encoded text. PongMessage = 10 ) func (s *Server) HandleWebsocket(w http.ResponseWriter, r *http.Request) { remote := GetRemoteFromReq(r) log.T.F("handling websocket connection from %s", remote) if len(s.Config.IPWhitelist) > 0 { for _, ip := range s.Config.IPWhitelist { log.T.F("checking IP whitelist: %s", ip) if strings.HasPrefix(remote, ip) { log.T.F("IP whitelisted %s", remote) goto whitelist } } log.T.F("IP not whitelisted: %s", remote) return } whitelist: ctx, cancel := context.WithCancel(s.Ctx) defer cancel() var err error var conn *websocket.Conn if conn, err = websocket.Accept( w, r, &websocket.AcceptOptions{OriginPatterns: []string{"*"}}, ); chk.E(err) { return } conn.SetReadLimit(DefaultMaxMessageSize) defer conn.CloseNow() listener := &Listener{ ctx: ctx, Server: s, conn: conn, remote: remote, req: r, } chal := make([]byte, 32) rand.Read(chal) listener.challenge.Store([]byte(hex.Enc(chal))) ticker := time.NewTicker(DefaultPingWait) go s.Pinger(ctx, conn, ticker, cancel) defer func() { log.D.F("closing websocket connection from %s", remote) cancel() ticker.Stop() listener.publishers.Receive(&W{Cancel: true}) }() for { select { case <-ctx.Done(): return default: } var typ websocket.MessageType var msg []byte log.T.F("waiting for message from %s", remote) if typ, msg, err = conn.Read(ctx); chk.E(err) { if strings.Contains( err.Error(), "use of closed network connection", ) { return } status := websocket.CloseStatus(err) switch status { case websocket.StatusNormalClosure, websocket.StatusGoingAway, websocket.StatusNoStatusRcvd, websocket.StatusAbnormalClosure, websocket.StatusProtocolError: default: log.E.F("unexpected close error from %s: %v", remote, err) } return } if typ == PingMessage { if err = conn.Write(ctx, PongMessage, msg); chk.E(err) { return } continue } go listener.HandleMessage(msg, remote) } } func (s *Server) Pinger( ctx context.Context, conn *websocket.Conn, ticker *time.Ticker, cancel context.CancelFunc, ) { defer func() { cancel() ticker.Stop() }() var err error for { select { case <-ticker.C: if err = conn.Ping(ctx); chk.E(err) { return } case <-ctx.Done(): return } } }