package uniws import ( "net/http" "sync" "time" "github.com/centrifugal/centrifugo/v6/internal/convert" "github.com/centrifugal/centrifugo/v6/internal/logging" "github.com/centrifugal/centrifugo/v6/internal/websocket" "github.com/centrifugal/centrifuge" "github.com/centrifugal/protocol" "github.com/rs/zerolog/log" "github.com/segmentio/encoding/json" ) // Handler handles WebSocket client connections. Usually WebSocket protocol // is a bidirectional connection between a client and a server for low-latency // communication. Here we utilize only one direction - giving users an additional // option for unidirectional transport. type Handler struct { node *centrifuge.Node upgrade *websocket.Upgrader config Config pingPong centrifuge.PingPongConfig } // Support passing connect request in URL params. // The value should be a properly encoded JSON object representing protocol.ConnectRequest. const connectUrlParam = "cf_connect" var writeBufferPool = &sync.Pool{} // NewHandler creates new Handler. func NewHandler( n *centrifuge.Node, c Config, CheckOrigin func(r *http.Request) bool, pingPong centrifuge.PingPongConfig, ) *Handler { upgrade := &websocket.Upgrader{ ReadBufferSize: c.ReadBufferSize, EnableCompression: c.Compression, } if c.UseWriteBufferPool { upgrade.WriteBufferPool = writeBufferPool } else { upgrade.WriteBufferSize = c.WriteBufferSize } if CheckOrigin != nil { upgrade.CheckOrigin = CheckOrigin } else { upgrade.CheckOrigin = sameHostOriginCheck() } return &Handler{ node: n, config: c, upgrade: upgrade, pingPong: pingPong, } } func (s *Handler) ServeHTTP(rw http.ResponseWriter, r *http.Request) { compression := s.config.Compression compressionLevel := s.config.CompressionLevel compressionMinSize := s.config.CompressionMinSize var req *protocol.ConnectRequest connectRequestString := r.URL.Query().Get(connectUrlParam) if connectRequestString != "" { _, err := json.Parse(convert.StringToBytes(connectRequestString), &req, json.ZeroCopy) if err != nil { log.Info().Err(err).Str("transport", transportName).Msg("error unmarshalling connect request") http.Error(rw, "invalid connect request", http.StatusBadRequest) return } } conn, _, err := s.upgrade.Upgrade(rw, r, nil) if err != nil { log.Error().Err(err).Str("transport", transportName).Msg("websocket upgrade error") return } if compression { err := conn.SetCompressionLevel(compressionLevel) if err != nil { log.Error().Err(err).Msg("websocket error setting compression level") } } framePingInterval := s.pingPong.PingInterval if framePingInterval <= 0 { // Always use frame ping/pong in unidirectional WebSocket transport. framePingInterval = DefaultWebsocketPingInterval } framePongTimeout := s.pingPong.PongTimeout if framePongTimeout <= 0 { framePongTimeout = DefaultWebsocketPongTimeout } writeTimeout := s.config.WriteTimeout.ToDuration() if writeTimeout == 0 { writeTimeout = DefaultWebsocketWriteTimeout } messageSizeLimit := s.config.MessageSizeLimit if messageSizeLimit == 0 { messageSizeLimit = DefaultWebsocketMessageSizeLimit } if messageSizeLimit > 0 { conn.SetReadLimit(int64(messageSizeLimit)) } if compression { decompressedMessageSizeLimit := s.config.DecompressedMessageSizeLimit if decompressedMessageSizeLimit == 0 { decompressedMessageSizeLimit = messageSizeLimit * DefaultWebsocketDecompressedMessageSizeLimitMultiplier } if decompressedMessageSizeLimit > 0 { conn.SetDecompressedReadLimit(int64(decompressedMessageSizeLimit)) } } // Separate goroutine for better GC of caller's data. go func() { // Record server-sent (outgoing) close codes for observability. Runs last // (registered first), after any teardown close frame has been written. // Client-supplied codes are skipped to keep metric cardinality bounded. defer func() { if code, incoming := conn.CloseCode(); code != 0 && !incoming { s.node.IncTransportOutgoingClose(transportName, code) } }() opts := websocketTransportOptions{ framePingInterval: framePingInterval, framePongTimeout: framePongTimeout, writeTimeout: writeTimeout, compressionMinSize: compressionMinSize, pingPongConfig: s.pingPong, joinMessages: s.config.JoinPushMessages, protoMajor: r.ProtoMajor, disableClosingHandshake: s.config.DisableClosingHandshake, disableDisconnectPush: s.config.DisableDisconnectPush, } graceCh := make(chan struct{}) transport := newWebsocketTransport(conn, opts, graceCh) select { case <-s.node.NotifyShutdown(): _ = transport.Close(centrifuge.DisconnectShutdown) return default: } ctxCh := make(chan struct{}) defer close(ctxCh) c, closeFn, err := centrifuge.NewClient(NewCancelContext(r.Context(), ctxCh), s.node, transport) if err != nil { log.Error().Err(err).Str("transport", transportName).Msg("error creating client") return } defer func() { _ = closeFn() }() if logging.Enabled(logging.DebugLevel) { log.Debug().Str("transport", transportName).Str("client", c.ID()).Msg("client connection established") defer func(started time.Time) { log.Debug().Str("transport", transportName).Str("client", c.ID()). Str("duration", time.Since(started).String()).Msg("client connection completed") }(time.Now()) } pongWait := framePingInterval + framePongTimeout _ = conn.SetReadDeadline(time.Now().Add(pongWait)) conn.SetPongHandler(func([]byte) error { _ = conn.SetReadDeadline(time.Now().Add(pongWait)) return nil }) waitClose := func() { // https://github.com/gorilla/websocket/issues/448 conn.SetPingHandler(nil) conn.SetPongHandler(nil) if s.config.DisableClosingHandshake { close(graceCh) return } _ = conn.SetReadDeadline(time.Now().Add(closeFrameWait)) for { if _, _, err := conn.NextReader(); err != nil { close(graceCh) return } } } if req == nil { _, data, err := conn.ReadMessage() if err != nil { waitClose() return } _, err = json.Parse(data, &req, json.ZeroCopy) if err != nil { log.Info().Err(err).Str("transport", transportName).Msg("error unmarshalling connect request") if !s.config.DisableClosingHandshake { err = conn.WriteControl( websocket.CloseMessage, websocket.FormatCloseMessage( int(centrifuge.DisconnectBadRequest.Code), centrifuge.DisconnectBadRequest.Reason, ), time.Now().Add(writeTimeout)) if err != nil { return } } waitClose() return } } if req == nil { // A literal `null` connect frame decodes to a nil request with no error; // treat it as an empty connect request. Passing nil to the connect // handler nil-derefs, and this runs in a bare goroutine (ServeHTTP has // already returned), so the panic would crash the whole process. req = &protocol.ConnectRequest{} } c.ProtocolConnect(req) for { _, _, err := conn.ReadMessage() if err != nil { break } } waitClose() }() }