Truncates HTTP responses in logging

logger.go

@@ -4,12 +4,14 @@
 package awsbase
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"fmt"
 	"io"
 	"log"
 	"net/http"
+	"net/textproto"
 	"strings"
 	"time"
 
@@ -146,16 +148,21 @@ func decomposeHTTPResponse(resp *http.Response, elapsed time.Duration) (map[stri
 	return result, nil
 }
 
-func decomposeResponseBody(resp *http.Response) (attribute.KeyValue, error) {
-	respBytes, err := io.ReadAll(resp.Body)
+func decomposeResponseBody(resp *http.Response) (kv attribute.KeyValue, err error) {
+	content, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return attribute.KeyValue{}, err
+		return kv, err
 	}
 
-	body := logging.MaskAWSAccessKey(string(respBytes))
-
 	// Restore the body reader
-	resp.Body = io.NopCloser(bytes.NewBuffer(respBytes))
+	resp.Body = io.NopCloser(bytes.NewBuffer(content))
+
+	reader := textproto.NewReader(bufio.NewReader(bytes.NewReader(content)))
+
+	body, err := logging.ReadTruncatedBody(reader, logging.MaxResponseBodyLen)
+	if err != nil {
+		return kv, err
+	}
 
 	return attribute.String("http.response.body", body), nil
 }

logging/http.go

@@ -23,7 +23,9 @@ import (
 )
 
 const (
-	maxRequestBodyLen = 512
+	maxRequestBodyLen = 1024
+
+	MaxResponseBodyLen = 4096
 )
 
 func DecomposeHTTPRequest(req *http.Request) (map[string]any, error) {
@@ -88,41 +90,27 @@ func decomposeRequestHeaders(req *http.Request) []attribute.KeyValue {
 	return results
 }
 
-func decomposeRequestBody(req *http.Request) (attribute.KeyValue, error) {
+func decomposeRequestBody(req *http.Request) (kv attribute.KeyValue, err error) {
 	reqBytes, err := httputil.DumpRequestOut(req, true)
 	if err != nil {
-		return attribute.KeyValue{}, err
+		return kv, err
 	}
 
 	reader := textproto.NewReader(bufio.NewReader(bytes.NewReader(reqBytes)))
 
 	if _, err = reader.ReadLine(); err != nil {
-		return attribute.KeyValue{}, err
+		return kv, err
 	}
 
 	if _, err = reader.ReadMIMEHeader(); err != nil {
-		return attribute.KeyValue{}, err
+		return kv, err
 	}
 
-	var builder strings.Builder
-	for {
-		line, err := reader.ReadContinuedLine()
-		if errors.Is(err, io.EOF) {
-			break
-		}
-		if err != nil {
-			return attribute.KeyValue{}, err
-		}
-		builder.WriteString(line)
-		if builder.Len() >= maxRequestBodyLen {
-			builder.WriteString("[truncated...]")
-			break
-		}
+	body, err := ReadTruncatedBody(reader, maxRequestBodyLen)
+	if err != nil {
+		return kv, err
 	}
 
-	body := builder.String()
-	body = MaskAWSAccessKey(body)
-
 	return attribute.String("http.request.body", body), nil
 }
 
@@ -222,3 +210,26 @@ func cleanUpHeaderAttributes(attrs []attribute.KeyValue) []attribute.KeyValue {
 		return attr
 	})
 }
+
+func ReadTruncatedBody(reader *textproto.Reader, len int) (string, error) {
+	var builder strings.Builder
+	for {
+		line, err := reader.ReadLine()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			return "", err
+		}
+		fmt.Fprintln(&builder, line)
+		if builder.Len() >= len {
+			fmt.Fprint(&builder, "[truncated...]")
+			break
+		}
+	}
+
+	body := builder.String()
+	body = MaskAWSAccessKey(body)
+
+	return body, nil
+}

v2/awsv1shim/logger.go

@@ -4,12 +4,14 @@
 package awsv1shim
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"fmt"
 	"io"
 	"log"
 	"net/http"
+	"net/textproto"
 	"strings"
 	"time"
 
@@ -21,6 +23,10 @@ import (
 	"go.opentelemetry.io/otel/semconv/v1.17.0/httpconv"
 )
 
+const (
+	responseBufferLen = logging.MaxResponseBodyLen + 1024
+)
+
 type debugLogger struct{}
 
 func (l debugLogger) Log(args ...interface{}) {
@@ -116,7 +122,7 @@ func logResponse(r *request.Request) {
 	bodyBuffer := bytes.NewBuffer(nil)
 
 	r.HTTPResponse.Body = &teeReaderCloser{
-		Reader: io.TeeReader(r.HTTPResponse.Body, bodyBuffer),
+		Reader: io.TeeReader(r.HTTPResponse.Body, limitWriter(bodyBuffer, responseBufferLen)),
 		Source: r.HTTPResponse.Body,
 	}
 
@@ -183,13 +189,44 @@ func decomposeHTTPResponse(resp *http.Response, body io.Reader, elapsed time.Dur
 	return result, nil
 }
 
-func decomposeResponseBody(bodyReader io.Reader) (attribute.KeyValue, error) {
-	respBytes, err := io.ReadAll(bodyReader)
+func decomposeResponseBody(bodyReader io.Reader) (kv attribute.KeyValue, err error) {
+	content, err := io.ReadAll(bodyReader)
 	if err != nil {
-		return attribute.KeyValue{}, err
+		return kv, err
 	}
 
-	body := logging.MaskAWSAccessKey(string(respBytes))
+	reader := textproto.NewReader(bufio.NewReader(bytes.NewReader(content)))
+
+	body, err := logging.ReadTruncatedBody(reader, logging.MaxResponseBodyLen)
+	if err != nil {
+		return kv, err
+	}
 
 	return attribute.String("http.response.body", body), nil
 }
+
+func limitWriter(w io.Writer, n int64) io.Writer {
+	return &limitedWriter{w, n}
+}
+
+type limitedWriter struct {
+	W io.Writer // the underlying writer
+	N int64     // max bytes remaining
+}
+
+// Write writes data into the wrapped Writer up to a limit of N bytes
+// Silently stops writing and returns full size of p to allow use with io.TeeReader
+func (w *limitedWriter) Write(p []byte) (int, error) {
+	if w.N <= 0 {
+		return len(p), nil
+	}
+	if int64(len(p)) > w.N {
+		n, err := w.W.Write(p[0:w.N])
+		w.N -= int64(n)
+		return len(p), err
+	} else {
+		n, err := w.W.Write(p)
+		w.N -= int64(n)
+		return n, err
+	}
+}