From 512d1e6da0680702df96bbc738650c7cf13045df Mon Sep 17 00:00:00 2001 From: "Lixia (Sylvia) Lei" Date: Thu, 5 Dec 2024 14:37:44 +0800 Subject: [PATCH] refactor per comments Signed-off-by: Lixia (Sylvia) Lei --- internal/trace/transport.go | 30 +++++++++++++----------------- internal/trace/transport_test.go | 8 ++++---- 2 files changed, 17 insertions(+), 21 deletions(-) diff --git a/internal/trace/transport.go b/internal/trace/transport.go index 3fbabc678..2ab827bfc 100644 --- a/internal/trace/transport.go +++ b/internal/trace/transport.go @@ -97,7 +97,7 @@ func logHeader(header http.Header) string { // logResponseBody prints out the response body if it is printable and within // the size limit. func logResponseBody(resp *http.Response) string { - if resp.Body == nil || resp.Body == http.NoBody || resp.ContentLength == 0 { + if resp.Body == nil || resp.Body == http.NoBody { return " No response body to print" } @@ -106,27 +106,23 @@ func logResponseBody(resp *http.Response) string { if !isPrintableContentType(contentType) { return fmt.Sprintf(" Response body of content type \"%s\" is not printed", contentType) } - if resp.ContentLength < 0 { - return " Response body of unknown content length is not printed" - } - if resp.ContentLength > payloadSizeLimit { - return fmt.Sprintf(" Response body larger than %d bytes is not printed", payloadSizeLimit) - } - // Note: If the actual body size mismatches the content length and exceeds the limit, - // the body will be truncated to the limit for seucrity consideration. - // In this case, the response processing subsequent to logging might be broken. - rc := resp.Body - defer rc.Close() - lr := io.LimitReader(rc, payloadSizeLimit) - bodyBytes, err := io.ReadAll(lr) + // read the body up to limit+1 to check if the body exceeds the limit + lr := io.LimitReader(resp.Body, payloadSizeLimit+1) + readBody, err := io.ReadAll(lr) if err != nil { return fmt.Sprintf(" Error reading response body: %v", err) } + // restore the body by concatenating the read body with the remaining body + resp.Body = io.NopCloser(io.MultiReader(bytes.NewReader(readBody), resp.Body)) - // reset the body for subsequent processing - resp.Body = io.NopCloser(bytes.NewReader(bodyBytes)) - return string(bodyBytes) + if len(readBody) == 0 { + return " Response body is empty" + } + if len(readBody) > int(payloadSizeLimit) { + return string(readBody[:payloadSizeLimit]) + "\n...(truncated)" + } + return string(readBody) } // isPrintableContentType returns true if the content of contentType is printable. diff --git a/internal/trace/transport_test.go b/internal/trace/transport_test.go index 36b4b232f..42c9a3db6 100644 --- a/internal/trace/transport_test.go +++ b/internal/trace/transport_test.go @@ -145,7 +145,7 @@ func Test_logResponseBody(t *testing.T) { ContentLength: 0, Header: http.Header{"Content-Type": []string{"text/plain"}}, }, - want: " No response body to print", + want: " Response body is empty", }, { name: "Unknown content length", @@ -154,7 +154,7 @@ func Test_logResponseBody(t *testing.T) { ContentLength: -1, Header: http.Header{"Content-Type": []string{"text/plain"}}, }, - want: " Response body of unknown content length is not printed", + want: "whatever", }, { name: "Non-printable content type", @@ -181,7 +181,7 @@ func Test_logResponseBody(t *testing.T) { ContentLength: payloadSizeLimit + 1, Header: http.Header{"Content-Type": []string{"text/plain"}}, }, - want: fmt.Sprintf(" Response body larger than %d bytes is not printed", payloadSizeLimit), + want: string(bytes.Repeat([]byte("a"), int(payloadSizeLimit))) + "\n...(truncated)", }, { name: "Printable content type within limit", @@ -208,7 +208,7 @@ func Test_logResponseBody(t *testing.T) { ContentLength: 1, // mismatched content length Header: http.Header{"Content-Type": []string{"text/plain"}}, }, - want: string(bytes.Repeat([]byte("a"), int(payloadSizeLimit))), + want: string(bytes.Repeat([]byte("a"), int(payloadSizeLimit))) + "\n...(truncated)", }, { name: "Actual body size is smaller than content length",