http2: add mechanism to send undeclared Trailers mid handler

This adds a way for http.Handlers to set trailers after the header has
already been flushed. This comment from the code explains:

// promoteUndeclaredTrailers permits http.Handlers to set trailers
// after the header has already been flushed. Because the Go
// ResponseWriter interface has no way to set Trailers (only the
// Header), and because we didn't want to expand the ResponseWriter
// interface, and because nobody used trailers, and because RFC 2616
// says you SHOULD (but not must) predeclare any trailers in the
// header, the official ResponseWriter rules said trailers in Go must
// be predeclared, and then we reuse the same ResponseWriter.Header()
// map to mean both Headers and Trailers.  When it's time to write the
// Trailers, we pick out the fields of Headers that were declared as
// trailers. That worked for a while, until we found the first major
// user of Trailers in the wild: gRPC (using them only over http2),
// and gRPC libraries permit setting trailers mid-stream without
// predeclarnig them. So: change of plans. We still permit the old
// way, but we also permit this hack: if a Header() key begins with
// "Trailer:", the suffix of that key is a Trailer. Because ':' is an
// invalid token byte anyway, there is no ambiguity. (And it's already
// filtered out) It's mildly hacky, but not terrible.

The official pre-declaring way still works. Example from net/http docs:
https://golang.org/pkg/net/http/#example_ResponseWriter_trailers
And ResponseWriter docs explaining it:
https://golang.org/pkg/net/http/#ResponseWriter

I don't want to add a new interface-assertable upgrade type (like
Hijacker or Flusher) for this because it's hurts composability and
makes everybody in the ecocsystem implement those, and two optional
interfaces is already bad enough. This is a weird enough feature
(Trailers by itself is weird enough), that I don't feel like a third
optional interface is worth it.

This code also filters invalid header fields (updates golang/go#14048),
since I had to update that code as part of this. But I want to later
return an error back to the user if possible. Or log something.

With this CL, all the grpc-go end2end tests pass with a new http2-based
Server implementation:
https://github.com/bradfitz/grpc-go/commit/a06f8f0593bfa8a0af72e57fb3916144f7b30121

Change-Id: I80f863d05a1810bd6f302f34932ad9df0a6646a6
Reviewed-on: https://go-review.googlesource.com/19131
Reviewed-by: Andrew Gerrand <adg@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

http2/server.go

@@ -51,6 +51,7 @@ import (
 	"os"
 	"reflect"
 	"runtime"
+	"sort"
 	"strconv"
 	"strings"
 	"sync"
@@ -1973,7 +1974,9 @@ func (rws *responseWriterState) declareTrailer(k string) {
 		// Forbidden by RFC 2616 14.40.
 		return
 	}
-	rws.trailers = append(rws.trailers, k)
+	if !strSliceContains(rws.trailers, k) {
+		rws.trailers = append(rws.trailers, k)
+	}
 }
 
 // writeChunk writes chunks from the bufio.Writer. But because
@@ -2041,6 +2044,10 @@ func (rws *responseWriterState) writeChunk(p []byte) (n int, err error) {
 		return 0, nil
 	}
 
+	if rws.handlerDone {
+		rws.promoteUndeclaredTrailers()
+	}
+
 	endStream := rws.handlerDone && !rws.hasTrailers()
 	if len(p) > 0 || endStream {
 		// only send a 0 byte DATA frame if we're ending the stream.
@@ -2061,6 +2068,53 @@ func (rws *responseWriterState) writeChunk(p []byte) (n int, err error) {
 	return len(p), nil
 }
 
+// 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:"
+
+// promoteUndeclaredTrailers permits http.Handlers to set trailers
+// after the header has already been flushed. Because the Go
+// ResponseWriter interface has no way to set Trailers (only the
+// Header), and because we didn't want to expand the ResponseWriter
+// interface, and because nobody used trailers, and because RFC 2616
+// says you SHOULD (but not must) predeclare any trailers in the
+// header, the official ResponseWriter rules said trailers in Go must
+// be predeclared, and then we reuse the same ResponseWriter.Header()
+// map to mean both Headers and Trailers.  When it's time to write the
+// Trailers, we pick out the fields of Headers that were declared as
+// trailers. That worked for a while, until we found the first major
+// user of Trailers in the wild: gRPC (using them only over http2),
+// and gRPC libraries permit setting trailers mid-stream without
+// predeclarnig them. So: change of plans. We still permit the old
+// way, but we also permit this hack: if a Header() key begins with
+// "Trailer:", the suffix of that key is a Trailer. Because ':' is an
+// invalid token byte anyway, there is no ambiguity. (And it's already
+// filtered out) It's mildly hacky, but not terrible.
+//
+// This method runs after the Handler is done and promotes any Header
+// fields to be trailers.
+func (rws *responseWriterState) promoteUndeclaredTrailers() {
+	for k, vv := range rws.handlerHeader {
+		if !strings.HasPrefix(k, TrailerPrefix) {
+			continue
+		}
+		trailerKey := strings.TrimPrefix(k, TrailerPrefix)
+		rws.declareTrailer(trailerKey)
+		rws.handlerHeader[http.CanonicalHeaderKey(trailerKey)] = vv
+	}
+	sort.Strings(rws.trailers)
+}
+
 func (w *responseWriter) Flush() {
 	rws := w.rws
 	if rws == nil {

http2/server_test.go

@@ -2647,7 +2647,13 @@ func testServerWritesTrailers(t *testing.T, withFlush bool) {
 		}
 		w.Header().Set("Server-Trailer-A", "valuea")
 		w.Header().Set("Server-Trailer-C", "valuec") // skipping B
+		// After a flush, random keys like Server-Surprise shouldn't show up:
 		w.Header().Set("Server-Surpise", "surprise! this isn't predeclared!")
+		// But we do permit promoting keys to trailers after a
+		// flush if they start with the magic
+		// otherwise-invalid "Trailer:" prefix:
+		w.Header().Set("Trailer:Post-Header-Trailer", "hi1")
+		w.Header().Set("Trailer:post-header-trailer2", "hi2")
 		w.Header().Set("Transfer-Encoding", "should not be included; Forbidden by RFC 2616 14.40")
 		w.Header().Set("Content-Length", "should not be included; Forbidden by RFC 2616 14.40")
 		w.Header().Set("Trailer", "should not be included; Forbidden by RFC 2616 14.40")
@@ -2689,6 +2695,8 @@ func testServerWritesTrailers(t *testing.T, withFlush bool) {
 			t.Fatalf("trailers HEADERS lacked END_HEADERS")
 		}
 		wanth = [][2]string{
+			{"post-header-trailer", "hi1"},
+			{"post-header-trailer2", "hi2"},
 			{"server-trailer-a", "valuea"},
 			{"server-trailer-c", "valuec"},
 		}
@@ -2699,6 +2707,37 @@ func testServerWritesTrailers(t *testing.T, withFlush bool) {
 	})
 }
 
+// validate transmitted header field names & values
+// golang.org/issue/14048
+func TestServerDoesntWriteInvalidHeaders(t *testing.T) {
+	testServerResponse(t, func(w http.ResponseWriter, r *http.Request) error {
+		w.Header().Add("OK1", "x")
+		w.Header().Add("Bad:Colon", "x") // colon (non-token byte) in key
+		w.Header().Add("Bad1\x00", "x")  // null in key
+		w.Header().Add("Bad2", "x\x00y") // null in value
+		return nil
+	}, func(st *serverTester) {
+		getSlash(st)
+		hf := st.wantHeaders()
+		if !hf.StreamEnded() {
+			t.Error("response HEADERS lacked END_STREAM")
+		}
+		if !hf.HeadersEnded() {
+			t.Fatal("response HEADERS didn't have END_HEADERS")
+		}
+		goth := st.decodeHeader(hf.HeaderBlockFragment())
+		wanth := [][2]string{
+			{":status", "200"},
+			{"ok1", "x"},
+			{"content-type", "text/plain; charset=utf-8"},
+			{"content-length", "0"},
+		}
+		if !reflect.DeepEqual(goth, wanth) {
+			t.Errorf("Header mismatch.\n got: %v\nwant: %v", goth, wanth)
+		}
+	})
+}
+
 func BenchmarkServerGets(b *testing.B) {
 	b.ReportAllocs()
 

http2/write.go

@@ -241,8 +241,18 @@ func encodeHeaders(enc *hpack.Encoder, h http.Header, keys []string) {
 	for _, k := range keys {
 		vv := h[k]
 		k = lowerHeader(k)
+		if !validHeaderFieldName(k) {
+			// TODO: return an error? golang.org/issue/14048
+			// For now just omit it.
+			continue
+		}
 		isTE := k == "transfer-encoding"
 		for _, v := range vv {
+			if !validHeaderFieldValue(v) {
+				// TODO: return an error? golang.org/issue/14048
+				// For now just omit it.
+				continue
+			}
 			// TODO: more of "8.1.2.2 Connection-Specific Header Fields"
 			if isTE && v != "trailers" {
 				continue