internal/http3: add server to client trailer header support

This change implements support for Server to send trailer headers, and
for ClientConn to receive said trailer headers. This is just like
go.dev/cl/743600, but in the opposite direction.

The bulk of the implementation relies on the trailer header encoding and
decoding support that was added to bodyWriter and bodyReader
respectively in go.dev/cl/743600.

For golang/go#70914

Change-Id: I0efded4b1ac3e3c6b9479f18402e02e9e764d4a2
Reviewed-on: https://go-review.googlesource.com/c/net/+/744220
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Nicholas Husin <[email protected]>
Reviewed-by: Damien Neil <[email protected]>

Author: nicholashusin

internal/http3/body.go

@@ -10,11 +10,29 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"net/textproto"
+	"strings"
 	"sync"
 
 	"golang.org/x/net/http/httpguts"
 )
 
+// extractTrailerFromHeader extracts the "Trailer" header values from a header
+// map, and populates a trailer map with those values as keys. The extracted
+// header values will be canonicalized.
+func extractTrailerFromHeader(header, trailer http.Header) {
+	for _, names := range header["Trailer"] {
+		names = textproto.TrimString(names)
+		for name := range strings.SplitSeq(names, ",") {
+			name = textproto.CanonicalMIMEHeaderKey(textproto.TrimString(name))
+			if !httpguts.ValidTrailerHeader(name) {
+				continue
+			}
+			trailer[name] = nil
+		}
+	}
+}
+
 // A bodyWriter writes a request or response body to a stream
 // as a series of DATA frames.
 type bodyWriter struct {

internal/http3/roundtrip.go

@@ -154,10 +154,16 @@ func (cc *ClientConn) RoundTrip(req *http.Request) (_ *http.Response, err error)
 			if err != nil {
 				return nil, err
 			}
-			if contentLength != 0 && req.Method != http.MethodHead {
+
+			trailer := make(http.Header)
+			extractTrailerFromHeader(h, trailer)
+			delete(h, "Trailer")
+
+			if (contentLength != 0 && req.Method != http.MethodHead) || len(trailer) > 0 {
 				rt.respBody = &bodyReader{
-					st:     st,
-					remain: contentLength,
+					st:      st,
+					remain:  contentLength,
+					trailer: trailer,
 				}
 			} else {
 				rt.respBody = http.NoBody
@@ -169,6 +175,7 @@ func (cc *ClientConn) RoundTrip(req *http.Request) (_ *http.Response, err error)
 				StatusCode:    statusCode,
 				Status:        strconv.Itoa(statusCode) + " " + http.StatusText(statusCode),
 				ContentLength: contentLength,
+				Trailer:       trailer,
 				Body:          (*transportResponseBody)(rt),
 			}
 			// TODO: Automatic Content-Type: gzip decoding.

internal/http3/roundtrip_test.go

@@ -546,3 +546,90 @@ func TestRoundTripWriteTrailerNoBody(t *testing.T) {
 		st.wantClosed("request is complete")
 	})
 }
+
+func TestRoundTripReadTrailer(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		tc := newTestClientConn(t)
+		tc.greet()
+
+		var req *http.Request
+		req, _ = http.NewRequest("GET", "https://example.tld/", nil)
+		rt := tc.roundTrip(req)
+		st := tc.wantStream(streamTypeRequest)
+
+		st.wantHeaders(nil)
+		st.writeHeaders(http.Header{
+			":status": {"200"},
+			"Trailer": {"Server-Trailer-A, Server-Trailer-B", "server-trailer-c"}, // Should be canonicalized.
+		})
+		body := []byte("body from server")
+		st.writeData(body)
+		st.writeHeaders(http.Header{
+			"Server-Trailer-A": {"valuea"},
+			// Note that Server-Trailer-B is skipped.
+			"Server-Trailer-C":   {"valuec"},
+			"Undeclared-Trailer": {"undeclared"}, // Should be ignored.
+		})
+
+		rt.wantStatus(200)
+		// Trailer is stripped off from http.Response.Header and given in http.Response.Trailer.
+		rt.wantHeaders(http.Header{})
+		rt.wantTrailers(http.Header{
+			"Server-Trailer-A": nil,
+			"Server-Trailer-B": nil,
+			"Server-Trailer-C": nil,
+		})
+
+		// Trailer updated after reading the body to EOF.
+		rt.wantBody(body)
+		rt.wantTrailers(http.Header{
+			"Server-Trailer-A": {"valuea"},
+			"Server-Trailer-B": nil,
+			"Server-Trailer-C": {"valuec"},
+		})
+		st.wantClosed("request is complete")
+	})
+}
+
+func TestRoundTripReadTrailerNoBody(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		tc := newTestClientConn(t)
+		tc.greet()
+
+		var req *http.Request
+		req, _ = http.NewRequest("GET", "https://example.tld/", nil)
+		rt := tc.roundTrip(req)
+		st := tc.wantStream(streamTypeRequest)
+
+		st.wantHeaders(nil)
+		st.writeHeaders(http.Header{
+			":status":        {"200"},
+			"Content-Length": {"0"},
+			"Trailer":        {"Server-Trailer-A, Server-Trailer-B", "server-trailer-c"}, // Should be canonicalized.
+		})
+		st.writeHeaders(http.Header{
+			"Server-Trailer-A": {"valuea"},
+			// Note that Server-Trailer-B is skipped.
+			"Server-Trailer-C":   {"valuec"},
+			"Undeclared-Trailer": {"undeclared"}, // Should be ignored.
+		})
+
+		rt.wantStatus(200)
+		// Trailer is stripped off from http.Response.Header and given in http.Response.Trailer.
+		rt.wantHeaders(http.Header{"Content-Length": {"0"}})
+		rt.wantTrailers(http.Header{
+			"Server-Trailer-A": nil,
+			"Server-Trailer-B": nil,
+			"Server-Trailer-C": nil,
+		})
+
+		// Trailer updated after reading the empty body to EOF.
+		rt.wantBody(make([]byte, 0))
+		rt.wantTrailers(http.Header{
+			"Server-Trailer-A": {"valuea"},
+			"Server-Trailer-B": nil,
+			"Server-Trailer-C": {"valuec"},
+		})
+		st.wantClosed("request is complete")
+	})
+}

internal/http3/server.go

@@ -7,8 +7,11 @@ package http3
 import (
 	"context"
 	"io"
+	"maps"
 	"net/http"
+	"slices"
 	"strconv"
+	"strings"
 	"sync"
 
 	"golang.org/x/net/http/httpguts"
@@ -252,12 +255,14 @@ func (sc *serverConn) handleRequestStream(st *stream) error {
 	rw := &responseWriter{
 		st:         st,
 		headers:    make(http.Header),
+		trailer:    make(http.Header),
 		isHeadResp: req.Method == "HEAD",
 		bw: &bodyWriter{
 			st:     st,
 			remain: -1,
 			flush:  false,
 			name:   "response",
+			enc:    &sc.enc,
 		},
 	}
 	defer rw.close()
@@ -290,6 +295,7 @@ type responseWriter struct {
 	bw          *bodyWriter
 	mu          sync.Mutex
 	headers     http.Header
+	trailer     http.Header
 	wroteHeader bool // Non-1xx header has been (logically) written.
 	isHeadResp  bool // response is for a HEAD request.
 }
@@ -298,12 +304,37 @@ func (rw *responseWriter) Header() http.Header {
 	return rw.headers
 }
 
+// prepareTrailerForWriteLocked populates any pre-declared trailer header with
+// its value, and passes it to bodyWriter so it can be written after body EOF.
+// Caller must hold rw.mu.
+func (rw *responseWriter) prepareTrailerForWriteLocked() {
+	for name := range rw.trailer {
+		if val, ok := rw.headers[name]; ok {
+			rw.trailer[name] = val
+		} else {
+			delete(rw.trailer, name)
+		}
+	}
+	if len(rw.trailer) > 0 {
+		rw.bw.trailer = rw.trailer
+	}
+}
+
 // Caller must hold rw.mu. If rw.wroteHeader is true, calling this method is a
 // no-op.
 func (rw *responseWriter) writeHeaderLockedOnce(statusCode int) {
 	if rw.wroteHeader {
 		return
 	}
+
+	// If there is any Trailer declared in headers, save them so we know which
+	// trailers have been pre-declared. Also, write back the extracted value,
+	// which is canonicalized, to rw.Header for consistency.
+	if _, ok := rw.headers["Trailer"]; ok {
+		extractTrailerFromHeader(rw.headers, rw.trailer)
+		rw.headers.Set("Trailer", strings.Join(slices.Sorted(maps.Keys(rw.trailer)), ", "))
+	}
+
 	enc := &qpackEncoder{}
 	enc.init()
 	encHeaders := enc.encode(func(f func(itype indexType, name, value string)) {
@@ -356,5 +387,10 @@ func (rw *responseWriter) close() error {
 	rw.mu.Lock()
 	defer rw.mu.Unlock()
 	rw.writeHeaderLockedOnce(http.StatusOK)
+	rw.prepareTrailerForWriteLocked()
+
+	if err := rw.bw.Close(); err != nil {
+		return err
+	}
 	return rw.st.stream.Close()
 }

internal/http3/server_test.go

@@ -489,6 +489,68 @@ func TestServerHandlerReadTrailerNoBody(t *testing.T) {
 	})
 }
 
+func TestServerHandlerWriteTrailer(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		body := []byte("some body")
+		ts := newTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Trailer", "server-trailer-a, server-trailer-b") // Trailer header will be canonicalized.
+			w.Header().Add("Trailer", "Server-Trailer-C")
+
+			w.Write(body)
+
+			w.Header().Set("server-trailer-a", "valuea") // Trailer header will be canonicalized.
+			w.Header().Set("Server-Trailer-C", "valuec") // skipping B
+			w.Header().Set("Server-Trailer-Not-Declared", "should be omitted")
+		}))
+		tc := ts.connect()
+		tc.greet()
+
+		reqStream := tc.newStream(streamTypeRequest)
+		reqStream.writeHeaders(requestHeader(nil))
+		synctest.Wait()
+		reqStream.wantHeaders(http.Header{
+			":status": {"200"},
+			"Trailer": {"Server-Trailer-A, Server-Trailer-B, Server-Trailer-C"},
+		})
+		reqStream.wantData(body)
+		reqStream.wantHeaders(http.Header{
+			"Server-Trailer-A": {"valuea"},
+			"Server-Trailer-C": {"valuec"},
+		})
+		reqStream.wantClosed("request is complete")
+	})
+}
+
+func TestServerHandlerWriteTrailerNoBody(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		ts := newTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Trailer", "server-trailer-a, server-trailer-b") // Trailer header will be canonicalized.
+			w.Header().Add("Trailer", "Server-Trailer-C")
+
+			w.(http.Flusher).Flush()
+
+			w.Header().Set("server-trailer-a", "valuea") // Trailer header will be canonicalized.
+			w.Header().Set("Server-Trailer-C", "valuec") // skipping B
+			w.Header().Set("Server-Trailer-Not-Declared", "should be omitted")
+		}))
+		tc := ts.connect()
+		tc.greet()
+
+		reqStream := tc.newStream(streamTypeRequest)
+		reqStream.writeHeaders(requestHeader(nil))
+		synctest.Wait()
+		reqStream.wantHeaders(http.Header{
+			":status": {"200"},
+			"Trailer": {"Server-Trailer-A, Server-Trailer-B, Server-Trailer-C"},
+		})
+		reqStream.wantHeaders(http.Header{
+			"Server-Trailer-A": {"valuea"},
+			"Server-Trailer-C": {"valuec"},
+		})
+		reqStream.wantClosed("request is complete")
+	})
+}
+
 type testServer struct {
 	t  testing.TB
 	s  *Server

internal/http3/transport_test.go

@@ -452,6 +452,13 @@ func (rt *testRoundTrip) wantHeaders(want http.Header) {
 	}
 }
 
+func (rt *testRoundTrip) wantTrailers(want http.Header) {
+	rt.t.Helper()
+	if diff := diffHeaders(rt.response().Trailer, want); diff != "" {
+		rt.t.Fatalf("unexpected response trailers:\n%v", diff)
+	}
+}
+
 // readBody reads the contents of the response body.
 func (rt *testRoundTrip) readBody() ([]byte, error) {
 	t := rt.t