internal/http3: fix Write in Server Handler returning the wrong value

When changing the approach used in go.dev/cl/746760 (i.e. adjusting
slices, rather than keeping track of active part of the slices in
Write), the return value was unfortunately not updated properly.

This CL fixes the problem and modifies the test to also verify that
Write returns the correct value.

For golang/go#70914

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

Author: nicholashusin

internal/http3/server.go

@@ -414,11 +414,12 @@ func (rw *responseWriter) Write(b []byte) (int, error) {
 	// As a special case, we always want to save b to the buffer even when b is
 	// big if we had yet to write our header, so we can infer headers like
 	// "Content-Type" with as much information as possible.
+	initialBLen := len(b)
 	initialBufLen := len(rw.bb)
 	if !rw.wroteHeader || len(b) <= cap(rw.bb)-len(rw.bb) {
 		b = rw.bb.write(b)
 		if len(b) == 0 {
-			return len(b), nil
+			return initialBLen, nil
 		}
 	}
 
@@ -430,13 +431,13 @@ func (rw *responseWriter) Write(b []byte) (int, error) {
 	// 3. Reset the current body buffer so it can be used again.
 	rw.writeHeaderLockedOnce()
 	if rw.cannotHaveBody {
-		return len(b), nil
+		return initialBLen, nil
 	}
 	if n, err := rw.bw.write(rw.bb, b); err != nil {
 		return max(0, n-initialBufLen), err
 	}
 	rw.bb.discard()
-	return len(b), nil
+	return initialBLen, nil
 }
 
 func (rw *responseWriter) Flush() {

internal/http3/server_test.go

@@ -755,7 +755,13 @@ func TestServerBuffersBodyWrite(t *testing.T) {
 			ts := newTestServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 				for n := 0; n < tt.bodyLen; n += tt.writeSize {
 					data := slices.Repeat([]byte("a"), min(tt.writeSize, tt.bodyLen-n))
-					w.Write(data)
+					n, err := w.Write(data)
+					if err != nil {
+						t.Fatal(err)
+					}
+					if n != len(data) {
+						t.Errorf("got %v bytes when writing in server handler, want %v", n, len(data))
+					}
 					if tt.flushes {
 						w.(http.Flusher).Flush()
 					}