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 <iant@golang.org>

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)
+	}
+}

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
 }
 

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

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)