From 67b0025f9cb203f5c030596e8faa1948966498df Mon Sep 17 00:00:00 2001 From: starainrt Date: Thu, 21 Aug 2025 19:17:19 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0content-length=E7=9A=84?= =?UTF-8?q?=E9=BB=98=E8=AE=A4=E5=A4=84=E7=90=86=E6=96=B9=E5=BC=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- curl.go | 68 ++++++++++++++++++++++++++++++++++++++++++---------- curl_test.go | 47 ++++++++++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 12 deletions(-) diff --git a/curl.go b/curl.go index 56c96a3..5b21038 100644 --- a/curl.go +++ b/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. -// also the memory usage will be higher -// Note that this function will not work if doRawRequest is true +// 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. +// So it may cause high memory usage if the request body is large. +// 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. -// also the memory usage will be higher -// Note that this function will not work if doRawRequest is true -func (r *RequestOpts) WithAutoCalcContentLength(autoCalcContentLength bool) RequestOpt { +// 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. +// So it may cause high memory usage if the request body is large. +// 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 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.Body = io.NopCloser(bytes.NewReader(data)) + req.ContentLength = int64(len(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{ diff --git a/curl_test.go b/curl_test.go index e34b79a..9c4d4d2 100644 --- a/curl_test.go +++ b/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)