upstream-centrifugo/internal/unihttpstream/handler.go

package unihttpstream

import (
	"io"
	"net/http"
	"time"

	"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,
	}
}

const streamWriteTimeout = time.Second

func (h Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodOptions {
		w.Header().Set("Access-Control-Max-Age", "300")
		w.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS")
		w.WriteHeader(http.StatusNoContent)
		return
	}
	if r.Method != http.MethodPost {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}
	_, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "expected http.ResponseWriter to be http.Flusher", http.StatusInternalServerError)
		return
	}
	var req *protocol.ConnectRequest

	maxBytesSize := int64(h.config.MaxRequestBodySize)
	r.Body = http.MaxBytesReader(w, r.Body, maxBytesSize)
	connectRequestData, err := io.ReadAll(r.Body)
	if err != nil {
		log.Info().Err(err).Str("transport", transportName).Msg("error reading uni http stream request body")
		if len(connectRequestData) >= int(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
	}

	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 := newStreamTransport(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("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)

	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(msg)
				if err != nil {
					sendAck()
					return
				}
				_, err = w.Write([]byte("\n"))
				if err != nil {
					sendAck()
					return
				}
			}
			_ = rc.Flush()
			_ = rc.SetWriteDeadline(time.Time{})
			sendAck()
		}
	}
}