Repository: quic-go/quic-go
Pull request #4412 (2024-04-05 06:25:00)

http3: sniff Content-Type when flushing the ResponseWriter
When caddy log is enabled, it can be observed `Content-Type` sniffing no longer works for http3, but works for h1/2. This is caused because when logging is enabled, caddy wraps all response writers with `responseRecorder` which will call `WriteHeader` [first](https://github.com/caddyserver/caddy/blob/master/modules/caddyhttp/responsewriter.go#L169).

Stdlib will try set `Content-Type` header in the same logic, for [http1](https://github.com/golang/go/blob/master/src/net/http/server.go#L1439), `*chunkWriter` is wrapped by `*response` to write the actual data. [http2](https://github.com/golang/net/blob/master/http2/server.go#L2630) will set this header in the same way.

Because this header is set once, it won't be overwritten automatically

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/quic-go/quic-go/pull/4412?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 84.87%. Comparing base [(`594440b`)](https://app.codecov.io/gh/quic-go/quic-go/commit/594440b04c7f385ad85fa53914dd02ce809da2ca?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go) to head [(`e3b1cc2`)](https://app.codecov.io/gh/quic-go/quic-go/pull/4412?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).
> Report is 88 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #4412      +/-   ##
==========================================
+ Coverage   84.09%   84.87%   +0.79%     
==========================================
  Files         150      153       +3     
  Lines       15401    14383    -1018     
==========================================
- Hits        12950    12207     -743     
+ Misses       1950     1672     -278     
- Partials      501      504       +3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/quic-go/quic-go/pull/4412?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=quic-go).

marten-seemann: @WeidiDeng The large refactor for HTTP Datagram support (#4452) was merged. Do you want to update this PR?

Patch:
diff --git a/connection.go b/connection.go
index 9c8821d9..b1e9aaee 100644
--- a/connection.go
+++ b/connection.go
@@ -2358,10 +2358,9 @@ func (s *connection) SendDatagram(p []byte) error {
 	}
 
 	f := &wire.DatagramFrame{DataLenPresent: true}
-	if protocol.ByteCount(len(p)) > f.MaxDataLen(s.peerParams.MaxDatagramFrameSize, s.version) {
-		return &DatagramTooLargeError{
-			PeerMaxDatagramFrameSize: int64(s.peerParams.MaxDatagramFrameSize),
-		}
+	maxDataLen := f.MaxDataLen(s.peerParams.MaxDatagramFrameSize, s.version)
+	if protocol.ByteCount(len(p)) > maxDataLen {
+		return &DatagramTooLargeError{MaxDatagramPayloadSize: int64(maxDataLen)}
 	}
 	f.Data = make([]byte, len(p))
 	copy(f.Data, p)
diff --git a/errors.go b/errors.go
index fda3c924..3fe1e0a9 100644
--- a/errors.go
+++ b/errors.go
@@ -64,7 +64,7 @@ func (e *StreamError) Error() string {
 
 // DatagramTooLargeError is returned from Connection.SendDatagram if the payload is too large to be sent.
 type DatagramTooLargeError struct {
-	PeerMaxDatagramFrameSize int64
+	MaxDatagramPayloadSize int64
 }
 
 func (e *DatagramTooLargeError) Is(target error) bool {
diff --git a/http3/body.go b/http3/body.go
index f496a959..fa023ce4 100644
--- a/http3/body.go
+++ b/http3/body.go
@@ -8,15 +8,6 @@ import (
 	"github.com/quic-go/quic-go"
 )
 
-// The HTTPStreamer allows taking over a HTTP/3 stream. The interface is implemented by:
-// * for the server: the http.Request.Body
-// * for the client: the http.Response.Body
-// On the client side, the stream will be closed for writing, unless the DontCloseRequestStream RoundTripOpt was set.
-// When a stream is taken over, it's the caller's responsibility to close the stream.
-type HTTPStreamer interface {
-	HTTPStream() Stream
-}
-
 // A Hijacker allows hijacking of the stream creating part of a quic.Session from a http.Response.Body.
 // It is used by WebTransport to create WebTransport streams after a session has been established.
 type Hijacker interface {
@@ -32,8 +23,6 @@ type body struct {
 	remainingContentLength int64
 	violatedContentLength  bool
 	hasContentLength       bool
-
-	wasHijacked bool // set when HTTPStream is called
 }
 
 func newBody(str *stream, contentLength int64) *body {
@@ -45,15 +34,7 @@ func newBody(str *stream, contentLength int64) *body {
 	return b
 }
 
-func (r *body) HTTPStream() Stream {
-	r.wasHijacked = true
-	return r.str
-}
-
 func (r *body) StreamID() quic.StreamID { return r.str.StreamID() }
-func (r *body) wasStreamHijacked() bool {
-	return r.wasHijacked
-}
 
 func (r *body) checkContentLengthViolation() error {
 	if !r.hasContentLength {
@@ -97,10 +78,7 @@ type requestBody struct {
 	getSettings  func() *Settings
 }
 
-var (
-	_ io.ReadCloser = &requestBody{}
-	_ HTTPStreamer  = &requestBody{}
-)
+var _ io.ReadCloser = &requestBody{}
 
 func newRequestBody(str *stream, contentLength int64, connCtx context.Context, rcvdSettings <-chan struct{}, getSettings func() *Settings) *requestBody {
 	return &requestBody{
diff --git a/http3/response_writer.go b/http3/response_writer.go
index 402938ba..8638ec57 100644
--- a/http3/response_writer.go
+++ b/http3/response_writer.go
@@ -12,6 +12,13 @@ import (
 	"github.com/quic-go/qpack"
 )
 
+// The HTTPStreamer allows taking over a HTTP/3 stream. The interface is implemented the http.Response.Body.
+// On the client side, the stream will be closed for writing, unless the DontCloseRequestStream RoundTripOpt was set.
+// When a stream is taken over, it's the caller's responsibility to close the stream.
+type HTTPStreamer interface {
+	HTTPStream() Stream
+}
+
 // The maximum length of an encoded HTTP/3 frame header is 16:
 // The frame has a type and length field, both QUIC varints (maximum 8 bytes in length)
 const frameHeaderLen = 16
@@ -36,6 +43,8 @@ type responseWriter struct {
 	headerWritten  bool  // set once the response header has been serialized to the stream
 	isHead         bool
 
+	hijacked bool // set on HTTPStream is called
+
 	logger *slog.Logger
 }
 
@@ -43,6 +52,7 @@ var (
 	_ http.ResponseWriter = &responseWriter{}
 	_ http.Flusher        = &responseWriter{}
 	_ Hijacker            = &responseWriter{}
+	_ HTTPStreamer        = &responseWriter{}
 )
 
 func newResponseWriter(str *stream, conn Connection, isHead bool, logger *slog.Logger) *responseWriter {
@@ -102,23 +112,24 @@ func (w *responseWriter) WriteHeader(status int) {
 	}
 }
 
+func (w *responseWriter) sniffContentType(p []byte) {
+	// If no content type, apply sniffing algorithm to body.
+	// We can't use `w.header.Get` here since if the Content-Type was set to nil, we shouldn't do sniffing.
+	_, haveType := w.header["Content-Type"]
+
+	// If the Transfer-Encoding or Content-Encoding was set and is non-blank,
+	// we shouldn't sniff the body.
+	hasTE := w.header.Get("Transfer-Encoding") != ""
+	hasCE := w.header.Get("Content-Encoding") != ""
+	if !hasCE && !haveType && !hasTE && len(p) > 0 {
+		w.header.Set("Content-Type", http.DetectContentType(p))
+	}
+}
+
 func (w *responseWriter) Write(p []byte) (int, error) {
 	bodyAllowed := bodyAllowedForStatus(w.status)
 	if !w.headerComplete {
-		// If body is not allowed, we don't need to (and we can't) sniff the content type.
-		if bodyAllowed {
-			// If no content type, apply sniffing algorithm to body.
-			// We can't use `w.header.Get` here since if the Content-Type was set to nil, we shoundn't do sniffing.
-			_, haveType := w.header["Content-Type"]
-
-			// If the Transfer-Encoding or Content-Encoding was set and is non-blank,
-			// we shouldn't sniff the body.
-			hasTE := w.header.Get("Transfer-Encoding") != ""
-			hasCE := w.header.Get("Content-Encoding") != ""
-			if !hasCE && !haveType && !hasTE && len(p) > 0 {
-				w.header.Set("Content-Type", http.DetectContentType(p))
-			}
-		}
+		w.sniffContentType(p)
 		w.WriteHeader(http.StatusOK)
 		bodyAllowed = true
 	}
@@ -148,6 +159,7 @@ func (w *responseWriter) Write(p []byte) (int, error) {
 
 func (w *responseWriter) doWrite(p []byte) (int, error) {
 	if !w.headerWritten {
+		w.sniffContentType(w.smallResponseBuf)
 		if err := w.writeHeader(w.status); err != nil {
 			return 0, maybeReplaceError(err)
 		}
@@ -220,6 +232,14 @@ func (w *responseWriter) Flush() {
 	}
 }
 
+func (w *responseWriter) HTTPStream() Stream {
+	w.hijacked = true
+	w.Flush()
+	return w.str
+}
+
+func (w *responseWriter) wasStreamHijacked() bool { return w.hijacked }
+
 func (w *responseWriter) Connection() Connection {
 	return w.conn
 }
diff --git a/http3/server.go b/http3/server.go
index 3f53db2c..f7ca9bab 100644
--- a/http3/server.go
+++ b/http3/server.go
@@ -572,7 +572,7 @@ func (s *Server) handleRequest(conn *connection, str quic.Stream, datagrams *dat
 		handler.ServeHTTP(r, req)
 	}()
 
-	if body.wasStreamHijacked() {
+	if r.wasStreamHijacked() {
 		return
 	}
 


Test patch:
diff --git a/connection_test.go b/connection_test.go
index c6d9d24e..7f9468d9 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -2477,6 +2477,40 @@ var _ = Describe("Connection", func() {
 		})
 	})
 
+	Context("datagrams", func() {
+		It("doesn't allow datagrams if the peer didn't enable support", func() {
+			conn.peerParams = &wire.TransportParameters{MaxDatagramFrameSize: 0}
+			Expect(conn.SendDatagram(make([]byte, 200))).To(MatchError("datagram support disabled"))
+		})
+
+		It("sends a datagram", func() {
+			conn.peerParams = &wire.TransportParameters{MaxDatagramFrameSize: 1000}
+			Expect(conn.SendDatagram([]byte("foobar"))).To(Succeed())
+			f := conn.datagramQueue.Peek()
+			Expect(f).ToNot(BeNil())
+			Expect(f.Data).To(Equal([]byte("foobar")))
+		})
+
+		It("says when a datagram is too big", func() {
+			conn.peerParams = &wire.TransportParameters{MaxDatagramFrameSize: 1000}
+			err := conn.SendDatagram(make([]byte, 2000))
+			Expect(err).To(HaveOccurred())
+			Expect(err).To(BeAssignableToTypeOf(&DatagramTooLargeError{}))
+			derr := err.(*DatagramTooLargeError)
+			Expect(derr.MaxDatagramPayloadSize).To(BeNumerically("<", 1000))
+			fmt.Println(derr.MaxDatagramPayloadSize)
+			Expect(conn.SendDatagram(make([]byte, derr.MaxDatagramPayloadSize))).To(Succeed())
+		})
+
+		It("receives datagrams", func() {
+			conn.config.EnableDatagrams = true
+			conn.datagramQueue.HandleDatagramFrame(&wire.DatagramFrame{Data: []byte("foobar")})
+			data, err := conn.ReceiveDatagram(context.Background())
+			Expect(err).ToNot(HaveOccurred())
+			Expect(data).To(Equal([]byte("foobar")))
+		})
+	})
+
 	It("returns the local address", func() {
 		Expect(conn.LocalAddr()).To(Equal(localAddr))
 	})
diff --git a/http3/response_writer_test.go b/http3/response_writer_test.go
index f49e0730..f68e5ace 100644
--- a/http3/response_writer_test.go
+++ b/http3/response_writer_test.go
@@ -137,9 +137,10 @@ var _ = Describe("Response Writer", func() {
 
 		// According to the spec, headers sent in the informational response must also be included in the final response
 		fields = decodeHeader(strBuf)
-		Expect(fields).To(HaveLen(3))
+		Expect(fields).To(HaveLen(4))
 		Expect(fields).To(HaveKeyWithValue(":status", []string{"200"}))
 		Expect(fields).To(HaveKey("date"))
+		Expect(fields).To(HaveKey("content-type"))
 		Expect(fields).To(HaveKeyWithValue("link", []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script"}))
 
 		Expect(getData(strBuf)).To(Equal([]byte("foobar")))
diff --git a/http3/server_test.go b/http3/server_test.go
index 892006db..1a30e9f7 100644
--- a/http3/server_test.go
+++ b/http3/server_test.go
@@ -207,6 +207,26 @@ var _ = Describe("Server", func() {
 			Expect(hfs).To(HaveLen(4))
 		})
 
+		It("sets Content-Type when WriteHeader is called but response is not flushed", func() {
+			s.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusNotFound)
+				w.Write([]byte("<html></html>"))
+			})
+
+			responseBuf := &bytes.Buffer{}
+			setRequest(encodeRequest(exampleGetRequest))
+			str.EXPECT().Context().Return(reqContext)
+			str.EXPECT().Write(gomock.Any()).DoAndReturn(responseBuf.Write).AnyTimes()
+			str.EXPECT().CancelRead(gomock.Any())
+			str.EXPECT().Close()
+
+			s.handleRequest(conn, str, nil, qpackDecoder)
+			hfs := decodeHeader(responseBuf)
+			Expect(hfs).To(HaveKeyWithValue(":status", []string{"404"}))
+			Expect(hfs).To(HaveKeyWithValue("content-length", []string{"13"}))
+			Expect(hfs).To(HaveKeyWithValue("content-type", []string{"text/html; charset=utf-8"}))
+		})
+
 		It("not sets Content-Length when the handler flushes to the client", func() {
 			s.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 				w.Write([]byte("foobar"))
@@ -581,7 +601,7 @@ var _ = Describe("Server", func() {
 				handlerCalled := make(chan struct{})
 				s.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 					defer close(handlerCalled)
-					r.Body.(HTTPStreamer).HTTPStream()
+					w.(HTTPStreamer).HTTPStream()
 					str.Write([]byte("foobar"))
 				})
 
@@ -590,10 +610,26 @@ var _ = Describe("Server", func() {
 				b = append(b, []byte("foobar")...)
 				setRequest(append(requestData, b...))
 				str.EXPECT().Context().Return(reqContext)
-				str.EXPECT().Write([]byte("foobar")).Return(6, nil)
+				var buf bytes.Buffer
+				str.EXPECT().Write(gomock.Any()).DoAndReturn(buf.Write).AnyTimes()
 
 				s.handleConn(conn)
 				Eventually(handlerCalled).Should(BeClosed())
+
+				// The buffer is expected to contain:
+				// 1. The response header (in a HEADERS frame)
+				// 2. the "foobar" (unframed)
+				frame, err := parseNextFrame(&buf, nil)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(frame).To(BeAssignableToTypeOf(&headersFrame{}))
+				df := frame.(*headersFrame)
+				data := make([]byte, df.Length)
+				_, err = io.ReadFull(&buf, data)
+				Expect(err).ToNot(HaveOccurred())
+				hdrs, err := qpackDecoder.DecodeFull(data)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(hdrs).To(ContainElement(qpack.HeaderField{Name: ":status", Value: "200"}))
+				Expect(buf.Bytes()).To(Equal([]byte("foobar")))
 			})
 
 			It("errors when the client sends a too large header frame", func() {
diff --git a/integrationtests/self/datagram_test.go b/integrationtests/self/datagram_test.go
index f9608272..1b4f0efd 100644
--- a/integrationtests/self/datagram_test.go
+++ b/integrationtests/self/datagram_test.go
@@ -64,7 +64,7 @@ var _ = Describe("Datagram test", func() {
 					maxDatagramMessageSize := f.MaxDataLen(maxDatagramSize, conn.ConnectionState().Version)
 					b := make([]byte, maxDatagramMessageSize+1)
 					Expect(conn.SendDatagram(b)).To(MatchError(&quic.DatagramTooLargeError{
-						PeerMaxDatagramFrameSize: int64(maxDatagramMessageSize),
+						MaxDatagramPayloadSize: int64(maxDatagramMessageSize),
 					}))
 					wg.Wait()
 				}
diff --git a/integrationtests/self/http_test.go b/integrationtests/self/http_test.go
index c8c018ff..5c08fb60 100644
--- a/integrationtests/self/http_test.go
+++ b/integrationtests/self/http_test.go
@@ -422,9 +422,8 @@ var _ = Describe("HTTP tests", func() {
 			defer GinkgoRecover()
 			close(handlerCalled)
 			w.WriteHeader(http.StatusOK)
-			w.(http.Flusher).Flush()
 
-			str := r.Body.(http3.HTTPStreamer).HTTPStream()
+			str := w.(http3.HTTPStreamer).HTTPStream()
 			str.Write([]byte("foobar"))
 
 			// Do this in a Go routine, so that the handler returns early.
@@ -734,9 +733,8 @@ var _ = Describe("HTTP tests", func() {
 				Eventually(conn.ReceivedSettings()).Should(BeClosed())
 				Expect(conn.Settings().EnableDatagrams).To(BeTrue())
 				w.WriteHeader(http.StatusOK)
-				w.(http.Flusher).Flush()
 
-				str := r.Body.(http3.HTTPStreamer).HTTPStream()
+				str := w.(http3.HTTPStreamer).HTTPStream()
 				go str.Read([]byte{0}) // need to continue reading from stream to observe state transitions
 
 				for {