更新content-length的默认处理方式

curl.go

@@ -12,7 +12,6 @@ import (
 	"net/http"
 	"net/url"
 	"os"
-	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -38,6 +37,7 @@ func (r *Request) Clone() *Request {
 			cookies:               CloneCookies(r.cookies),
 			bodyFormData:          CloneStringMapSlice(r.bodyFormData),
 			bodyFileData:          CloneFiles(r.bodyFileData),
+			contentLength:         r.contentLength,
 			queries:               CloneStringMapSlice(r.queries),
 			bodyDataBytes:         CloneByteSlice(r.bodyDataBytes),
 			customTransport:       r.customTransport,
@@ -274,6 +274,20 @@ type RequestOpts struct {
 	customDNS             []string
 	basicAuth             [2]string
 	autoCalcContentLength bool
+	contentLength         int64 // 人工设置
+}
+
+func (r *RequestOpts) ContentLength() int64 {
+	return r.contentLength
+}
+
+// SetContentLength sets the Content-Length header for the request.
+// This function will overwrite any existing or auto calculated Content-Length header.
+// if the length is less than 0, it will not set the Content-Length header. chunked transfer encoding will be used instead.
+// chunked transfer encoding may cause some servers to reject the request if they do not support it.
+// Note that this function will not work if doRawRequest is true
+func (r *RequestOpts) SetContentLength(contextLength int64) {
+	r.contentLength = contextLength
 }
 
 func (r *RequestOpts) CustomTransport() bool {
@@ -306,9 +320,10 @@ func (r *Request) AutoCalcContentLength() bool {
 }
 
 // SetAutoCalcContentLength sets whether to automatically calculate the Content-Length header based on the request body.
-// WARN: If set to true, the Content-Length header will be set to the length of the request body, which may cause issues with chunked transfer encoding.
+// WARN: If set to true, the Content-Length header will be set to the length of the request body, data will be cached in memory.
-// also the memory usage will be higher
+// So it may cause high memory usage if the request body is large.
-// Note that this function will not work if doRawRequest is true
+// If set to false, the Content-Length header will not be set,unless the request body is a byte slice or bytes.Buffer which has a specific length.
+// Note that this function will not work if doRawRequest is true or the ContentLength is already set.
 func (r *Request) SetAutoCalcContentLength(autoCalcContentLength bool) error {
 	if r.doRawRequest {
 		return fmt.Errorf("doRawRequest is true, cannot set autoCalcContentLength")
@@ -1342,12 +1357,25 @@ func WithAddCustomDNS(customDNS string) RequestOpt {
 }
 
 // WithAutoCalcContentLength sets whether to automatically calculate the Content-Length header based on the request body.
-// WARN: If set to true, the Content-Length header will be set to the length of the request body, which may cause issues with chunked transfer encoding.
+// WARN: If set to true, the Content-Length header will be set to the length of the request body, data will be cached in memory.
-// also the memory usage will be higher
+// So it may cause high memory usage if the request body is large.
-// Note that this function will not work if doRawRequest is true
+// If set to false, the Content-Length header will not be set,unless the request body is a byte slice or bytes.Buffer which has a specific length.
-func (r *RequestOpts) WithAutoCalcContentLength(autoCalcContentLength bool) RequestOpt {
+// Note that this function will not work if doRawRequest is true or ContentLength already set
+func WithAutoCalcContentLength(autoCalcContentLength bool) RequestOpt {
 	return func(opt *RequestOpts) error {
-		r.autoCalcContentLength = autoCalcContentLength
+		opt.autoCalcContentLength = autoCalcContentLength
+		return nil
+	}
+}
+
+// WithContentLength sets the Content-Length for the request.
+// This function will overwrite any existing or auto calculated Content-Length header.
+// if the length is less than 0, it will not set the Content-Length header. chunked transfer encoding will be used instead.
+// chunked transfer encoding may cause some servers to reject the request if they do not support it.
+// Note that this function will not work if doRawRequest is true
+func WithContentLength(length int64) RequestOpt {
+	return func(opt *RequestOpts) error {
+		opt.contentLength = length
 		return nil
 	}
 }
@@ -1572,11 +1600,21 @@ func newRequest(ctx context.Context, uri string, method string, opts ...RequestO
 
 func applyDataReader(r *Request) error {
 	// 优先度为：bodyDataReader > bodyDataBytes > bodyFormData > bodyFileData
+	r.rawRequest.ContentLength = 0
 	if r.bodyDataReader != nil {
+		switch v := r.bodyDataReader.(type) {
+		case *bytes.Buffer:
+			r.rawRequest.ContentLength = int64(v.Len())
+		case *bytes.Reader:
+			r.rawRequest.ContentLength = int64(v.Len())
+		case *strings.Reader:
+			r.rawRequest.ContentLength = int64(v.Len())
+		}
 		r.rawRequest.Body = io.NopCloser(r.bodyDataReader)
 		return nil
 	}
 	if len(r.bodyDataBytes) != 0 {
+		r.rawRequest.ContentLength = int64(len(r.bodyDataBytes))
 		r.rawRequest.Body = io.NopCloser(bytes.NewReader(r.bodyDataBytes))
 		return nil
 	}
@@ -1587,6 +1625,7 @@ func applyDataReader(r *Request) error {
 				body.Add(k, vv)
 			}
 		}
+		r.rawRequest.ContentLength = int64(len(body.Encode()))
 		r.rawRequest.Body = io.NopCloser(strings.NewReader(body.Encode()))
 		return nil
 	}
@@ -1596,7 +1635,6 @@ func applyDataReader(r *Request) error {
 		r.rawRequest.Header.Set("Content-Type", w.FormDataContentType())
 		go func() {
 			defer pw.Close() // ensure pipe writer is closed
-
 			if len(r.bodyFormData) != 0 {
 				for k, v := range r.bodyFormData {
 					for _, vv := range v {
@@ -1688,10 +1726,16 @@ func applyOptions(r *Request) error {
 				if err != nil {
 					return fmt.Errorf("read data error: %s", err)
 				}
-				req.Header.Set("Content-Length", strconv.Itoa(len(data)))
+				req.ContentLength = int64(len(data))
-				req.Body = io.NopCloser(bytes.NewReader(data))
+				req.Body = io.NopCloser(bytes.NewBuffer(data))
 			}
 		}
+		if r.contentLength > 0 {
+			req.ContentLength = r.contentLength
+		} else if r.contentLength < 0 {
+			//force use chunked transfer encoding
+			req.ContentLength = 0
+		}
 	}
 	if !r.alreadySetLookUpIPfn && len(r.customDNS) > 0 {
 		resolver := net.Resolver{

curl_test.go

@@ -603,6 +603,53 @@ func TestTlsConfig(t *testing.T) {
 	}
 }
 
+func TestHttpPostAndChunked(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		if req.Method != http.MethodPost {
+			t.Errorf("Expected 'POST', got %v", req.Method)
+		}
+		buf := make([]byte, 1024)
+		n, _ := req.Body.Read(buf)
+		if string(buf[:n]) != "hello world" {
+			t.Errorf("Expected body to be 'hello world', got %s", string(buf[:n]))
+		}
+
+		if req.Header.Get("chunked") == "true" {
+			if req.TransferEncoding[0] != "chunked" {
+				t.Errorf("Expected Transfer-Encoding to be 'chunked', got %s", req.Header.Get("Transfer-Encoding"))
+			}
+		} else {
+			if len(req.TransferEncoding) > 0 && req.TransferEncoding[0] == "chunked" {
+				t.Errorf("Expected Transfer-Encoding to not be 'chunked', got %s", req.Header.Get("Transfer-Encoding"))
+			}
+		}
+		rw.Write([]byte(`OK`))
+	}))
+	defer server.Close()
+
+	resp, err := Post(server.URL, WithBytes([]byte("hello world")), WithContentLength(-1), WithHeader("Content-Type", "text/plain"),
+		WithHeader("chunked", "true"))
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+	body := resp.Body().String()
+	if body != "OK" {
+		t.Errorf("Expected OK, got %v", body)
+	}
+	resp.Close()
+
+	resp, err = Post(server.URL, WithBytes([]byte("hello world")), WithHeader("Content-Type", "text/plain"),
+		WithHeader("chunked", "false"))
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+	defer resp.Close()
+	body = resp.Body().String()
+	if body != "OK" {
+		t.Errorf("Expected OK, got %v", body)
+	}
+}
+
 func TestWithTimeout(t *testing.T) {
 	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		time.Sleep(time.Second * 30)