From b2c54afe1451f93e1fbbad257a151d8425cd308d Mon Sep 17 00:00:00 2001 From: Brad Fitzpatrick Date: Tue, 1 Nov 2016 15:24:11 +0000 Subject: [PATCH] net/http, net/http/httptest: make http2's TrailerPrefix work for http1 Go's http1 implementation originally had a mechanism to send HTTP trailers based on pre-declaring the trailer keys whose values you'd later let after the header was written. http2 copied the same mechanism, but it was found to be unsufficient for gRPC's wire protocol. A second trailer mechanism was added later (but only to http2) for handlers that want to send a trailer without knowing in advance they'd need to. Copy the same mechanism back to http1 and document it. Fixes #15754 Change-Id: I8c40d55e28b0e5b7087d3d1a904a392c56ee1f9b Reviewed-on: https://go-review.googlesource.com/32479 Reviewed-by: Ian Lance Taylor --- src/net/http/clientserver_test.go | 31 ++++++++++ src/net/http/httptest/recorder.go | 11 ++++ src/net/http/httptest/recorder_test.go | 2 + src/net/http/server.go | 79 ++++++++++++++++++++++---- 4 files changed, 111 insertions(+), 12 deletions(-) diff --git a/src/net/http/clientserver_test.go b/src/net/http/clientserver_test.go index f5524241897..d01e7558dca 100644 --- a/src/net/http/clientserver_test.go +++ b/src/net/http/clientserver_test.go @@ -1270,3 +1270,34 @@ func testNoSniffExpectRequestBody(t *testing.T, h2 bool) { t.Errorf("status code = %v; want %v", res.StatusCode, StatusUnauthorized) } } + +func TestServerUndeclaredTrailers_h1(t *testing.T) { testServerUndeclaredTrailers(t, h1Mode) } +func TestServerUndeclaredTrailers_h2(t *testing.T) { testServerUndeclaredTrailers(t, h2Mode) } +func testServerUndeclaredTrailers(t *testing.T, h2 bool) { + defer afterTest(t) + cst := newClientServerTest(t, h2, HandlerFunc(func(w ResponseWriter, r *Request) { + w.Header().Set("Foo", "Bar") + w.Header().Set("Trailer:Foo", "Baz") + w.(Flusher).Flush() + w.Header().Add("Trailer:Foo", "Baz2") + w.Header().Set("Trailer:Bar", "Quux") + })) + defer cst.close() + res, err := cst.c.Get(cst.ts.URL) + if err != nil { + t.Fatal(err) + } + if _, err := io.Copy(ioutil.Discard, res.Body); err != nil { + t.Fatal(err) + } + res.Body.Close() + delete(res.Header, "Date") + delete(res.Header, "Content-Type") + + if want := (Header{"Foo": {"Bar"}}); !reflect.DeepEqual(res.Header, want) { + t.Errorf("Header = %#v; want %#v", res.Header, want) + } + if want := (Header{"Foo": {"Baz", "Baz2"}, "Bar": {"Quux"}}); !reflect.DeepEqual(res.Trailer, want) { + t.Errorf("Trailer = %#v; want %#v", res.Trailer, want) + } +} diff --git a/src/net/http/httptest/recorder.go b/src/net/http/httptest/recorder.go index 24653031bd6..5f1aa6af479 100644 --- a/src/net/http/httptest/recorder.go +++ b/src/net/http/httptest/recorder.go @@ -203,6 +203,17 @@ func (rw *ResponseRecorder) Result() *http.Response { res.Trailer[k] = vv2 } } + for k, vv := range rw.HeaderMap { + if !strings.HasPrefix(k, http.TrailerPrefix) { + continue + } + if res.Trailer == nil { + res.Trailer = make(http.Header) + } + for _, v := range vv { + res.Trailer.Add(strings.TrimPrefix(k, http.TrailerPrefix), v) + } + } return res } diff --git a/src/net/http/httptest/recorder_test.go b/src/net/http/httptest/recorder_test.go index ff9b9911a86..9afba4e556a 100644 --- a/src/net/http/httptest/recorder_test.go +++ b/src/net/http/httptest/recorder_test.go @@ -207,6 +207,7 @@ func TestRecorder(t *testing.T) { w.Header().Set("Trailer-A", "valuea") w.Header().Set("Trailer-C", "valuec") w.Header().Set("Trailer-NotDeclared", "should be omitted") + w.Header().Set("Trailer:Trailer-D", "with prefix") }, check( hasStatus(200), @@ -216,6 +217,7 @@ func TestRecorder(t *testing.T) { hasTrailer("Trailer-A", "valuea"), hasTrailer("Trailer-C", "valuec"), hasNotTrailers("Non-Trailer", "Trailer-B", "Trailer-NotDeclared"), + hasTrailer("Trailer-D", "with prefix"), ), }, { diff --git a/src/net/http/server.go b/src/net/http/server.go index eae065f6730..c527ea8eeff 100644 --- a/src/net/http/server.go +++ b/src/net/http/server.go @@ -87,11 +87,25 @@ type Handler interface { // has returned. type ResponseWriter interface { // Header returns the header map that will be sent by - // WriteHeader. Changing the header after a call to - // WriteHeader (or Write) has no effect unless the modified - // headers were declared as trailers by setting the - // "Trailer" header before the call to WriteHeader (see example). - // To suppress implicit response headers, set their value to nil. + // WriteHeader. The Header map also is the mechanism with which + // Handlers can set HTTP trailers. + // + // Changing the header map after a call to WriteHeader (or + // Write) has no effect unless the modified headers are + // trailers. + // + // There are two ways to set Trailers. The preferred way is to + // predeclare in the headers which trailers you will later + // send by setting the "Trailer" header to the names of the + // trailer keys which will come later. In this case, those + // keys of the Header map are treated as if they were + // trailers. See the example. The second way, for trailer + // keys not known to the Handler until after the first Write, + // is to prefix the Header map keys with the TrailerPrefix + // constant value. See TrailerPrefix. + // + // To suppress implicit response headers (such as "Date"), set + // their value to nil. Header() Header // Write writes the data to the connection as part of an HTTP reply. @@ -358,13 +372,7 @@ func (cw *chunkWriter) close() { bw := cw.res.conn.bufw // conn's bufio writer // zero chunk to mark EOF bw.WriteString("0\r\n") - if len(cw.res.trailers) > 0 { - trailers := make(Header) - for _, h := range cw.res.trailers { - if vv := cw.res.handlerHeader[h]; len(vv) > 0 { - trailers[h] = vv - } - } + if trailers := cw.res.finalTrailers(); trailers != nil { trailers.Write(bw) // the writer handles noting errors } // final blank line after the trailers (whether @@ -432,6 +440,43 @@ type response struct { didCloseNotify int32 // atomic (only 0->1 winner should send) } +// TrailerPrefix is a magic prefix for ResponseWriter.Header map keys +// that, if present, signals that the map entry is actually for +// the response trailers, and not the response headers. The prefix +// is stripped after the ServeHTTP call finishes and the values are +// sent in the trailers. +// +// This mechanism is intended only for trailers that are not known +// prior to the headers being written. If the set of trailers is fixed +// or known before the header is written, the normal Go trailers mechanism +// is preferred: +// https://golang.org/pkg/net/http/#ResponseWriter +// https://golang.org/pkg/net/http/#example_ResponseWriter_trailers +const TrailerPrefix = "Trailer:" + +// finalTrailers is called after the Handler exits and returns a non-nil +// value if the Handler set any trailers. +func (w *response) finalTrailers() Header { + var t Header + for k, vv := range w.handlerHeader { + if strings.HasPrefix(k, TrailerPrefix) { + if t == nil { + t = make(Header) + } + t[strings.TrimPrefix(k, TrailerPrefix)] = vv + } + } + for _, k := range w.trailers { + if t == nil { + t = make(Header) + } + for _, v := range w.handlerHeader[k] { + t.Add(k, v) + } + } + return t +} + type atomicBool int32 func (b *atomicBool) isSet() bool { return atomic.LoadInt32((*int32)(b)) != 0 } @@ -1105,7 +1150,17 @@ func (cw *chunkWriter) writeHeader(p []byte) { } var setHeader extraHeader + // Don't write out the fake "Trailer:foo" keys. See TrailerPrefix. trailers := false + for k := range cw.header { + if strings.HasPrefix(k, TrailerPrefix) { + if excludeHeader == nil { + excludeHeader = make(map[string]bool) + } + excludeHeader[k] = true + trailers = true + } + } for _, v := range cw.header["Trailer"] { trailers = true foreachHeaderElement(v, cw.res.declareTrailer)