feat: stream native request bodies

transport/native/server.go

@@ -410,12 +410,49 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
 			}
 		}
 		handler := s.handler
-		disconnectWatch := newClientDisconnectWatch(reader, state.conn, cancel)
+		var disconnectWatch *clientDisconnectWatch
+		if body, ok := request.Body.(*requestBody); ok {
+			// A body reader and the disconnect watcher must never read the
+			// bufio.Reader concurrently. The body starts the watcher once it
+			// reaches EOF; body read errors cancel the request directly.
+			body.cancel = cancel
+			body.onComplete = func() {
+				disconnectWatch = newClientDisconnectWatch(reader, state.conn, cancel)
+			}
+		} else {
+			disconnectWatch = newClientDisconnectWatch(reader, state.conn, cancel)
+		}
 		handler.ServeHTTP(writer, request)
-		disconnectWatch.stop()
+		bodyErr := request.Body.Close()
+		if disconnectWatch != nil {
+			disconnectWatch.stop()
+		}
 		cancel()
 
-		closing := s.isClosing() || state.closeWhenIdle.Load()
+		closing := s.isClosing() || state.closeWhenIdle.Load() || bodyErr != nil
+		if bodyErr != nil {
+			var protocolErr *protocolError
+			if !errors.As(bodyErr, &protocolErr) {
+				if requestSlot {
+					<-s.requestSlots
+				}
+				writer.reset("")
+				s.writerPool.Put(writer)
+				state.active.Store(false)
+				return
+			}
+			// Request bodies are streamed into the handler, so a malformed
+			// chunk or an over-limit body can be discovered while draining
+			// after the handler returns. Replace the buffered response with
+			// the same protocol error that eager parsing used to return.
+			if !writer.streaming {
+				writer.reset(request.Method)
+				writer.configure(request, state.conn, true, s.WriteTimeout, cancel)
+				writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
+				writer.WriteHeader(protocolErr.status)
+				writer.body.WriteString(protocolErr.message + "\n")
+			}
+		}
 		if err := writer.writeTo(state.conn, closeAfter || closing, s.WriteTimeout); err != nil {
 			if requestSlot {
 				<-s.requestSlots
@@ -577,27 +614,23 @@ func readRequest(reader *bufio.Reader, conn net.Conn, maxHeaderBytes int, maxBod
 	if !chunked && contentLength > maxBodyBytes {
 		return nil, true, &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
 	}
-	var body []byte
+	var body io.ReadCloser = http.NoBody
 	var trailer http.Header
 	if chunked {
-		body, trailer, err = readChunkedBody(reader, maxBodyBytes, setBodyDeadline, readLine)
-		if err != nil {
-			return nil, true, err
+		trailer = make(http.Header)
+		body = &requestBody{
+			reader: reader, maxBytes: maxBodyBytes, chunked: true,
+			trailer: trailer, setDeadline: setBodyDeadline, readLine: readLine,
 		}
 		// net/http uses -1 to indicate that the decoded length was not known
-		// from the request headers. The body is already buffered for the
-		// handler, but retaining this value keeps the request contract
-		// compatible with the standard server.
+		// from the request headers.
 		contentLength = -1
 	} else {
-		body = make([]byte, contentLength)
 		if contentLength > 0 {
 			if setBodyDeadline != nil {
 				setBodyDeadline()
 			}
-			if _, err := io.ReadFull(reader, body); err != nil {
-				return nil, true, err
-			}
+			body = &requestBody{reader: reader, remaining: contentLength, setDeadline: setBodyDeadline}
 		}
 	}
 	request := &http.Request{
@@ -607,7 +640,7 @@ func readRequest(reader *bufio.Reader, conn net.Conn, maxHeaderBytes int, maxBod
 		ProtoMajor:    protoMajor,
 		ProtoMinor:    protoMinor,
 		Header:        header,
-		Body:          io.NopCloser(bytes.NewReader(body)),
+		Body:          body,
 		ContentLength: contentLength,
 		Host:          header.Get("Host"),
 		RequestURI:    target,
@@ -627,78 +660,199 @@ func readRequest(reader *bufio.Reader, conn net.Conn, maxHeaderBytes int, maxBod
 	return request, closeAfter, nil
 }
 
-func readChunkedBody(reader *bufio.Reader, maxBodyBytes int64, setBodyDeadline func(), readLine func() (string, error)) ([]byte, http.Header, error) {
-	var body bytes.Buffer
-	initialCapacity := int64(64 << 10)
-	if maxBodyBytes < initialCapacity {
-		initialCapacity = maxBodyBytes
-	}
-	if initialCapacity > 0 && initialCapacity <= int64(int(^uint(0)>>1)) {
-		body.Grow(int(initialCapacity))
-	}
+type requestBody struct {
+	reader      *bufio.Reader
+	remaining   int64
+	maxBytes    int64
+	readBytes   int64
+	chunked     bool
+	chunkRemain int64
+	needSuffix  bool
+	trailer     http.Header
+	readLine    func() (string, error)
+	setDeadline func()
+	started     bool
+	cancel      context.CancelFunc
+	onComplete  func()
+	done        bool
+	closed      bool
+	err         error
+}
 
+func (b *requestBody) Read(p []byte) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+	if b.err != nil {
+		return 0, b.err
+	}
+	if b.done {
+		return 0, io.EOF
+	}
+	if !b.chunked {
+		return b.readContentLength(p)
+	}
+	return b.readChunked(p)
+}
+
+func (b *requestBody) readContentLength(p []byte) (int, error) {
+	if b.remaining == 0 {
+		b.complete()
+		return 0, io.EOF
+	}
+	if !b.started {
+		b.started = true
+		if b.setDeadline != nil {
+			b.setDeadline()
+		}
+	}
+	if int64(len(p)) > b.remaining {
+		p = p[:b.remaining]
+	}
+	n, err := b.reader.Read(p)
+	b.remaining -= int64(n)
+	b.readBytes += int64(n)
+	if b.remaining == 0 {
+		b.complete()
+		if n > 0 {
+			return n, nil
+		}
+		return 0, io.EOF
+	}
+	if err != nil {
+		b.fail(err)
+	}
+	return n, err
+}
+
+func (b *requestBody) readChunked(p []byte) (int, error) {
 	for {
-		line, err := readLine()
-		if err != nil {
-			return nil, nil, err
-		}
-		chunkSizeText := line
-		if extension := strings.IndexByte(chunkSizeText, ';'); extension >= 0 {
-			chunkSizeText = chunkSizeText[:extension]
-		}
-		chunkSizeText = strings.TrimSpace(chunkSizeText)
-		if chunkSizeText == "" || !validHexNumber(chunkSizeText) {
-			return nil, nil, &protocolError{status: http.StatusBadRequest, message: "invalid chunk size"}
-		}
-		chunkSize, parseErr := strconv.ParseUint(chunkSizeText, 16, 63)
-		if parseErr != nil {
-			return nil, nil, &protocolError{status: http.StatusBadRequest, message: "invalid chunk size"}
-		}
-		if chunkSize == 0 {
-			trailer := make(http.Header)
-			for {
-				line, err = readLine()
-				if err != nil {
-					return nil, nil, err
+		if b.chunkRemain == 0 {
+			if b.needSuffix {
+				var suffix [2]byte
+				if _, err := io.ReadFull(b.reader, suffix[:]); err != nil {
+					b.fail(err)
+					return 0, err
 				}
-				if line == "" {
-					return body.Bytes(), trailer, nil
+				if suffix[0] != '\r' || suffix[1] != '\n' {
+					err := &protocolError{status: http.StatusBadRequest, message: "malformed chunk terminator"}
+					b.fail(err)
+					return 0, err
 				}
-				if strings.HasPrefix(line, " ") || strings.HasPrefix(line, "\t") {
-					return nil, nil, &protocolError{status: http.StatusBadRequest, message: "folded trailers are not supported"}
-				}
-				index := strings.IndexByte(line, ':')
-				if index <= 0 {
-					return nil, nil, &protocolError{status: http.StatusBadRequest, message: "malformed trailer"}
-				}
-				name := strings.TrimSpace(line[:index])
-				value := strings.TrimSpace(line[index+1:])
-				if !validToken(name) || !validHeaderValue(value) {
-					return nil, nil, &protocolError{status: http.StatusBadRequest, message: "malformed trailer"}
-				}
-				trailer.Add(name, value)
+				b.needSuffix = false
+			}
+			if err := b.beginChunk(); err != nil {
+				b.fail(err)
+				return 0, err
+			}
+			if b.done {
+				return 0, io.EOF
 			}
 		}
-		if int64(body.Len()) > maxBodyBytes-int64(chunkSize) {
-			return nil, nil, &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
+		if int64(len(p)) > b.chunkRemain {
+			p = p[:b.chunkRemain]
 		}
-		if setBodyDeadline != nil {
-			setBodyDeadline()
+		n, err := b.reader.Read(p)
+		b.chunkRemain -= int64(n)
+		b.readBytes += int64(n)
+		if b.chunkRemain == 0 {
+			b.needSuffix = true
 		}
-		if _, err := io.CopyN(&body, reader, int64(chunkSize)); err != nil {
-			return nil, nil, err
+		if n > 0 {
+			if err != nil && b.chunkRemain > 0 {
+				b.fail(err)
+				return n, err
+			}
+			return n, nil
 		}
-		if setBodyDeadline != nil {
-			setBodyDeadline()
+		if err != nil {
+			b.fail(err)
+			return 0, err
 		}
-		var suffix [2]byte
-		if _, err := io.ReadFull(reader, suffix[:]); err != nil {
-			return nil, nil, err
-		}
-		if suffix[0] != '\r' || suffix[1] != '\n' {
-			return nil, nil, &protocolError{status: http.StatusBadRequest, message: "malformed chunk terminator"}
+	}
+}
+
+func (b *requestBody) beginChunk() error {
+	if b.setDeadline != nil {
+		b.setDeadline()
+	}
+	line, err := b.readLine()
+	if err != nil {
+		return err
+	}
+	chunkSizeText := line
+	if extension := strings.IndexByte(chunkSizeText, ';'); extension >= 0 {
+		chunkSizeText = chunkSizeText[:extension]
+	}
+	chunkSizeText = strings.TrimSpace(chunkSizeText)
+	if chunkSizeText == "" || !validHexNumber(chunkSizeText) {
+		return &protocolError{status: http.StatusBadRequest, message: "invalid chunk size"}
+	}
+	chunkSize, parseErr := strconv.ParseUint(chunkSizeText, 16, 63)
+	if parseErr != nil {
+		return &protocolError{status: http.StatusBadRequest, message: "invalid chunk size"}
+	}
+	if chunkSize == 0 {
+		for {
+			line, err = b.readLine()
+			if err != nil {
+				return err
+			}
+			if line == "" {
+				b.complete()
+				return nil
+			}
+			if strings.HasPrefix(line, " ") || strings.HasPrefix(line, "\t") {
+				return &protocolError{status: http.StatusBadRequest, message: "folded trailers are not supported"}
+			}
+			index := strings.IndexByte(line, ':')
+			if index <= 0 {
+				return &protocolError{status: http.StatusBadRequest, message: "malformed trailer"}
+			}
+			name := strings.TrimSpace(line[:index])
+			value := strings.TrimSpace(line[index+1:])
+			if !validToken(name) || !validHeaderValue(value) {
+				return &protocolError{status: http.StatusBadRequest, message: "malformed trailer"}
+			}
+			b.trailer.Add(name, value)
 		}
 	}
+	if int64(chunkSize) > b.maxBytes-b.readBytes {
+		return &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
+	}
+	b.chunkRemain = int64(chunkSize)
+	return nil
+}
+
+func (b *requestBody) Close() error {
+	if b.closed {
+		return b.err
+	}
+	b.closed = true
+	if !b.done && b.err == nil {
+		_, _ = io.Copy(io.Discard, b)
+	}
+	return b.err
+}
+
+func (b *requestBody) complete() {
+	if b.done {
+		return
+	}
+	b.done = true
+	if b.onComplete != nil {
+		b.onComplete()
+	}
+}
+
+func (b *requestBody) fail(err error) {
+	if b.err == nil {
+		b.err = err
+	}
+	b.done = true
+	if b.cancel != nil {
+		b.cancel()
+	}
 }
 
 func validHexNumber(value string) bool {

transport/native/server_test.go

@@ -107,6 +107,47 @@ func TestServerRejectsMalformedBodyAndLimits(t *testing.T) {
 	}
 }
 
+func TestServerStreamsContentLengthBodyToHandler(t *testing.T) {
+	entered := make(chan struct{})
+	application := omni.NewWithRouter(netrouter.New())
+	application.POST("/body", func(ctx *omni.Context) error {
+		close(entered)
+		body, err := io.ReadAll(ctx.Request().Body)
+		if err != nil {
+			return err
+		}
+		return ctx.OK(map[string]string{"body": string(body)})
+	})
+	server, listener, done := startServer(t, application)
+	defer stopServer(t, server, listener, done)
+
+	conn, err := net.Dial("tcp", listener.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	if _, err := io.WriteString(conn, "POST /body HTTP/1.1\r\nHost: test\r\nContent-Length: 5\r\n\r\nhel"); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-entered:
+	case <-time.After(time.Second):
+		t.Fatal("handler did not start before the complete body arrived")
+	}
+	if _, err := io.WriteString(conn, "lo"); err != nil {
+		t.Fatal(err)
+	}
+	response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodPost})
+	if err != nil {
+		t.Fatal(err)
+	}
+	body, err := io.ReadAll(response.Body)
+	response.Body.Close()
+	if err != nil || response.StatusCode != http.StatusOK || !strings.Contains(string(body), `"body":"hello"`) {
+		t.Fatalf("unexpected streamed body response: status=%d body=%s error=%v", response.StatusCode, body, err)
+	}
+}
+
 func TestServerChunkedBodyAndTrailers(t *testing.T) {
 	application := omni.NewWithRouter(netrouter.New())
 	application.POST("/body", func(ctx *omni.Context) error {
@@ -154,11 +195,18 @@ func TestServerChunkedBodyLimit(t *testing.T) {
 	}
 }
 
-func TestServerSlowBodyTimesOutBeforeHandler(t *testing.T) {
+func TestServerSlowBodyTimesOutWithoutSuccessfulHandler(t *testing.T) {
 	var handlerCalls atomic.Int32
 	application := omni.NewWithRouter(netrouter.New())
 	application.POST("/body", func(ctx *omni.Context) error {
+		body, err := io.ReadAll(ctx.Request().Body)
+		if err != nil {
+			return err
+		}
 		handlerCalls.Add(1)
+		if len(body) != 4 {
+			return fmt.Errorf("unexpected body length %d", len(body))
+		}
 		return ctx.OK("ok")
 	})
 	server := native.New(application)
@@ -178,14 +226,10 @@ func TestServerSlowBodyTimesOutBeforeHandler(t *testing.T) {
 		t.Fatal(err)
 	}
 	defer conn.Close()
-	conn.SetReadDeadline(time.Now().Add(time.Second))
 	_, _ = io.WriteString(conn, "POST /body HTTP/1.1\r\nHost: test\r\nContent-Length: 4\r\n\r\na")
-	buffer := make([]byte, 1)
-	if _, err := conn.Read(buffer); err == nil {
-		t.Fatal("expected slow request body to time out and close")
-	}
+	requirePeerClose(t, conn)
 	if handlerCalls.Load() != 0 {
-		t.Fatal("handler ran before the request body was complete")
+		t.Fatal("handler reported success before the request body was complete")
 	}
 }
 

transport/native/slow_client_test.go

@@ -43,8 +43,10 @@ func TestServerSlowHeadersCloseBeforeHandler(t *testing.T) {
 
 func TestServerSlowChunkedBodyClosesBeforeHandler(t *testing.T) {
 	var handlerCalls atomic.Int32
-	server := native.New(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
-		handlerCalls.Add(1)
+	server := native.New(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
+		if _, err := io.ReadAll(r.Body); err == nil {
+			handlerCalls.Add(1)
+		}
 	}))
 	server.ReadTimeout = 40 * time.Millisecond
 	listener, err := net.Listen("tcp", "127.0.0.1:0")