Skip to content

Commit

Permalink
Added option to limit response body size on debug log (#99)
Browse files Browse the repository at this point in the history
And improvements:
  - Refactor `if-else-if-else` chain to use switch(recommended by Effective Go).
  - Add cases `/json`, `/json-invalid`, `/long-text`, `/long-json`
  • Loading branch information
sudo-suhas authored and jeevatkm committed Oct 18, 2017
1 parent b081b82 commit 8b5e3f9
Show file tree
Hide file tree
Showing 6 changed files with 116 additions and 44 deletions.
33 changes: 21 additions & 12 deletions client.go
Original file line number Diff line number Diff line change
Expand Up @@ -85,18 +85,19 @@ type Client struct {
JSONMarshal func(v interface{}) ([]byte, error)
JSONUnmarshal func(data []byte, v interface{}) error

httpClient *http.Client
setContentLength bool
isHTTPMode bool
outputDirectory string
scheme string
proxyURL *url.URL
closeConnection bool
notParseResponse bool
beforeRequest []func(*Client, *Request) error
udBeforeRequest []func(*Client, *Request) error
preReqHook func(*Client, *Request) error
afterResponse []func(*Client, *Response) error
httpClient *http.Client
setContentLength bool
isHTTPMode bool
outputDirectory string
scheme string
proxyURL *url.URL
closeConnection bool
notParseResponse bool
debugBodySizeLimit int64
beforeRequest []func(*Client, *Request) error
udBeforeRequest []func(*Client, *Request) error
preReqHook func(*Client, *Request) error
afterResponse []func(*Client, *Response) error
}

// User type is to hold an username and password information
Expand Down Expand Up @@ -372,6 +373,14 @@ 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.
// resty.SetDebugBodyLimit(1000000)
//
func (c *Client) SetDebugBodyLimit(sl int64) *Client {
c.debugBodySizeLimit = sl
return c
}

// SetDisableWarn method disables the warning message on `go-resty` client.
// For example: go-resty warns the user when BasicAuth used on HTTP mode.
// resty.SetDisableWarn(true)
Expand Down
41 changes: 41 additions & 0 deletions client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
package resty

import (
"bytes"
"crypto/tls"
"errors"
"io/ioutil"
Expand Down Expand Up @@ -340,6 +341,10 @@ func TestClientOptions(t *testing.T) {
SetDebug(true)
assertEqual(t, DefaultClient.Debug, true)

var sl int64 = 1000000
SetDebugBodyLimit(sl)
assertEqual(t, DefaultClient.debugBodySizeLimit, sl)

SetAllowGetMethodPayload(true)
assertEqual(t, DefaultClient.AllowGetMethodPayload, true)

Expand Down Expand Up @@ -411,6 +416,42 @@ func TestNewRequest(t *testing.T) {
assertNotNil(t, request)
}

func TestDebugBodySizeLimit(t *testing.T) {
ts := createGetServer(t)
defer ts.Close()

var lgr bytes.Buffer
c := dc()
c.SetDebug(true)
c.SetLogger(&lgr)
c.SetDebugBodyLimit(30)

testcases := []struct{ url, want string }{
// Text, does not exceed limit.
{ts.URL, "TestGet: text response"},
// Empty response.
{ts.URL + "/no-content", "***** NO CONTENT *****"},
// JSON, does not exceed limit.
{ts.URL + "/json", "{\n \"TestGet\": \"JSON response\"\n}"},
// Invalid JSON, does not exceed limit.
{ts.URL + "/json-invalid", "TestGet: Invalid JSON"},
// Text, exceeds limit.
{ts.URL + "/long-text", "RESPONSE TOO LARGE"},
// JSON, exceeds limit.
{ts.URL + "/long-json", "RESPONSE TOO LARGE"},
}

for _, tc := range testcases {
_, err := c.R().Get(tc.url)
assertError(t, err)
debugLog := lgr.String()
if !strings.Contains(debugLog, tc.want) {
t.Errorf("Expected logs to contain [%v], got [\n%v]", tc.want, debugLog)
}
lgr.Reset()
}
}

// CustomRoundTripper just for test
type CustomRoundTripper struct {
}
Expand Down
37 changes: 22 additions & 15 deletions default.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"crypto/tls"
"encoding/json"
"io"
"math"
"net/http"
"net/http/cookiejar"
"net/url"
Expand All @@ -25,21 +26,22 @@ func New() *Client {
cookieJar, _ := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})

c := &Client{
HostURL: "",
QueryParam: url.Values{},
FormData: url.Values{},
Header: http.Header{},
UserInfo: nil,
Token: "",
Cookies: make([]*http.Cookie, 0),
Debug: false,
Log: getLogger(os.Stderr),
RetryCount: 0,
RetryWaitTime: defaultWaitTime,
RetryMaxWaitTime: defaultMaxWaitTime,
JSONMarshal: json.Marshal,
JSONUnmarshal: json.Unmarshal,
httpClient: &http.Client{Jar: cookieJar},
HostURL: "",
QueryParam: url.Values{},
FormData: url.Values{},
Header: http.Header{},
UserInfo: nil,
Token: "",
Cookies: make([]*http.Cookie, 0),
Debug: false,
Log: getLogger(os.Stderr),
RetryCount: 0,
RetryWaitTime: defaultWaitTime,
RetryMaxWaitTime: defaultMaxWaitTime,
JSONMarshal: json.Marshal,
JSONUnmarshal: json.Unmarshal,
httpClient: &http.Client{Jar: cookieJar},
debugBodySizeLimit: math.MaxInt32,
}

// Default transport
Expand Down Expand Up @@ -158,6 +160,11 @@ func SetDebug(d bool) *Client {
return DefaultClient.SetDebug(d)
}

// SetDebugBodyLimit method sets the response body limit for debug mode. See `Client.SetDebugBodyLimit` for more information.
func SetDebugBodyLimit(sl int64) *Client {
return DefaultClient.SetDebugBodyLimit(sl)
}

// SetAllowGetMethodPayload method allows the GET method with payload. See `Client.SetAllowGetMethodPayload` for more information.
func SetAllowGetMethodPayload(a bool) *Client {
return DefaultClient.SetAllowGetMethodPayload(a)
Expand Down
2 changes: 1 addition & 1 deletion middleware.go
Original file line number Diff line number Diff line change
Expand Up @@ -227,7 +227,7 @@ func responseLogger(c *Client, res *Response) error {
if res.Request.isSaveResponse {
c.Log.Printf("BODY :\n***** RESPONSE WRITTEN INTO FILE *****")
} else {
c.Log.Printf("BODY :\n%v", res.fmtBodyString())
c.Log.Printf("BODY :\n%v", res.fmtBodyString(c.debugBodySizeLimit))
}
c.Log.Println("----------------------------------------------------------")
c.enableLogPrefix()
Expand Down
14 changes: 8 additions & 6 deletions response.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ package resty

import (
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
Expand Down Expand Up @@ -119,20 +120,21 @@ func (r *Response) RawBody() io.ReadCloser {
return r.RawResponse.Body
}

func (r *Response) fmtBodyString() string {
bodyStr := "***** NO CONTENT *****"
func (r *Response) fmtBodyString(sl int64) string {
if r.body != nil {
if int64(len(r.body)) > sl {
return fmt.Sprintf("***** RESPONSE TOO LARGE (size - %d) *****", len(r.body))
}
ct := r.Header().Get(hdrContentTypeKey)
if IsJSONType(ct) {
out := acquireBuffer()
defer releaseBuffer(out)
if err := json.Indent(out, r.body, "", " "); err == nil {
bodyStr = string(out.Bytes())
return out.String()
}
} else {
bodyStr = r.String()
}
return r.String()
}

return bodyStr
return "***** NO CONTENT *****"
}
33 changes: 23 additions & 10 deletions resty_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1232,19 +1232,33 @@ func createGetServer(t *testing.T) *httptest.Server {
t.Logf("Path: %v", r.URL.Path)

if r.Method == MethodGet {
if r.URL.Path == "/" {
switch r.URL.Path {
case "/":
_, _ = w.Write([]byte("TestGet: text response"))
} else if r.URL.Path == "/mypage" {
case "/no-content":
_, _ = w.Write([]byte(""))
case "/json":
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"TestGet": "JSON response"}`))
case "/json-invalid":
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte("TestGet: Invalid JSON"))
case "/long-text":
_, _ = w.Write([]byte("TestGet: text response with size > 30"))
case "/long-json":
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"TestGet": "JSON response with size > 30"}`))
case "/mypage":
w.WriteHeader(http.StatusBadRequest)
} else if r.URL.Path == "/mypage2" {
case "/mypage2":
_, _ = w.Write([]byte("TestGet: text response from mypage2"))
} else if r.URL.Path == "/set-retrycount-test" {
case "/set-retrycount-test":
attp := atomic.AddInt32(&attempt, 1)
if attp <= 3 {
time.Sleep(time.Second * 6)
}
_, _ = w.Write([]byte("TestClientRetry page"))
} else if r.URL.Path == "/set-retrywaittime-test" {
case "/set-retrywaittime-test":
// Returns time.Duration since last request here
// or 0 for the very first request
if atomic.LoadInt32(&attempt) == 0 {
Expand All @@ -1257,20 +1271,19 @@ func createGetServer(t *testing.T) *httptest.Server {
_, _ = fmt.Fprintf(w, "%d", uint64(sinceLastRequest))
}
atomic.AddInt32(&attempt, 1)
} else if r.URL.Path == "/set-timeout-test-with-sequence" {
case "/set-timeout-test-with-sequence":
seq := atomic.AddInt32(&sequence, 1)
time.Sleep(time.Second * 2)
_, _ = fmt.Fprintf(w, "%d", seq)
} else if r.URL.Path == "/set-timeout-test" {
case "/set-timeout-test":
time.Sleep(time.Second * 6)
_, _ = w.Write([]byte("TestClientTimeout page"))

} else if r.URL.Path == "/my-image.png" {
case "/my-image.png":
fileBytes, _ := ioutil.ReadFile(getTestDataPath() + "/test-img.png")
w.Header().Set("Content-Type", "image/png")
w.Header().Set("Content-Length", strconv.Itoa(len(fileBytes)))
_, _ = w.Write(fileBytes)
} else if r.URL.Path == "/get-method-payload-test" {
case "/get-method-payload-test":
body, err := ioutil.ReadAll(r.Body)
if err != nil {
t.Errorf("Error: could not read get body: %s", err.Error())
Expand Down

0 comments on commit 8b5e3f9

Please sign in to comment.