From 3b17108465e70a78066674d2514b8c9e8a7ed570 Mon Sep 17 00:00:00 2001 From: Alex CHEN Date: Sun, 12 Jan 2020 21:17:19 +0800 Subject: [PATCH] Apply debugBodySizeLimit to Request debug logger --- client.go | 2 +- middleware.go | 2 +- request.go | 89 ++++++++++++++++++++++++++++--------------------- request_test.go | 79 +++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 132 insertions(+), 40 deletions(-) diff --git a/client.go b/client.go index 5a6e6331..bc60487c 100644 --- a/client.go +++ b/client.go @@ -366,7 +366,7 @@ func (c *Client) SetDebug(d bool) *Client { return c } -// SetDebugBodyLimit sets the maximum size for which the response body will be logged in debug mode. +// SetDebugBodyLimit sets the maximum size for which the response and request body will be logged in debug mode. // client.SetDebugBodyLimit(1000000) func (c *Client) SetDebugBodyLimit(sl int64) *Client { c.debugBodySizeLimit = sl diff --git a/middleware.go b/middleware.go index de3debd4..4d5c41f8 100644 --- a/middleware.go +++ b/middleware.go @@ -265,7 +265,7 @@ func addCredentials(c *Client, r *Request) error { func requestLogger(c *Client, r *Request) error { if c.Debug { rr := r.RawRequest - rl := &RequestLog{Header: copyHeaders(rr.Header), Body: r.fmtBodyString()} + rl := &RequestLog{Header: copyHeaders(rr.Header), Body: r.fmtBodyString(c.debugBodySizeLimit)} if c.requestLog != nil { if err := c.requestLog(rl); err != nil { return err diff --git a/request.go b/request.go index d67fb735..32481ea2 100644 --- a/request.go +++ b/request.go @@ -7,7 +7,6 @@ package resty import ( "bytes" "context" - "encoding/base64" "encoding/json" "encoding/xml" "fmt" @@ -662,51 +661,65 @@ type SRVRecord struct { // Request Unexported methods //_______________________________________________________________________ -func (r *Request) fmtBodyString() (body string) { +func (r *Request) fmtBodyString(sl int64) (body string) { body = "***** NO CONTENT *****" - if isPayloadSupported(r.Method, r.client.AllowGetMethodPayload) { - if _, ok := r.Body.(io.Reader); ok { - body = "***** BODY IS io.Reader *****" - return - } + if !isPayloadSupported(r.Method, r.client.AllowGetMethodPayload) { + return + } - // multipart or form-data - if r.isMultiPart || r.isFormData { - body = r.bodyBuf.String() - return - } + if _, ok := r.Body.(io.Reader); ok { + body = "***** BODY IS io.Reader *****" + return + } - // request body data - if r.Body == nil { + // multipart or form-data + if r.isMultiPart || r.isFormData { + bodySize := int64(r.bodyBuf.Len()) + if bodySize > sl { + body = fmt.Sprintf("***** REQUEST TOO LARGE (size - %d) *****", bodySize) return } - var prtBodyBytes []byte - var err error - - contentType := r.Header.Get(hdrContentTypeKey) - kind := kindOf(r.Body) - if canJSONMarshal(contentType, kind) { - prtBodyBytes, err = json.MarshalIndent(&r.Body, "", " ") - } else if IsXMLType(contentType) && (kind == reflect.Struct) { - prtBodyBytes, err = xml.MarshalIndent(&r.Body, "", " ") - } else if b, ok := r.Body.(string); ok { - if IsJSONType(contentType) { - bodyBytes := []byte(b) - out := acquireBuffer() - defer releaseBuffer(out) - if err = json.Indent(out, bodyBytes, "", " "); err == nil { - prtBodyBytes = out.Bytes() - } - } else { - body = b - return + body = r.bodyBuf.String() + return + } + + // request body data + if r.Body == nil { + return + } + var prtBodyBytes []byte + var err error + + contentType := r.Header.Get(hdrContentTypeKey) + kind := kindOf(r.Body) + if canJSONMarshal(contentType, kind) { + prtBodyBytes, err = json.MarshalIndent(&r.Body, "", " ") + } else if IsXMLType(contentType) && (kind == reflect.Struct) { + prtBodyBytes, err = xml.MarshalIndent(&r.Body, "", " ") + } else if b, ok := r.Body.(string); ok { + if IsJSONType(contentType) { + bodyBytes := []byte(b) + out := acquireBuffer() + defer releaseBuffer(out) + if err = json.Indent(out, bodyBytes, "", " "); err == nil { + prtBodyBytes = out.Bytes() } - } else if b, ok := r.Body.([]byte); ok { - body = base64.StdEncoding.EncodeToString(b) + } else { + body = b } + } else if b, ok := r.Body.([]byte); ok { + body = fmt.Sprintf("***** BODY IS byte(s) (size - %d) *****", len(b)) + return + } + + if prtBodyBytes != nil && err == nil { + body = string(prtBodyBytes) + } - if prtBodyBytes != nil && err == nil { - body = string(prtBodyBytes) + if len(body) > 0 { + bodySize := int64(len([]byte(body))) + if bodySize > sl { + body = fmt.Sprintf("***** REQUEST TOO LARGE (size - %d) *****", bodySize) } } diff --git a/request_test.go b/request_test.go index b74a18e3..2cb3c0fd 100644 --- a/request_test.go +++ b/request_test.go @@ -1576,3 +1576,82 @@ func TestTraceInfoWithoutEnableTrace(t *testing.T) { assertEqual(t, true, tr.TotalTime == 0) } } + +func TestDebugLoggerRequestBodyTooLarge(t *testing.T) { + ts := createFilePostServer(t) + defer ts.Close() + + debugBodySizeLimit := int64(512) + + // upload an image with more than 512 bytes + output := bytes.NewBufferString("") + resp, err := New().SetDebug(true).outputLogTo(output).SetDebugBodyLimit(debugBodySizeLimit).R(). + SetFile("file", filepath.Join(getTestDataPath(), "test-img.png")). + SetHeader("Content-Type", "image/png"). + Post(ts.URL + "/upload") + assertNil(t, err) + assertNotNil(t, resp) + assertEqual(t, true, strings.Contains(output.String(), "REQUEST TOO LARGE")) + + // upload a text file with no more than 512 bytes + output = bytes.NewBufferString("") + resp, err = New().SetDebug(true).outputLogTo(output).SetDebugBodyLimit(debugBodySizeLimit).R(). + SetFile("file", filepath.Join(getTestDataPath(), "text-file.txt")). + SetHeader("Content-Type", "text/plain"). + Post(ts.URL + "/upload") + assertNil(t, err) + assertNotNil(t, resp) + assertEqual(t, true, strings.Contains(output.String(), " THIS IS TEXT FILE FOR MULTIPART UPLOAD TEST ")) + + formTs := createFormPostServer(t) + defer formTs.Close() + + // post form with more than 512 bytes data + output = bytes.NewBufferString("") + resp, err = New().SetDebug(true).outputLogTo(output).SetDebugBodyLimit(debugBodySizeLimit).R(). + SetFormData(map[string]string{ + "first_name": "Alex", + "last_name": strings.Repeat("C", int(debugBodySizeLimit)), + "zip_code": "00001"}). + SetBasicAuth("myuser", "mypass"). + Post(formTs.URL + "/profile") + assertNil(t, err) + assertNotNil(t, resp) + assertEqual(t, true, strings.Contains(output.String(), "REQUEST TOO LARGE")) + + // post form with no more than 512 bytes data + output = bytes.NewBufferString("") + resp, err = New().SetDebug(true).outputLogTo(output).SetDebugBodyLimit(debugBodySizeLimit).R(). + SetFormData(map[string]string{ + "first_name": "Alex", + "last_name": "C", + "zip_code": "00001"}). + SetBasicAuth("myuser", "mypass"). + Post(formTs.URL + "/profile") + assertNil(t, err) + assertNotNil(t, resp) + assertEqual(t, true, strings.Contains(output.String(), "Alex")) + + // post string with more than 512 bytes data + output = bytes.NewBufferString("") + resp, err = New().SetDebug(true).outputLogTo(output).SetDebugBodyLimit(debugBodySizeLimit).R(). + SetBody(`{ + "first_name": "Alex", + "last_name": "` + strings.Repeat("C", int(debugBodySizeLimit)) + `C", + "zip_code": "00001"}`). + SetBasicAuth("myuser", "mypass"). + Post(formTs.URL + "/profile") + assertNil(t, err) + assertNotNil(t, resp) + assertEqual(t, true, strings.Contains(output.String(), "REQUEST TOO LARGE")) + + // post slice with more than 512 bytes data + output = bytes.NewBufferString("") + resp, err = New().SetDebug(true).outputLogTo(output).SetDebugBodyLimit(debugBodySizeLimit).R(). + SetBody([]string{strings.Repeat("C", int(debugBodySizeLimit))}). + SetBasicAuth("myuser", "mypass"). + Post(formTs.URL + "/profile") + assertNil(t, err) + assertNotNil(t, resp) + assertEqual(t, true, strings.Contains(output.String(), "REQUEST TOO LARGE")) +}