internal/http3: make responseWriter behave closer to other http.ResponseWriter

While running net/http tests against our HTTP/3 implementation locally,
some tests fail due to slight behavior differences in responseWriter
compared to other http.ResponseWriter implementations:

- responseWriter does not return a 200 OK response if a server handler
  is completely empty.
- responseWriter does not have a Flush method, and therefore does not
  implement http.Flusher.

There are surely more differences, but these are straightforward to fix
right now.

For golang/go#70914

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

internal/http3/server.go

@@ -251,8 +251,9 @@ func (rw *responseWriter) Header() http.Header {
 	return rw.headers
 }
 
-// Caller must hold rw.mu.
-func (rw *responseWriter) writeHeaderLocked(statusCode int) {
+// Caller must hold rw.mu. If rw.wroteHeader is true, calling this method is a
+// no-op.
+func (rw *responseWriter) writeHeaderLockedOnce(statusCode int) {
 	// TODO: support trailer header.
 	if rw.wroteHeader {
 		return
@@ -283,21 +284,26 @@ func (rw *responseWriter) writeHeaderLocked(statusCode int) {
 func (rw *responseWriter) WriteHeader(statusCode int) {
 	rw.mu.Lock()
 	defer rw.mu.Unlock()
-	rw.writeHeaderLocked(statusCode)
+	rw.writeHeaderLockedOnce(statusCode)
 }
 
 func (rw *responseWriter) Write(b []byte) (int, error) {
 	rw.mu.Lock()
 	defer rw.mu.Unlock()
-	if !rw.wroteHeader {
-		rw.writeHeaderLocked(http.StatusOK)
-	}
+	rw.writeHeaderLockedOnce(http.StatusOK)
 	if rw.isHeadResp {
 		return 0, nil
 	}
 	return rw.bw.Write(b)
 }
 
+func (rw *responseWriter) Flush() {
+	rw.bw.st.Flush()
+}
+
 func (rw *responseWriter) close() error {
+	rw.mu.Lock()
+	defer rw.mu.Unlock()
+	rw.writeHeaderLockedOnce(http.StatusOK)
 	return rw.st.stream.Close()
 }

internal/http3/server_test.go

@@ -12,6 +12,7 @@ import (
 	"net/netip"
 	"testing"
 	"testing/synctest"
+	"time"
 
 	"golang.org/x/net/internal/quic/quicwire"
 	"golang.org/x/net/quic"
@@ -174,6 +175,65 @@ func TestServerHeadResponseNoBody(t *testing.T) {
 	})
 }
 
+func TestServerHandlerEmpty(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		ts := newTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// Empty handler should return a 200 OK
+		}))
+		tc := ts.connect()
+		tc.greet()
+
+		reqStream := tc.newStream(streamTypeRequest)
+		reqStream.writeHeaders(http.Header{":method": {http.MethodGet}})
+		synctest.Wait()
+		reqStream.wantHeaders(http.Header{":status": {"200"}})
+		reqStream.wantClosed("request is complete")
+	})
+}
+
+func TestServerHandlerFlushing(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		ts := newTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			time.Sleep(time.Second)
+			w.Write([]byte("first"))
+
+			time.Sleep(time.Second)
+			w.Write([]byte("second"))
+			w.(http.Flusher).Flush()
+
+			time.Sleep(time.Second)
+			w.Write([]byte("third"))
+		}))
+		tc := ts.connect()
+		tc.greet()
+
+		reqStream := tc.newStream(streamTypeRequest)
+		reqStream.writeHeaders(http.Header{":method": {http.MethodGet}})
+		synctest.Wait()
+
+		respBody := make([]byte, 100)
+
+		time.Sleep(time.Second)
+		synctest.Wait()
+		if n, err := reqStream.Read(respBody); err == nil {
+			t.Errorf("want no message yet, got %v bytes read", n)
+		}
+
+		time.Sleep(time.Second)
+		synctest.Wait()
+		if _, err := reqStream.Read(respBody); err != nil {
+			t.Errorf("failed to read partial response from server, got err: %v", err)
+		}
+
+		time.Sleep(time.Second)
+		synctest.Wait()
+		if _, err := reqStream.Read(respBody); err != io.EOF {
+			t.Errorf("expected EOF, got err: %v", err)
+		}
+		reqStream.wantClosed("request is complete")
+	})
+}
+
 type testServer struct {
 	t  testing.TB
 	s  *Server