diff --git a/adapters/humafiber/humafiber.go b/adapters/humafiber/humafiber.go index 535ae3bf..ae2e4aeb 100644 --- a/adapters/humafiber/humafiber.go +++ b/adapters/humafiber/humafiber.go @@ -1,11 +1,13 @@ package humafiber import ( + "bufio" "bytes" "context" "crypto/tls" "io" "mime/multipart" + "net" "net/http" "net/url" "strings" @@ -145,6 +147,16 @@ func (c *fiberWrapper) BodyWriter() io.Writer { return c.orig.RequestCtx() } +// StreamBody streams the response body via Fiber/fasthttp's stream writer. It +// is the optional streaming hook huma's SSE support uses because fasthttp can't +// flush the response writer synchronously from within the handler. +func (c *fiberWrapper) StreamBody(fn func(io.Writer)) { + rc := c.orig.RequestCtx() + rc.SetBodyStreamWriter(func(bw *bufio.Writer) { + fn(&fiberStreamWriter{bw: bw, conn: rc.Conn()}) + }) +} + func (c *fiberWrapper) TLS() *tls.ConnectionState { return c.orig.RequestCtx().TLSConnectionState() } @@ -262,3 +274,26 @@ func New(r *fiber.App, config huma.Config) huma.API { func NewWithGroup(r *fiber.App, g fiber.Router, config huma.Config) huma.API { return huma.NewAPI(config, &fiberAdapter{tester: r, router: g}) } + +// fiberStreamWriter adapts fasthttp's buffered stream writer to the io.Writer, +// http.Flusher, and write-deadline interfaces the streaming code expects. It is +// shared by the Fiber v2 and v3 adapters. +type fiberStreamWriter struct { + bw *bufio.Writer + conn net.Conn +} + +func (w *fiberStreamWriter) Write(p []byte) (int, error) { + return w.bw.Write(p) +} + +func (w *fiberStreamWriter) Flush() { + _ = w.bw.Flush() +} + +func (w *fiberStreamWriter) SetWriteDeadline(t time.Time) error { + if w.conn == nil { + return nil + } + return w.conn.SetWriteDeadline(t) +} diff --git a/adapters/humafiber/humafiber_test.go b/adapters/humafiber/humafiber_test.go index dab7ac67..ff9957c3 100644 --- a/adapters/humafiber/humafiber_test.go +++ b/adapters/humafiber/humafiber_test.go @@ -2,14 +2,60 @@ package humafiber import ( "context" + "io" "net/http" "net/http/httptest" "testing" "github.com/danielgtaylor/huma/v2" + "github.com/danielgtaylor/huma/v2/sse" "github.com/gofiber/fiber/v3" ) +// TestSSE exercises the Fiber-specific streaming path (the StreamBody hook / +// fasthttp SetBodyStreamWriter) so Server-Sent Events flush to the client +// instead of failing with "unable to flush". Regression test for #888. +func TestSSE(t *testing.T) { + app := fiber.New() + api := New(app, huma.DefaultConfig("Test", "1.0.0")) + + type Message struct { + Text string `json:"text"` + } + + sse.Register(api, huma.Operation{ + OperationID: "sse", + Method: http.MethodGet, + Path: "/sse", + }, map[string]any{ + "message": Message{}, + }, func(ctx context.Context, input *struct{}, send sse.Sender) { + _ = send.Data(Message{Text: "hello"}) + _ = send.Comment("heartbeat") + }) + + resp, err := app.Test(httptest.NewRequest(http.MethodGet, "http://example.com/sse", nil)) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + t.Fatalf("status = %d, want 200", resp.StatusCode) + } + if ct := resp.Header.Get("Content-Type"); ct != "text/event-stream" { + t.Fatalf("Content-Type = %q, want text/event-stream", ct) + } + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + want := "data: {\"text\":\"hello\"}\n\n: heartbeat\n\n" + if string(body) != want { + t.Fatalf("body = %q, want %q", body, want) + } +} + // TestNewContext covers the exported NewContext constructor. Unlike the other // adapters, the Fiber adapter's Handle wraps the context to expose fasthttp // user values, so it can't call NewContext itself; exercise it directly here. diff --git a/adapters/humafiber/humafiber_v2.go b/adapters/humafiber/humafiber_v2.go index bca6681c..4d9b2463 100644 --- a/adapters/humafiber/humafiber_v2.go +++ b/adapters/humafiber/humafiber_v2.go @@ -1,6 +1,7 @@ package humafiber import ( + "bufio" "bytes" "context" "crypto/tls" @@ -144,6 +145,16 @@ func (c *fiberV2Wrapper) BodyWriter() io.Writer { return c.orig.Context() } +// StreamBody streams the response body via Fiber/fasthttp's stream writer. It +// is the optional streaming hook huma's SSE support uses because fasthttp can't +// flush the response writer synchronously from within the handler. +func (c *fiberV2Wrapper) StreamBody(fn func(io.Writer)) { + rc := c.orig.Context() + rc.SetBodyStreamWriter(func(bw *bufio.Writer) { + fn(&fiberStreamWriter{bw: bw, conn: rc.Conn()}) + }) +} + func (c *fiberV2Wrapper) TLS() *tls.ConnectionState { return c.orig.Context().TLSConnectionState() } diff --git a/adapters/humafiber/humafiber_v2_test.go b/adapters/humafiber/humafiber_v2_test.go index 20dbdfe0..0670e9c6 100644 --- a/adapters/humafiber/humafiber_v2_test.go +++ b/adapters/humafiber/humafiber_v2_test.go @@ -9,10 +9,45 @@ import ( "testing" "github.com/danielgtaylor/huma/v2" + "github.com/danielgtaylor/huma/v2/sse" fiberV2 "github.com/gofiber/fiber/v2" "github.com/stretchr/testify/require" ) +// TestSSEV2 exercises the Fiber v2 streaming path (the StreamBody hook / +// fasthttp SetBodyStreamWriter) so Server-Sent Events flush to the client. +// Regression test for #888. +func TestSSEV2(t *testing.T) { + app := fiberV2.New() + api := NewV2(app, huma.DefaultConfig("Test", "1.0.0")) + + type Message struct { + Text string `json:"text"` + } + + sse.Register(api, huma.Operation{ + OperationID: "sse", + Method: http.MethodGet, + Path: "/sse", + }, map[string]any{ + "message": Message{}, + }, func(ctx context.Context, input *struct{}, send sse.Sender) { + _ = send.Data(Message{Text: "hello"}) + _ = send.Comment("heartbeat") + }) + + req, _ := http.NewRequest(http.MethodGet, "http://example.com/sse", nil) + resp, err := app.Test(req) + require.NoError(t, err) + defer resp.Body.Close() + + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, "text/event-stream", resp.Header.Get("Content-Type")) + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, "data: {\"text\":\"hello\"}\n\n: heartbeat\n\n", string(body)) +} + // TestFiberV2EachHeaderAndCookie ensures EachHeader yields one callback per // header value (not one per byte) so header-based helpers like huma.ReadCookie // continue to work under the Fiber v2 adapter. Regression test for #1055. diff --git a/sse/sse.go b/sse/sse.go index f8b71faa..c6a11eea 100644 --- a/sse/sse.go +++ b/sse/sse.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "os" "reflect" @@ -35,6 +36,16 @@ type writeDeadliner interface { SetWriteDeadline(time.Time) error } +// bodyStreamer is an optional interface an adapter's context may implement when +// its response writer can't be flushed synchronously from within the handler +// (e.g. fasthttp/Fiber). StreamBody is called with a writer used to stream the +// response body; the response stays open until the callback returns. The writer +// flushes on demand when it implements http.Flusher and honors write deadlines +// when it implements writeDeadliner. +type bodyStreamer interface { + StreamBody(func(io.Writer)) +} + // Message is a single SSE message. There is no `event` field as this is // handled by the `eventTypeMap` when registering the operation. type Message struct { @@ -147,110 +158,133 @@ func Register[I any](api huma.API, op huma.Operation, eventTypeMap map[string]an // Commit response headers immediately so the client's // EventSource.onopen fires without waiting for the first event. ctx.SetStatus(http.StatusOK) - bw := ctx.BodyWriter() - encoder := json.NewEncoder(bw) - - // Get the flusher/deadliner from the response writer if possible. - var flusher http.Flusher - flushCheck := bw - for { - if f, ok := flushCheck.(http.Flusher); ok { - flusher = f - break - } - if u, ok := flushCheck.(unwrapper); ok { - flushCheck = u.Unwrap() - } else { - break - } - } - if flusher != nil { - flusher.Flush() - } - var deadliner writeDeadliner - deadlineCheck := bw - for { - if d, ok := deadlineCheck.(writeDeadliner); ok { - deadliner = d - break - } - if u, ok := deadlineCheck.(unwrapper); ok { - deadlineCheck = u.Unwrap() - } else { - break - } - } + // Adapters whose response writer can't be flushed synchronously + // (e.g. Fiber/fasthttp) implement bodyStreamer and stream through + // a callback; everything else writes to BodyWriter directly. + if bs, ok := ctx.(bodyStreamer); ok { + bs.StreamBody(func(w io.Writer) { + stream(ctx.Context(), w, typeToEvent, input, f) + }) - flush := func() error { - if flusher == nil { - fmt.Fprintln(os.Stderr, "error: unable to flush") - return fmt.Errorf("unable to flush: %w", http.ErrNotSupported) - } - flusher.Flush() - return nil + return } - send := func(msg Message) error { - if deadliner != nil { - if err := deadliner.SetWriteDeadline(time.Now().Add(WriteTimeout)); err != nil { - fmt.Fprintf(os.Stderr, "warning: unable to set write deadline: %v\n", err) - } - } else { - fmt.Fprintln(os.Stderr, "write deadline not supported by underlying writer") - } - - // Write optional fields - if msg.ID > 0 { - bw.Write(fmt.Appendf(nil, "id: %d\n", msg.ID)) - } - if msg.Retry > 0 { - bw.Write(fmt.Appendf(nil, "retry: %d\n", msg.Retry)) - } - - if msg.Comment != "" { - // CR, LF, and CRLF are all SSE line terminators. Normalize - // them to LF so every line of the comment is re-emitted as - // its own ": " comment line and can't inject other fields. - comment := strings.ReplaceAll(msg.Comment, "\r\n", "\n") - comment = strings.ReplaceAll(comment, "\r", "\n") - for line := range strings.SplitSeq(comment, "\n") { - bw.Write(fmt.Appendf(nil, ": %s\n", line)) - } - } - - if msg.Data == nil { - bw.Write([]byte("\n")) - return flush() - } - - event, ok := typeToEvent[deref(reflect.TypeOf(msg.Data))] - if !ok { - fmt.Fprintf(os.Stderr, "error: unknown event type %v\n", reflect.TypeOf(msg.Data)) - debug.PrintStack() - } - if event != "" && event != "message" { - // `message` is the default, so no need to transmit it. - bw.Write([]byte("event: " + event + "\n")) - } - - // Write the message data. - if _, err := bw.Write([]byte("data: ")); err != nil { - return err - } - if err := encoder.Encode(msg.Data); err != nil { - bw.Write([]byte(`{"error": "encode error: `)) - bw.Write([]byte(err.Error())) - bw.Write([]byte("\"}\n\n")) - return err - } - bw.Write([]byte("\n")) - return flush() - } - - // Call the user-provided SSE handler. - f(ctx.Context(), input, send) + stream(ctx.Context(), ctx.BodyWriter(), typeToEvent, input, f) }, }, nil }) } + +// stream runs the SSE send loop against a single writer: it discovers the +// writer's flush and write-deadline support, builds the send function, and +// invokes the user handler. It is shared by the direct-write path and the +// bodyStreamer callback path. +func stream[I any](reqCtx context.Context, w io.Writer, typeToEvent map[reflect.Type]string, input *I, f func(ctx context.Context, input *I, send Sender)) { + encoder := json.NewEncoder(w) + + // Get the flusher/deadliner from the writer if possible. + var flusher http.Flusher + flushCheck := w + + for { + if fl, ok := flushCheck.(http.Flusher); ok { + flusher = fl + break + } + if u, ok := flushCheck.(unwrapper); ok { + flushCheck = u.Unwrap() + } else { + break + } + } + if flusher != nil { + flusher.Flush() + } + + var deadliner writeDeadliner + deadlineCheck := w + + for { + if d, ok := deadlineCheck.(writeDeadliner); ok { + deadliner = d + break + } + if u, ok := deadlineCheck.(unwrapper); ok { + deadlineCheck = u.Unwrap() + } else { + break + } + } + + flush := func() error { + if flusher == nil { + fmt.Fprintln(os.Stderr, "error: unable to flush") + return fmt.Errorf("unable to flush: %w", http.ErrNotSupported) + } + + flusher.Flush() + + return nil + } + + send := func(msg Message) error { + if deadliner != nil { + if err := deadliner.SetWriteDeadline(time.Now().Add(WriteTimeout)); err != nil { + fmt.Fprintf(os.Stderr, "warning: unable to set write deadline: %v\n", err) + } + } else { + fmt.Fprintln(os.Stderr, "write deadline not supported by underlying writer") + } + + // Write optional fields. + if msg.ID > 0 { + w.Write(fmt.Appendf(nil, "id: %d\n", msg.ID)) + } + if msg.Retry > 0 { + w.Write(fmt.Appendf(nil, "retry: %d\n", msg.Retry)) + } + + if msg.Comment != "" { + // CR, LF, and CRLF are all SSE line terminators. Normalize them to + // LF so every line of the comment is re-emitted as its own ": " + // comment line and can't inject other fields. + comment := strings.ReplaceAll(msg.Comment, "\r\n", "\n") + comment = strings.ReplaceAll(comment, "\r", "\n") + for line := range strings.SplitSeq(comment, "\n") { + w.Write(fmt.Appendf(nil, ": %s\n", line)) + } + } + + if msg.Data == nil { + w.Write([]byte("\n")) + return flush() + } + + event, ok := typeToEvent[deref(reflect.TypeOf(msg.Data))] + if !ok { + fmt.Fprintf(os.Stderr, "error: unknown event type %v\n", reflect.TypeOf(msg.Data)) + debug.PrintStack() + } + if event != "" && event != "message" { + // `message` is the default, so no need to transmit it. + w.Write([]byte("event: " + event + "\n")) + } + + // Write the message data. + if _, err := w.Write([]byte("data: ")); err != nil { + return err + } + if err := encoder.Encode(msg.Data); err != nil { + w.Write([]byte(`{"error": "encode error: `)) + w.Write([]byte(err.Error())) + w.Write([]byte("\"}\n\n")) + return err + } + w.Write([]byte("\n")) + return flush() + } + + // Call the user-provided SSE handler. + f(reqCtx, input, send) +} diff --git a/sse/sse_internal_test.go b/sse/sse_internal_test.go new file mode 100644 index 00000000..3a63fd89 --- /dev/null +++ b/sse/sse_internal_test.go @@ -0,0 +1,104 @@ +package sse + +import ( + "bytes" + "context" + "crypto/tls" + "io" + "mime/multipart" + "net/http" + "net/url" + "testing" + "time" + + "github.com/danielgtaylor/huma/v2" + "github.com/stretchr/testify/assert" +) + +// flushRecorder is an io.Writer that also implements http.Flusher, standing in +// for the stream writer an adapter hands to bodyStreamer.StreamBody. +type flushRecorder struct { + bytes.Buffer + flushes int +} + +func (f *flushRecorder) Flush() { f.flushes++ } + +// streamContext is a minimal huma.Context that also implements bodyStreamer, so +// the Fiber/fasthttp streaming path can be exercised without importing an +// adapter. It records what the SSE code writes to its stream writer. +type streamContext struct { + op *huma.Operation + header http.Header + status int + body *flushRecorder +} + +func (c *streamContext) Operation() *huma.Operation { return c.op } +func (c *streamContext) Context() context.Context { return context.Background() } +func (c *streamContext) TLS() *tls.ConnectionState { return nil } +func (c *streamContext) Version() huma.ProtoVersion { return huma.ProtoVersion{} } +func (c *streamContext) Method() string { return http.MethodGet } +func (c *streamContext) Host() string { return "" } +func (c *streamContext) RemoteAddr() string { return "" } +func (c *streamContext) URL() url.URL { return url.URL{} } +func (c *streamContext) Param(string) string { return "" } +func (c *streamContext) Query(string) string { return "" } +func (c *streamContext) Header(string) string { return "" } +func (c *streamContext) EachHeader(func(string, string)) {} +func (c *streamContext) BodyReader() io.Reader { return http.NoBody } +func (c *streamContext) GetMultipartForm() (*multipart.Form, error) { return nil, nil } +func (c *streamContext) SetReadDeadline(time.Time) error { return nil } +func (c *streamContext) SetStatus(code int) { c.status = code } +func (c *streamContext) Status() int { return c.status } +func (c *streamContext) SetHeader(name, value string) { c.header.Set(name, value) } +func (c *streamContext) AppendHeader(name, value string) { c.header.Add(name, value) } +func (c *streamContext) BodyWriter() io.Writer { return c.body } +func (c *streamContext) StreamBody(fn func(io.Writer)) { fn(c.body) } + +// streamAdapter is a huma.Adapter that drives a single registered operation +// through a streamContext, exercising the real Register wiring. +type streamAdapter struct { + handler func(huma.Context) + ctx *streamContext +} + +func (a *streamAdapter) Handle(_ *huma.Operation, handler func(huma.Context)) { + a.handler = handler +} + +func (a *streamAdapter) ServeHTTP(http.ResponseWriter, *http.Request) { + a.handler(a.ctx) +} + +// TestRegisterStreamsViaBodyStreamer covers the bodyStreamer branch in Register, +// which real adapters only reach via Fiber/fasthttp. It confirms events and +// comments are written to and flushed through the adapter's stream writer. +func TestRegisterStreamsViaBodyStreamer(t *testing.T) { + type message struct { + Text string `json:"text"` + } + + adapter := &streamAdapter{ctx: &streamContext{ + op: &huma.Operation{OperationID: "sse", Method: http.MethodGet, Path: "/sse"}, + header: http.Header{}, + body: &flushRecorder{}, + }} + api := huma.NewAPI(huma.DefaultConfig("Test", "1.0.0"), adapter) + + Register(api, huma.Operation{ + OperationID: "sse", + Method: http.MethodGet, + Path: "/sse", + }, map[string]any{"message": message{}}, func(_ context.Context, _ *struct{}, send Sender) { + _ = send.Data(message{Text: "hi"}) + _ = send.Comment("ping") + }) + + adapter.ServeHTTP(nil, nil) + + assert.Equal(t, "text/event-stream", adapter.ctx.header.Get("Content-Type")) + assert.Equal(t, http.StatusOK, adapter.ctx.status) + assert.Equal(t, "data: {\"text\":\"hi\"}\n\n: ping\n\n", adapter.ctx.body.String()) + assert.Positive(t, adapter.ctx.body.flushes) +}