feat: support bounded chunked request bodies

2026-09-21 03:47:55 +08:00
parent b3cb9e8857
commit 22f1c9ebfb
6 changed files with 195 additions and 26 deletions
--- a/transport/native/server.go
+++ b/transport/native/server.go
@@ -2,8 +2,8 @@
 //
 // It deliberately keeps the public handler boundary as net/http so existing
 // GoOmni business code remains unchanged while connection management and
-// protocol I/O move out of net/http.Server. The parser currently supports
-// Content-Length request bodies and rejects chunked requests.
+// protocol I/O move out of net/http.Server. The parser supports Content-Length
+// and chunked request bodies with configurable size limits.
 package native

 import (
@@ -493,14 +493,23 @@ func readRequest(reader *bufio.Reader, conn net.Conn, maxHeaderBytes int, maxBod
 			return nil, true, &protocolError{status: http.StatusExpectationFailed, message: "unsupported expectation"}
 		}
 	}
+	chunked := false
 	for _, value := range header.Values("Transfer-Encoding") {
 		for _, coding := range strings.Split(value, ",") {
 			coding = strings.TrimSpace(coding)
-			if coding != "" && !strings.EqualFold(coding, "identity") {
-				return nil, true, &protocolError{status: http.StatusNotImplemented, message: "transfer encoding is not supported"}
+			if coding == "" || strings.EqualFold(coding, "identity") {
+				continue
 			}
+			if strings.EqualFold(coding, "chunked") && !chunked {
+				chunked = true
+				continue
+			}
+			return nil, true, &protocolError{status: http.StatusNotImplemented, message: "transfer encoding is not supported"}
 		}
 	}
+	if chunked && len(header.Values("Content-Length")) > 0 {
+		return nil, true, &protocolError{status: http.StatusBadRequest, message: "content length with chunked encoding"}
+	}
 	contentLength := int64(0)
 	contentLengthSeen := false
 	for _, value := range header.Values("Content-Length") {
@@ -514,17 +523,31 @@ func readRequest(reader *bufio.Reader, conn net.Conn, maxHeaderBytes int, maxBod
 			contentLengthSeen = true
 		}
 	}
-	if contentLength > maxBodyBytes {
+	if !chunked && contentLength > maxBodyBytes {
 		return nil, true, &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
 	}
-	body := make([]byte, contentLength)
-	if contentLength > 0 {
-		if setBodyDeadline != nil {
-			setBodyDeadline()
-		}
-		if _, err := io.ReadFull(reader, body); err != nil {
+	var body []byte
+	var trailer http.Header
+	if chunked {
+		body, trailer, err = readChunkedBody(reader, maxBodyBytes, setBodyDeadline, readLine)
+		if err != nil {
 			return nil, true, err
 		}
+		// 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.
+		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
+			}
+		}
 	}
 	request := &http.Request{
 		Method:        parts[0],
@@ -538,6 +561,10 @@ func readRequest(reader *bufio.Reader, conn net.Conn, maxHeaderBytes int, maxBod
 		Host:          header.Get("Host"),
 		RequestURI:    target,
 		RemoteAddr:    conn.RemoteAddr().String(),
+		Trailer:       trailer,
+	}
+	if chunked {
+		request.TransferEncoding = []string{"chunked"}
 	}
 	closeAfter := headerHasToken(header.Values("Connection"), "close") ||
 		(protoMajor == 1 && protoMinor == 0 && !headerHasToken(header.Values("Connection"), "keep-alive"))
@@ -545,6 +572,90 @@ 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))
+	}
+
+	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 line == "" {
+					return body.Bytes(), trailer, nil
+				}
+				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)
+			}
+		}
+		if int64(body.Len()) > maxBodyBytes-int64(chunkSize) {
+			return nil, nil, &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
+		}
+		if setBodyDeadline != nil {
+			setBodyDeadline()
+		}
+		if _, err := io.CopyN(&body, reader, int64(chunkSize)); err != nil {
+			return nil, nil, err
+		}
+		if setBodyDeadline != nil {
+			setBodyDeadline()
+		}
+		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 validHexNumber(value string) bool {
+	for index := 0; index < len(value); index++ {
+		char := value[index]
+		if !((char >= '0' && char <= '9') || (char >= 'a' && char <= 'f') || (char >= 'A' && char <= 'F')) {
+			return false
+		}
+	}
+	return true
+}
+
 func parseProtocol(value string) (int, int, bool) {
 	if value != "HTTP/1.1" && value != "HTTP/1.0" {
 		return 0, 0, false
--- a/transport/native/server_test.go
+++ b/transport/native/server_test.go
@@ -53,7 +53,7 @@ func TestServerKeepAliveAndApplicationBridge(t *testing.T) {
 	}
 }

-func TestServerRejectsUnsupportedBodyAndLimits(t *testing.T) {
+func TestServerRejectsMalformedBodyAndLimits(t *testing.T) {
 	application := omni.NewWithRouter(netrouter.New())
 	application.POST("/body", func(ctx *omni.Context) error { return ctx.OK("ok") })

@@ -67,28 +67,85 @@ func TestServerRejectsUnsupportedBodyAndLimits(t *testing.T) {
 	go func() { done <- server.Serve(context.Background(), listener) }()
 	defer stopServer(t, server, listener, done)

-	for _, requestText := range []string{
-		"POST /body HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\n\r\n4\r\ntest\r\n0\r\n\r\n",
-		"POST /body HTTP/1.1\r\nHost: test\r\nContent-Length: 5\r\n\r\ntest!",
-	} {
+	tests := []struct {
+		request string
+		status  int
+	}{
+		{
+			request: "POST /body HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\n\r\nZ\r\nbad\r\n0\r\n\r\n",
+			status:  http.StatusBadRequest,
+		},
+		{
+			request: "POST /body HTTP/1.1\r\nHost: test\r\nContent-Length: 5\r\n\r\ntest!",
+			status:  http.StatusRequestEntityTooLarge,
+		},
+	}
+	for _, test := range tests {
 		conn, err := net.Dial("tcp", listener.Addr().String())
 		if err != nil {
 			t.Fatal(err)
 		}
-		_, _ = io.WriteString(conn, requestText)
+		_, _ = io.WriteString(conn, test.request)
 		response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodPost})
 		if err != nil {
 			conn.Close()
 			t.Fatal(err)
 		}
 		response.Body.Close()
-		if response.StatusCode != http.StatusNotImplemented && response.StatusCode != http.StatusRequestEntityTooLarge {
-			t.Fatalf("unexpected protocol status %d", response.StatusCode)
+		if response.StatusCode != test.status {
+			t.Fatalf("expected status %d, got %d", test.status, response.StatusCode)
 		}
 		conn.Close()
 	}
 }

+func TestServerChunkedBodyAndTrailers(t *testing.T) {
+	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
+		}
+		return ctx.OK(map[string]string{
+			"body":     string(body),
+			"trailer":  ctx.Request().Trailer.Get("X-Trace"),
+			"transfer": strings.Join(ctx.Request().TransferEncoding, ","),
+		})
+	})
+	server, listener, done := startServer(t, application)
+	defer stopServer(t, server, listener, done)
+
+	request := "POST /body HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n" +
+		"4;source=wire\r\nWiki\r\n5\r\npedia\r\n0\r\nX-Trace: yes\r\n\r\n"
+	response := rawResponseMethod(t, listener.Addr().String(), http.MethodPost, request)
+	defer response.Body.Close()
+	body, err := io.ReadAll(response.Body)
+	if err != nil || response.StatusCode != http.StatusOK || !strings.Contains(string(body), `"body":"Wikipedia"`) || !strings.Contains(string(body), `"trailer":"yes"`) || !strings.Contains(string(body), `"transfer":"chunked"`) {
+		t.Fatalf("unexpected chunked response: status=%d body=%s error=%v", response.StatusCode, body, err)
+	}
+}
+
+func TestServerChunkedBodyLimit(t *testing.T) {
+	application := omni.NewWithRouter(netrouter.New())
+	application.POST("/body", func(ctx *omni.Context) error { return ctx.OK("unexpected") })
+	server := native.New(application)
+	server.MaxBodyBytes = 4
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	done := make(chan error, 1)
+	go func() { done <- server.Serve(context.Background(), listener) }()
+	defer stopServer(t, server, listener, done)
+
+	request := "POST /body HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\n\r\n5\r\nhello\r\n0\r\n\r\n"
+	response := rawResponseMethod(t, listener.Addr().String(), http.MethodPost, request)
+	defer response.Body.Close()
+	if response.StatusCode != http.StatusRequestEntityTooLarge {
+		t.Fatalf("expected chunked body limit status 413, got %d", response.StatusCode)
+	}
+}
+
 func TestServerSlowBodyTimesOutBeforeHandler(t *testing.T) {
 	var handlerCalls atomic.Int32
 	application := omni.NewWithRouter(netrouter.New())
@@ -210,9 +267,9 @@ func TestServerProtocolBoundaries(t *testing.T) {
 			status:  http.StatusBadRequest,
 		},
 		{
-			name:    "chunked",
-			request: "POST /health HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\n\r\n",
-			status:  http.StatusNotImplemented,
+			name:    "malformed chunked",
+			request: "POST /health HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\n\r\nZ\r\n\r\n",
+			status:  http.StatusBadRequest,
 		},
 		{
 			name:    "expectation",