feat: redact secrets from other authentication headers when using debug option

Author: stainless-app[bot]

internal/debugmiddleware/debug_middleware.go

@@ -1,9 +1,11 @@
 package debugmiddleware
 
 import (
+	"log"
 	"net/http"
 	"net/http/httputil"
 	"strings"
+	"sync"
 )
 
 // For the time being these type definitions are duplicated here so that we can
@@ -15,14 +17,27 @@ type (
 
 const redactedPlaceholder = "<REDACTED>"
 
-// DebugMiddleware returns a middleware that logs HTTP requests and responses.
-//
-// logWriter is log.Default() under most circumstances, but made low level so we
-// can more easily inject a buffer to check in tests.
-func DebugMiddleware(logger interface{ Printf(string, ...any) }) Middleware {
+// Headers known to contain sensitive information like an API key.
+var sensitiveHeaders = []string{}
+
+// RequestLogger is a middleware that logs HTTP requests and responses.
+type RequestLogger struct {
+	logger           interface{ Printf(string, ...any) } // field for testability; usually log.Default()
+	sensitiveHeaders []string                            // field for testability; usually sensitiveHeaders
+}
+
+// NewRequestLogger returns a new RequestLogger instance with default options.
+func NewRequestLogger() *RequestLogger {
+	return &RequestLogger{
+		logger:           log.Default(),
+		sensitiveHeaders: sensitiveHeaders,
+	}
+}
+
+func (m *RequestLogger) Middleware() Middleware {
 	return func(req *http.Request, mn MiddlewareNext) (*http.Response, error) {
-		if reqBytes, err := httputil.DumpRequest(redactRequest(req), true); err == nil {
-			logger.Printf("Request Content:\n%s\n", reqBytes)
+		if reqBytes, err := httputil.DumpRequest(m.redactRequest(req), true); err == nil {
+			m.logger.Printf("Request Content:\n%s\n", reqBytes)
 		}
 
 		resp, err := mn(req)
@@ -31,7 +46,7 @@ func DebugMiddleware(logger interface{ Printf(string, ...any) }) Middleware {
 		}
 
 		if respBytes, err := httputil.DumpResponse(resp, true); err == nil {
-			logger.Printf("Response Content:\n%s\n", respBytes)
+			m.logger.Printf("Response Content:\n%s\n", respBytes)
 		}
 
 		return resp, err
@@ -42,17 +57,40 @@ func DebugMiddleware(logger interface{ Printf(string, ...any) }) Middleware {
 // purposes. If redaction is necessary, the request is cloned before mutating
 // the original and that clone is returned. As a small optimization, the
 // original is request is returned unchanged if no redaction is necessary.
-func redactRequest(req *http.Request) *http.Request {
-	if auth := req.Header.Get("Authorization"); auth != "" {
+func (m *RequestLogger) redactRequest(req *http.Request) *http.Request {
+	cloneReq := sync.OnceFunc(func() {
 		req = req.Clone(req.Context())
+	})
+
+	// Notably, the clauses below are written so they can redact multiple
+	// headers of the same name if necessary.
+	if values := req.Header.Values("Authorization"); len(values) > 0 {
+		cloneReq()
+		req.Header.Del("Authorization")
+
+		for _, value := range values {
+			// In case we're using something like a bearer token (e.g. `Bearer
+			// <my_token>`), keep the `Bearer` part for more debugging
+			// information.
+			if authKind, _, ok := strings.Cut(value, " "); ok {
+				req.Header.Add("Authorization", authKind+" "+redactedPlaceholder)
+			} else {
+				req.Header.Add("Authorization", redactedPlaceholder)
+			}
+		}
+	}
+
+	for _, header := range m.sensitiveHeaders {
+		values := req.Header.Values(header)
+		if len(values) == 0 {
+			continue
+		}
+
+		cloneReq()
+		req.Header.Del(header)
 
-		// In case we're using something like a bearer token (e.g. `Bearer
-		// <my_token>`), keep the `Bearer` part for more debugging
-		// information.
-		if authKind, _, ok := strings.Cut(auth, " "); ok {
-			req.Header.Set("Authorization", authKind+" "+redactedPlaceholder)
-		} else {
-			req.Header.Set("Authorization", redactedPlaceholder)
+		for range values {
+			req.Header.Add(header, redactedPlaceholder)
 		}
 	}
 

internal/debugmiddleware/debug_middleware_test.go

@@ -5,16 +5,21 @@ import (
 	"log"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
 	"strings"
 	"testing"
 )
 
 func TestDebugMiddleware(t *testing.T) {
 	t.Parallel()
 
-	setup := func() (Middleware, *bytes.Buffer) {
-		var logBuf bytes.Buffer
-		return DebugMiddleware(log.New(&logBuf, "", 0)), &logBuf
+	setup := func() (*RequestLogger, *bytes.Buffer) {
+		var (
+			logBuf     bytes.Buffer
+			middleware = NewRequestLogger()
+		)
+		middleware.logger = log.New(&logBuf, "", 0)
+		return middleware, &logBuf
 	}
 
 	t.Run("DoesNotRedactMostHeaders", func(t *testing.T) {
@@ -28,7 +33,7 @@ func TestDebugMiddleware(t *testing.T) {
 		req.Header.Set("User-Agent", stainlessUserAgent)
 
 		var nextMiddlewareRan bool
-		middleware(req, func(req *http.Request) (*http.Response, error) {
+		middleware.Middleware()(req, func(req *http.Request) (*http.Response, error) {
 			nextMiddlewareRan = true
 
 			// The request sent down through middleware shouldn't be mutated.
@@ -59,7 +64,7 @@ func TestDebugMiddleware(t *testing.T) {
 		req.Header.Set("Authorization", secretToken)
 
 		var nextMiddlewareRan bool
-		middleware(req, func(req *http.Request) (*http.Response, error) {
+		middleware.Middleware()(req, func(req *http.Request) (*http.Response, error) {
 			nextMiddlewareRan = true
 
 			// The request sent down through middleware shouldn't be mutated.
@@ -88,7 +93,7 @@ func TestDebugMiddleware(t *testing.T) {
 		req.Header.Set("Authorization", "Bearer "+secretToken)
 
 		var nextMiddlewareRan bool
-		middleware(req, func(req *http.Request) (*http.Response, error) {
+		middleware.Middleware()(req, func(req *http.Request) (*http.Response, error) {
 			nextMiddlewareRan = true
 
 			return &http.Response{}, nil
@@ -102,4 +107,99 @@ func TestDebugMiddleware(t *testing.T) {
 			t.Error("expected authorization header to be redacted")
 		}
 	})
+
+	t.Run("RedactsMultipleAuthorizationHeaders", func(t *testing.T) {
+		t.Parallel()
+
+		middleware, logBuf := setup()
+
+		req := httptest.NewRequest("GET", "https://example.com", nil)
+		req.Header.Add("Authorization", secretToken+"1")
+		req.Header.Add("Authorization", secretToken+"2")
+
+		var nextMiddlewareRan bool
+		middleware.Middleware()(req, func(req *http.Request) (*http.Response, error) {
+			nextMiddlewareRan = true
+
+			// The request sent down through middleware shouldn't be mutated.
+			if !reflect.DeepEqual(req.Header.Values("Authorization"), []string{secretToken + "1", secretToken + "2"}) {
+				t.Errorf("expected original request to be unmodified")
+			}
+
+			return &http.Response{}, nil
+		})
+
+		if !nextMiddlewareRan {
+			t.Error("expected next middleware to have been run")
+		}
+
+		if strings.Count(logBuf.String(), "Authorization: "+redactedPlaceholder) != 2 {
+			t.Error("expected exactly two redacted placeholders in authorization headers")
+		}
+	})
+
+	const customAPIKeyHeader = "X-My-Api-Key"
+
+	t.Run("RedactsSensitiveHeaders", func(t *testing.T) {
+		t.Parallel()
+
+		middleware, logBuf := setup()
+
+		middleware.sensitiveHeaders = []string{customAPIKeyHeader}
+
+		req := httptest.NewRequest("GET", "https://example.com", nil)
+		req.Header.Set(customAPIKeyHeader, secretToken)
+
+		var nextMiddlewareRan bool
+		middleware.Middleware()(req, func(req *http.Request) (*http.Response, error) {
+			nextMiddlewareRan = true
+
+			// The request sent down through middleware shouldn't be mutated.
+			if req.Header.Get(customAPIKeyHeader) != secretToken {
+				t.Error("expected original request to be unmodified")
+			}
+
+			return &http.Response{}, nil
+		})
+
+		if !nextMiddlewareRan {
+			t.Error("expected next middleware to have been run")
+		}
+
+		if !strings.Contains(logBuf.String(), customAPIKeyHeader+": "+redactedPlaceholder) {
+			t.Errorf("expected %s header to be redacted", customAPIKeyHeader)
+		}
+	})
+
+	t.Run("RedactsMultipleSensitiveHeaders", func(t *testing.T) {
+		t.Parallel()
+
+		middleware, logBuf := setup()
+
+		middleware.sensitiveHeaders = []string{customAPIKeyHeader}
+
+		req := httptest.NewRequest("GET", "https://example.com", nil)
+		req.Header.Add(customAPIKeyHeader, secretToken+"1")
+		req.Header.Add(customAPIKeyHeader, secretToken+"2")
+
+		var nextMiddlewareRan bool
+		middleware.Middleware()(req, func(req *http.Request) (*http.Response, error) {
+			nextMiddlewareRan = true
+
+			// The request sent down through middleware shouldn't be mutated.
+			if !reflect.DeepEqual(req.Header.Values(customAPIKeyHeader), []string{secretToken + "1", secretToken + "2"}) {
+				t.Error("expected original request to be unmodified")
+			}
+
+			return &http.Response{}, nil
+		})
+
+		if !nextMiddlewareRan {
+			t.Error("expected next middleware to have been run")
+		}
+
+		if strings.Count(logBuf.String(), customAPIKeyHeader+": "+redactedPlaceholder) != 2 {
+			t.Errorf("expected %s header to be redacted", customAPIKeyHeader)
+		}
+	})
 }

pkg/cmd/flagoptions.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"encoding/json"
 	"io"
-	"log"
 	"mime/multipart"
 	"os"
 
@@ -32,7 +31,7 @@ func flagOptions(
 ) ([]option.RequestOption, error) {
 	var options []option.RequestOption
 	if cmd.Bool("debug") {
-		options = append(options, option.WithMiddleware(debugmiddleware.DebugMiddleware(log.Default())))
+		options = append(options, option.WithMiddleware(debugmiddleware.NewRequestLogger().Middleware()))
 	}
 
 	queries := make(map[string]any)