Repository: danielgtaylor/huma
Pull request #36 (2022-04-05 03:23:02)

Feat: Added WriteContent to huma context in order to support reading from streams directly into http response
This is useful for large data requests, especially ones that need to be chunked into multiple range requests like direct video streaming.  Since `http.ServeContent` already does exactly this and meticulously follows the associated RFCs, I decided to write this as a wrapper on top of that function.


Patch:
diff --git a/README.md b/README.md
index b684e44..e39b0bb 100644
--- a/README.md
+++ b/README.md
@@ -276,6 +276,60 @@ app.Resource("/exhaustive").Get("exhaustive", "Exhastive errors example",
 })
 ```
 
+## WriteContent
+
+Write contents allows you to write content in the provided ReadSeeker in the response.  It will handle Range, If-Match, If-Unmodified-Since, If-None-Match, If-Modified-Since, and if-Range requests for caching and large object partial content responses.
+
+example of a partial content response sending a large video file:
+```go
+package main
+
+import (
+	"net/http"
+	"os"
+
+	"github.com/danielgtaylor/huma"
+	"github.com/danielgtaylor/huma/cli"
+	"github.com/danielgtaylor/huma/responses"
+)
+
+func main() {
+	// Create a new router & CLI with default middleware.
+	app := cli.NewRouter("Video Content", "1.0.0")
+
+	// Declare the /vid resource and a GET operation on it.
+	app.Resource("/vid").Get("get-vid", "Get video content",
+		// ServeContent returns all the responses needed for WriteContent
+		responses.ServeContent()...,
+	).Run(func(ctx huma.Context) {
+		// This is he handler function for the operation. Write the response.
+		ctx.Header().Set("Content-Type", "video/mp4")
+		f, err := os.Open("./vid.mp4")
+		if err != nil {
+			ctx.WriteError(http.StatusInternalServerError, "Error while opening video")
+			return
+		}
+		defer f.Close()
+
+		fStat, err := os.Stat("./vid.mp4")
+		if err != nil {
+			ctx.WriteError(http.StatusInternalServerError, "Error while attempting to get Last Modified time")
+			return
+		}
+
+		// Note that name (ex: vid.mp4) and lastModified (ex: fState.ModTime()) are both optional
+		// if name is "" Content-Type must be set by the handler
+		// if lastModified is time.Time{} Modified request headers will not be respected by WriteContent
+		ctx.WriteContent("vid.mp4", f, fStat.ModTime())
+	})
+
+	// Run the CLI. When passed no arguments, it starts the server.
+	app.Run()
+}
+```
+
+Note that `WriteContent` does not automatically set the mime type.  You should set the `Content-Type` response header directly.  Also in order for `WriteContent` to respect the `Modified` headers you must call `SetContentLastModified`.  This is optional and if not set `WriteContent` will simply not respect the `Modified` request headers.
+
 ## Request Inputs
 
 Requests can have parameters and/or a body as input to the handler function. Like responses, inputs use standard Go structs but the tags are different. Here are the available tags:
diff --git a/context.go b/context.go
index c84102c..b524c47 100644
--- a/context.go
+++ b/context.go
@@ -4,9 +4,11 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"reflect"
 	"strings"
+	"time"
 
 	"github.com/danielgtaylor/huma/negotiation"
 	"github.com/fxamacker/cbor/v2"
@@ -65,6 +67,10 @@ type Context interface {
 	// content negotiation (e.g. JSON or CBOR). This must match the registered
 	// response status code & type.
 	WriteModel(status int, model interface{})
+
+	// WriteContent wraps http.ServeContent in order to handle serving streams
+	// it will handle Range and Modified (like If-Unmodified-Since) headers.
+	WriteContent(name string, content io.ReadSeeker, lastModified time.Time)
 }
 
 type hcontext struct {
@@ -141,7 +147,7 @@ func (c *hcontext) WriteHeader(status int) {
 
 func (c *hcontext) Write(data []byte) (int, error) {
 	if c.closed {
-		panic(fmt.Errorf("Trying to write to response after WriteModel or WriteError for %s %s", c.r.Method, c.r.URL.Path))
+		panic(fmt.Errorf("Trying to write to response after WriteModel, WriteError, or WriteContent for %s %s", c.r.Method, c.r.URL.Path))
 	}
 
 	return c.ResponseWriter.Write(data)
@@ -149,7 +155,7 @@ func (c *hcontext) Write(data []byte) (int, error) {
 
 func (c *hcontext) WriteError(status int, message string, errors ...error) {
 	if c.closed {
-		panic(fmt.Errorf("Trying to write to response after WriteModel or WriteError for %s %s", c.r.Method, c.r.URL.Path))
+		panic(fmt.Errorf("Trying to write to response after WriteModel, WriteError, or WriteContent for %s %s", c.r.Method, c.r.URL.Path))
 	}
 
 	details := []*ErrorDetail{}
@@ -186,7 +192,7 @@ func (c *hcontext) WriteError(status int, message string, errors ...error) {
 
 func (c *hcontext) WriteModel(status int, model interface{}) {
 	if c.closed {
-		panic(fmt.Errorf("Trying to write to response after WriteModel or WriteError for %s %s", c.r.Method, c.r.URL.Path))
+		panic(fmt.Errorf("Trying to write to response after WriteModel, WriteError, or WriteContent for %s %s", c.r.Method, c.r.URL.Path))
 	}
 
 	// Get the negotiated content type the client wants and we are willing to
@@ -373,3 +379,12 @@ func selectContentType(r *http.Request) string {
 
 	return ct
 }
+
+func (c *hcontext) WriteContent(name string, content io.ReadSeeker, lastModified time.Time) {
+	if c.closed {
+		panic(fmt.Errorf("Trying to write to response after WriteModel, WriteError, or WriteContent for %s %s", c.r.Method, c.r.URL.Path))
+	}
+
+	http.ServeContent(c.ResponseWriter, c.r, name, lastModified, content)
+	c.closed = true
+}
diff --git a/responses/responses.go b/responses/responses.go
index db69aa1..a25cf58 100644
--- a/responses/responses.go
+++ b/responses/responses.go
@@ -38,6 +38,11 @@ func NoContent() huma.Response {
 	return response(http.StatusNoContent)
 }
 
+// PartialContent HTTP 206 response
+func PartialContent() huma.Response {
+	return response(http.StatusPartialContent)
+}
+
 // MovedPermanently HTTP 301 response.
 func MovedPermanently() huma.Response {
 	return response(http.StatusMovedPermanently)
@@ -143,3 +148,22 @@ func GatewayTimeout() huma.Response {
 func String(status int) huma.Response {
 	return response(status).ContentType("text/plain")
 }
+
+// ServeContent returns a slice containing all valid responses for
+// context.ServeContent
+func ServeContent() []huma.Response {
+	return []huma.Response{
+		OK().Headers("Last-Modified", "Content-Type"),
+		PartialContent().Headers(
+			"Last-Modified",
+			"Content-Type",
+			"Content-Range",
+			"Content-Length",
+			"multipart/byteranges",
+			"Accept-Ranges",
+			"Content-Encoding"),
+		NotModified().Headers("Last-Modified"),
+		PreconditionFailed().Headers("Last-Modified", "Content-Type"),
+		InternalServerError(),
+	}
+}


Test patch:
diff --git a/context_test.go b/context_test.go
index 37a835f..c9d767c 100644
--- a/context_test.go
+++ b/context_test.go
@@ -1,12 +1,14 @@
 package huma
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/fxamacker/cbor/v2"
 	"github.com/goccy/go-yaml"
@@ -268,3 +270,100 @@ func TestValue(t *testing.T) {
 	r, _ := http.NewRequest(http.MethodGet, "/", nil)
 	handler(w, r)
 }
+
+func TestWriteContent(t *testing.T) {
+	app := newTestRouter()
+
+	b := []byte("Test Byte Data")
+
+	app.Resource("/content").Get("test", "Test",
+		NewResponse(200, "desc").Model(Response{}),
+	).Run(func(ctx Context) {
+		ctx.Header().Set("Content-Type", "application/octet-stream")
+		content := bytes.NewReader(b)
+		ctx.WriteContent("", content, time.Time{})
+	})
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest(http.MethodGet, "/content", nil)
+	app.ServeHTTP(w, req)
+
+	assert.Equal(t, 200, w.Result().StatusCode)
+	assert.Equal(t, "application/octet-stream", w.Header().Get("Content-Type"))
+	assert.Equal(t, b, w.Body.Bytes())
+}
+
+func TestWriteContentRespectsRange(t *testing.T) {
+	app := newTestRouter()
+
+	b := []byte("Test Byte Data")
+
+	app.Resource("/content").Get("test", "Test",
+		NewResponse(206, "desc").Model(Response{}),
+	).Run(func(ctx Context) {
+		ctx.Header().Set("Content-Type", "application/octet-stream")
+		content := bytes.NewReader(b)
+		ctx.WriteContent("", content, time.Time{})
+	})
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest(http.MethodGet, "/content", nil)
+	req.Header.Set("Range", "bytes=0-5")
+	app.ServeHTTP(w, req)
+
+	// confirms that Range is properly being forwarded to ServeContent.
+	// we'll assume more advanced range use cases are properly tested
+	// in the http library.
+	assert.Equal(t, w.Result().StatusCode, 206)
+	assert.Equal(t, []byte("Test B"), w.Body.Bytes())
+}
+
+func TestWriteContentLastModified(t *testing.T) {
+	app := newTestRouter()
+
+	b := []byte("Test Byte Data")
+	modTime := time.Now()
+
+	app.Resource("/content").Get("test", "Test",
+		NewResponse(206, "desc").Model(Response{}),
+	).Run(func(ctx Context) {
+		ctx.Header().Set("Content-Type", "application/octet-stream")
+		content := bytes.NewReader(b)
+		ctx.WriteContent("", content, modTime)
+	})
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest(http.MethodGet, "/content", nil)
+	app.ServeHTTP(w, req)
+
+	// confirms that modTime is properly being forwarded to ServeContent.
+	// We'll assume the more advanced modTime use cases are properly tested
+	// in http library.
+	strTime := modTime.UTC().Format(http.TimeFormat)
+
+	assert.Equal(t, strTime, w.Header().Get("Last-Modified"))
+
+}
+
+func TestWriteContentName(t *testing.T) {
+	app := newTestRouter()
+
+	b := []byte("Test Byte Data")
+
+	app.Resource("/content").Get("test", "Test",
+		NewResponse(206, "desc").Model(Response{}),
+	).Run(func(ctx Context) {
+
+		content := bytes.NewReader(b)
+		ctx.WriteContent("/path/with/content.mp4", content, time.Time{})
+	})
+
+	w := httptest.NewRecorder()
+	req, _ := http.NewRequest(http.MethodGet, "/content", nil)
+	app.ServeHTTP(w, req)
+
+	// confirms that name is properly being forwarded to ServeContent.
+	// We'll assume the more advanced modTime use cases are properly tested
+	// in http library.
+	assert.Equal(t, "video/mp4", w.Header().Get("Content-Type"))
+}
diff --git a/responses/responses_test.go b/responses/responses_test.go
index a61fe42..c390e72 100644
--- a/responses/responses_test.go
+++ b/responses/responses_test.go
@@ -19,6 +19,7 @@ var funcs = struct {
 		Created,
 		Accepted,
 		NoContent,
+		PartialContent,
 		MovedPermanently,
 		Found,
 		NotModified,
@@ -55,6 +56,7 @@ func TestResponses(t *testing.T) {
 		http.StatusCreated,
 		http.StatusAccepted,
 		http.StatusNoContent,
+		http.StatusPartialContent,
 		http.StatusMovedPermanently,
 		http.StatusFound,
 		http.StatusNotModified,