internal/http3: make responseWriter.Flush write headers if not done yet

In net/http and x/net/http2, flushing an http.ResponseWriter will also
trigger headers to be written if none has been written yet at that
point. This behavior is sometimes relied on when implementing streaming
responses (see http://go.dev/cl/4552062); therefore, let's do this in
x/net/internal/http3 too.

For golang/go#70914

Change-Id: Ib914afc85df21a1cddd1d5019311d3f25d943f8a
Reviewed-on: https://go-review.googlesource.com/c/net/+/742160
Reviewed-by: Damien Neil <dneil@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Nicholas Husin <husin@google.com>

Author: nicholashusin

internal/http3/server.go

@@ -298,6 +298,9 @@ func (rw *responseWriter) Write(b []byte) (int, error) {
 }
 
 func (rw *responseWriter) Flush() {
+	rw.mu.Lock()
+	rw.writeHeaderLockedOnce(http.StatusOK)
+	rw.mu.Unlock()
 	rw.bw.st.Flush()
 }
 

internal/http3/server_test.go

@@ -65,7 +65,7 @@ func TestServerHeader(t *testing.T) {
 			"header-from-client": {"that", "should", "be", "echoed"},
 		})
 		synctest.Wait()
-		reqStream.wantHeaders(map[string][]string{
+		reqStream.wantHeaders(http.Header{
 			":status":            {"204"},
 			"Header-From-Client": {"that", "should", "be", "echoed"},
 		})
@@ -93,7 +93,7 @@ func TestServerPseudoHeader(t *testing.T) {
 		reqStream := tc.newStream(streamTypeRequest)
 		reqStream.writeHeaders(http.Header{":method": {"GET"}})
 		synctest.Wait()
-		reqStream.wantHeaders(map[string][]string{":status": {"321"}})
+		reqStream.wantHeaders(http.Header{":status": {"321"}})
 		reqStream.wantClosed("request is complete")
 	})
 }
@@ -114,7 +114,7 @@ func TestServerInvalidHeader(t *testing.T) {
 		reqStream := tc.newStream(streamTypeRequest)
 		reqStream.writeHeaders(http.Header{})
 		synctest.Wait()
-		reqStream.wantHeaders(map[string][]string{
+		reqStream.wantHeaders(http.Header{
 			":status":      {"200"},
 			"Valid-Name":   {"valid value"},
 			"Valid-Name-2": {"valid value 2"},
@@ -234,6 +234,35 @@ func TestServerHandlerFlushing(t *testing.T) {
 	})
 }
 
+func TestServerHandlerStreaming(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		stream := make(chan string)
+		ts := newTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// Flushing when we have not written anything yet implicitly calls
+			// w.WriteHeader(200).
+			w.(http.Flusher).Flush()
+			for str := range stream {
+				w.Write([]byte(str))
+				w.(http.Flusher).Flush()
+			}
+		}))
+		tc := ts.connect()
+		tc.greet()
+
+		reqStream := tc.newStream(streamTypeRequest)
+		reqStream.writeHeaders(http.Header{":method": {http.MethodGet}})
+		synctest.Wait()
+		reqStream.wantHeaders(http.Header{":status": {"200"}})
+
+		for _, data := range []string{"a", "bunch", "of", "things", "to", "stream"} {
+			stream <- data
+			reqStream.wantData([]byte(data))
+		}
+		close(stream)
+		reqStream.wantClosed("request is complete")
+	})
+}
+
 type testServer struct {
 	t  testing.TB
 	s  *Server