package unisse import ( "io" "net/http" "time" "github.com/centrifugal/centrifugo/v6/internal/convert" "github.com/centrifugal/centrifugo/v6/internal/logging" "github.com/centrifugal/centrifugo/v6/internal/tools" "github.com/centrifugal/centrifuge" "github.com/centrifugal/protocol" "github.com/rs/zerolog/log" "github.com/segmentio/encoding/json" ) type Handler struct { node *centrifuge.Node config Config pingPong centrifuge.PingPongConfig } func NewHandler(n *centrifuge.Node, c Config, pingPong centrifuge.PingPongConfig) *Handler { return &Handler{ node: n, config: c, pingPong: pingPong, } } // Since SSE is a GET request we are looking for connect request in URL params. // This should be a properly encoded JSON object. const connectUrlParam = "cf_connect" const streamWriteTimeout = time.Second func (h Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { _, ok := w.(http.Flusher) if !ok { http.Error(w, "expected http.ResponseWriter to be http.Flusher", http.StatusInternalServerError) return } var req *protocol.ConnectRequest if r.Method == http.MethodGet { 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") w.WriteHeader(http.StatusBadRequest) return } } else { req = &protocol.ConnectRequest{} } } else if r.Method == http.MethodPost { maxBytesSize := h.config.MaxRequestBodySize r.Body = http.MaxBytesReader(w, r.Body, int64(maxBytesSize)) connectRequestData, err := io.ReadAll(r.Body) if err != nil { log.Error().Err(err).Msg("error reading uni sse request body") if len(connectRequestData) >= maxBytesSize { w.WriteHeader(http.StatusRequestEntityTooLarge) return } return } _, err = json.Parse(connectRequestData, &req, json.ZeroCopy) if err != nil { if logging.Enabled(logging.DebugLevel) { log.Debug().Err(err).Str("transport", transportName).Msg("malformed connect request") } w.WriteHeader(http.StatusBadRequest) return } } else { w.WriteHeader(http.StatusMethodNotAllowed) return } if req == nil { // A literal `null` connect request decodes to a nil request with no error; // treat it as an empty request rather than passing nil to the connect // handler, which would nil-deref. req = &protocol.ConnectRequest{} } ack := make(chan struct{}) transport := newEventsourceTransport(r, h.pingPong, ack) c, closeFn, err := centrifuge.NewClient(r.Context(), h.node, transport) if err != nil { log.Error().Err(err).Str("transport", transportName).Msg("error create client") return } defer func() { _ = closeFn() }() defer close(transport.closedCh) // need to execute this after client 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()) } if h.config.ConnectCodeToHTTPResponse.Enabled { err = c.ProtocolConnectNoErrorToDisconnect(req) if err != nil { resp, ok := tools.ConnectErrorToToHTTPResponse(err, h.config.ConnectCodeToHTTPResponse.Transforms) if ok { w.WriteHeader(resp.StatusCode) _, _ = w.Write([]byte(resp.Body)) return } w.WriteHeader(http.StatusInternalServerError) _, _ = w.Write([]byte(http.StatusText(http.StatusInternalServerError))) return } } else { c.ProtocolConnect(req) } if r.ProtoMajor == 1 { // An endpoint MUST NOT generate an HTTP/2 message containing connection-specific header fields. // Source: RFC7540. w.Header().Set("Connection", "keep-alive") } w.Header().Set("X-Accel-Buffering", "no") w.Header().Set("Content-Type", "text/event-stream; charset=utf-8") w.Header().Set("Cache-Control", "private, no-cache, no-store, must-revalidate, max-age=0") w.Header().Set("Pragma", "no-cache") w.Header().Set("Expire", "0") w.WriteHeader(http.StatusOK) rc := http.NewResponseController(w) _ = rc.SetWriteDeadline(time.Now().Add(streamWriteTimeout)) _, err = w.Write([]byte("\r\n")) if err != nil { return } _ = rc.Flush() sendAck := func() { select { case ack <- struct{}{}: case <-r.Context().Done(): } } for { select { case <-r.Context().Done(): return case <-transport.disconnectCh: return case messages, messagesOK := <-transport.messages: if !messagesOK { sendAck() return } _ = rc.SetWriteDeadline(time.Now().Add(streamWriteTimeout)) for _, msg := range messages { _, err = w.Write(convert.StringToBytes("data: " + convert.BytesToString(msg) + "\n\n")) if err != nil { sendAck() return } } _ = rc.Flush() _ = rc.SetWriteDeadline(time.Time{}) sendAck() } } }