From 22a09c88be4e96dccbf357764bb71f0aa29b39df Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Wed, 8 Apr 2026 22:31:29 +0000 Subject: [PATCH 1/3] feat(coderd/x/chatd/chatdebug): add recorder, transport, and redaction Change-Id: Ibbc67a85ba78201c0778ccb5c8675b15e90b1cdf Signed-off-by: Thomas Kosiewski --- coderd/x/chatd/chatdebug/recorder.go | 277 +++++++ coderd/x/chatd/chatdebug/recorder_test.go | 174 +++++ coderd/x/chatd/chatdebug/redaction.go | 227 ++++++ coderd/x/chatd/chatdebug/redaction_test.go | 277 +++++++ coderd/x/chatd/chatdebug/reuse_step_test.go | 74 ++ coderd/x/chatd/chatdebug/stubs.go | 205 ++--- .../x/chatd/chatdebug/stubs_internal_test.go | 32 - coderd/x/chatd/chatdebug/transport.go | 382 +++++++++ coderd/x/chatd/chatdebug/transport_test.go | 737 ++++++++++++++++++ 9 files changed, 2208 insertions(+), 177 deletions(-) create mode 100644 coderd/x/chatd/chatdebug/recorder.go create mode 100644 coderd/x/chatd/chatdebug/recorder_test.go create mode 100644 coderd/x/chatd/chatdebug/redaction.go create mode 100644 coderd/x/chatd/chatdebug/redaction_test.go create mode 100644 coderd/x/chatd/chatdebug/reuse_step_test.go create mode 100644 coderd/x/chatd/chatdebug/transport.go create mode 100644 coderd/x/chatd/chatdebug/transport_test.go diff --git a/coderd/x/chatd/chatdebug/recorder.go b/coderd/x/chatd/chatdebug/recorder.go new file mode 100644 index 0000000000000..c3cc4d2c2aafb --- /dev/null +++ b/coderd/x/chatd/chatdebug/recorder.go @@ -0,0 +1,277 @@ +package chatdebug + +import ( + "context" + "sync" + "sync/atomic" + "time" + + "charm.land/fantasy" + "github.com/google/uuid" + + "cdr.dev/slog/v3" +) + +// RecorderOptions identifies the chat/model context for debug recording. +type RecorderOptions struct { + ChatID uuid.UUID + OwnerID uuid.UUID + Provider string + Model string +} + +// WrapModel returns model unchanged when debug recording is disabled, or a +// debug wrapper when a service is available. +func WrapModel( + model fantasy.LanguageModel, + svc *Service, + opts RecorderOptions, +) fantasy.LanguageModel { + if model == nil { + panic("chatdebug: nil LanguageModel") + } + if svc == nil { + return model + } + return &debugModel{inner: model, svc: svc, opts: opts} +} + +type attemptSink struct { + mu sync.Mutex + attempts []Attempt + attemptCounter atomic.Int32 +} + +func (s *attemptSink) nextAttemptNumber() int { + if s == nil { + panic("chatdebug: nil attemptSink") + } + return int(s.attemptCounter.Add(1)) +} + +func (s *attemptSink) record(a Attempt) { + s.mu.Lock() + defer s.mu.Unlock() + + s.attempts = append(s.attempts, a) +} + +func (s *attemptSink) snapshot() []Attempt { + s.mu.Lock() + defer s.mu.Unlock() + + attempts := make([]Attempt, len(s.attempts)) + copy(attempts, s.attempts) + return attempts +} + +type attemptSinkKey struct{} + +func withAttemptSink(ctx context.Context, sink *attemptSink) context.Context { + if sink == nil { + panic("chatdebug: nil attemptSink") + } + return context.WithValue(ctx, attemptSinkKey{}, sink) +} + +func attemptSinkFromContext(ctx context.Context) *attemptSink { + sink, _ := ctx.Value(attemptSinkKey{}).(*attemptSink) + return sink +} + +var stepCounters sync.Map // map[uuid.UUID]*atomic.Int32 + +func nextStepNumber(runID uuid.UUID) int32 { + val, _ := stepCounters.LoadOrStore(runID, &atomic.Int32{}) + counter, ok := val.(*atomic.Int32) + if !ok { + panic("chatdebug: invalid step counter type") + } + return counter.Add(1) +} + +// CleanupStepCounter removes per-run step counter and reference count +// state. This is used by tests and later stacked branches that have a +// real run lifecycle. +func CleanupStepCounter(runID uuid.UUID) { + stepCounters.Delete(runID) + runRefCounts.Delete(runID) +} + +const stepFinalizeTimeout = 5 * time.Second + +func stepFinalizeContext(ctx context.Context) (context.Context, context.CancelFunc) { + if ctx == nil { + panic("chatdebug: nil context") + } + return context.WithTimeout(context.WithoutCancel(ctx), stepFinalizeTimeout) +} + +func syncStepCounter(runID uuid.UUID, stepNumber int32) { + val, _ := stepCounters.LoadOrStore(runID, &atomic.Int32{}) + counter, ok := val.(*atomic.Int32) + if !ok { + panic("chatdebug: invalid step counter type") + } + for { + current := counter.Load() + if current >= stepNumber { + return + } + if counter.CompareAndSwap(current, stepNumber) { + return + } + } +} + +type stepHandle struct { + stepCtx *StepContext + sink *attemptSink + svc *Service + opts RecorderOptions + once sync.Once + mu sync.Mutex + status Status + response any + usage any + err any + metadata any +} + +// beginStep validates preconditions, creates a debug step, and returns a +// handle plus an enriched context carrying StepContext and attemptSink. +// Returns (nil, original ctx) when debug recording should be skipped. +func beginStep( + ctx context.Context, + svc *Service, + opts RecorderOptions, + op Operation, + normalizedReq any, +) (*stepHandle, context.Context) { + if svc == nil { + return nil, ctx + } + + rc, ok := RunFromContext(ctx) + if !ok || rc.RunID == uuid.Nil { + return nil, ctx + } + + chatID := opts.ChatID + if chatID == uuid.Nil { + chatID = rc.ChatID + } + if !svc.IsEnabled(ctx, chatID, opts.OwnerID) { + return nil, ctx + } + + holder, reuseStep := reuseHolderFromContext(ctx) + if reuseStep { + holder.mu.Lock() + defer holder.mu.Unlock() + // Only reuse the cached handle if it belongs to the same run. + // A different RunContext means a new logical run, so we must + // create a fresh step to avoid cross-run attribution. + if holder.handle != nil && holder.handle.stepCtx.RunID == rc.RunID { + enriched := ContextWithStep(ctx, holder.handle.stepCtx) + enriched = withAttemptSink(enriched, holder.handle.sink) + return holder.handle, enriched + } + } + + stepNum := nextStepNumber(rc.RunID) + step, err := svc.CreateStep(ctx, CreateStepParams{ + RunID: rc.RunID, + ChatID: chatID, + StepNumber: stepNum, + Operation: op, + Status: StatusInProgress, + HistoryTipMessageID: rc.HistoryTipMessageID, + NormalizedRequest: normalizedReq, + }) + if err != nil { + svc.log.Warn(ctx, "failed to create chat debug step", + slog.Error(err), + slog.F("chat_id", chatID), + slog.F("run_id", rc.RunID), + slog.F("operation", op), + ) + return nil, ctx + } + + syncStepCounter(rc.RunID, step.StepNumber) + actualStepNumber := step.StepNumber + if actualStepNumber == 0 { + actualStepNumber = stepNum + } + + sc := &StepContext{ + StepID: step.ID, + RunID: rc.RunID, + ChatID: chatID, + StepNumber: actualStepNumber, + Operation: op, + HistoryTipMessageID: rc.HistoryTipMessageID, + } + handle := &stepHandle{stepCtx: sc, sink: &attemptSink{}, svc: svc, opts: opts} + enriched := ContextWithStep(ctx, handle.stepCtx) + enriched = withAttemptSink(enriched, handle.sink) + if reuseStep { + holder.handle = handle + } + + return handle, enriched +} + +// finish updates the debug step with final status and data. +// sync.Once prevents data races when concurrent callers (e.g. +// retried stream wrappers sharing a reuse handle) both attempt +// to finalize the same step. Only the first finish call takes +// effect. +func (h *stepHandle) finish( + ctx context.Context, + status Status, + response any, + usage any, + errPayload any, + metadata any, +) { + if h == nil || h.stepCtx == nil { + return + } + + h.once.Do(func() { + h.mu.Lock() + h.status = status + h.response = response + h.usage = usage + h.err = errPayload + h.metadata = metadata + h.mu.Unlock() + if h.svc == nil { + return + } + + updateCtx, cancel := stepFinalizeContext(ctx) + defer cancel() + + if _, updateErr := h.svc.UpdateStep(updateCtx, UpdateStepParams{ + ID: h.stepCtx.StepID, + ChatID: h.stepCtx.ChatID, + Status: status, + NormalizedResponse: response, + Usage: usage, + Attempts: h.sink.snapshot(), + Error: errPayload, + Metadata: metadata, + FinishedAt: time.Now(), + }); updateErr != nil { + h.svc.log.Warn(updateCtx, "failed to finalize chat debug step", + slog.Error(updateErr), + slog.F("step_id", h.stepCtx.StepID), + slog.F("chat_id", h.stepCtx.ChatID), + slog.F("status", status), + ) + } + }) +} diff --git a/coderd/x/chatd/chatdebug/recorder_test.go b/coderd/x/chatd/chatdebug/recorder_test.go new file mode 100644 index 0000000000000..a2c3c88846a3d --- /dev/null +++ b/coderd/x/chatd/chatdebug/recorder_test.go @@ -0,0 +1,174 @@ +package chatdebug //nolint:testpackage // Uses unexported recorder helpers. + +import ( + "context" + "sort" + "sync" + "testing" + + "charm.land/fantasy" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/x/chatd/chattest" +) + +func TestAttemptSink_ThreadSafe(t *testing.T) { + t.Parallel() + + const n = 256 + + sink := &attemptSink{} + var wg sync.WaitGroup + wg.Add(n) + + for i := range n { + go func() { + defer wg.Done() + sink.record(Attempt{Number: i + 1, ResponseStatus: 200 + i}) + }() + } + + wg.Wait() + + attempts := sink.snapshot() + require.Len(t, attempts, n) + + numbers := make([]int, 0, n) + statuses := make([]int, 0, n) + for _, attempt := range attempts { + numbers = append(numbers, attempt.Number) + statuses = append(statuses, attempt.ResponseStatus) + } + sort.Ints(numbers) + sort.Ints(statuses) + + for i := range n { + require.Equal(t, i+1, numbers[i]) + require.Equal(t, 200+i, statuses[i]) + } +} + +func TestAttemptSinkContext(t *testing.T) { + t.Parallel() + + ctx := context.Background() + require.Nil(t, attemptSinkFromContext(ctx)) + + sink := &attemptSink{} + ctx = withAttemptSink(ctx, sink) + require.Same(t, sink, attemptSinkFromContext(ctx)) +} + +func TestWrapModel_NilModel(t *testing.T) { + t.Parallel() + + require.Panics(t, func() { + WrapModel(nil, &Service{}, RecorderOptions{}) + }) +} + +func TestWrapModel_NilService(t *testing.T) { + t.Parallel() + + model := &chattest.FakeModel{ProviderName: "provider", ModelName: "model"} + wrapped := WrapModel(model, nil, RecorderOptions{}) + require.Same(t, model, wrapped) +} + +func TestNextStepNumber_Concurrent(t *testing.T) { + t.Parallel() + + const n = 256 + + runID := uuid.New() + results := make([]int, n) + var wg sync.WaitGroup + wg.Add(n) + + for i := range n { + go func() { + defer wg.Done() + results[i] = int(nextStepNumber(runID)) + }() + } + + wg.Wait() + + sort.Ints(results) + for i := range n { + require.Equal(t, i+1, results[i]) + } +} + +func TestStepFinalizeContext_StripsCancellation(t *testing.T) { + t.Parallel() + + baseCtx, cancelBase := context.WithCancel(context.Background()) + cancelBase() + require.ErrorIs(t, baseCtx.Err(), context.Canceled) + + finalizeCtx, cancelFinalize := stepFinalizeContext(baseCtx) + defer cancelFinalize() + + require.NoError(t, finalizeCtx.Err()) + _, hasDeadline := finalizeCtx.Deadline() + require.True(t, hasDeadline) +} + +func TestSyncStepCounter_AdvancesCounter(t *testing.T) { + t.Parallel() + + runID := uuid.New() + t.Cleanup(func() { CleanupStepCounter(runID) }) + + syncStepCounter(runID, 7) + require.Equal(t, int32(8), nextStepNumber(runID)) +} + +func TestStepHandleFinish_NilHandle(t *testing.T) { + t.Parallel() + + var handle *stepHandle + handle.finish(context.Background(), StatusCompleted, nil, nil, nil, nil) +} + +func TestBeginStep_NilService(t *testing.T) { + t.Parallel() + + ctx := context.Background() + handle, enriched := beginStep(ctx, nil, RecorderOptions{}, OperationGenerate, nil) + require.Nil(t, handle) + require.Nil(t, attemptSinkFromContext(enriched)) + _, ok := StepFromContext(enriched) + require.False(t, ok) +} + +func TestBeginStep_FallsBackToRunChatID(t *testing.T) { + t.Parallel() + + runID := uuid.New() + runChatID := uuid.New() + ctx := ContextWithRun(context.Background(), &RunContext{RunID: runID, ChatID: runChatID}) + + handle, enriched := beginStep(ctx, &Service{}, RecorderOptions{}, OperationGenerate, nil) + require.NotNil(t, handle) + require.Equal(t, runChatID, handle.stepCtx.ChatID) + + stepCtx, ok := StepFromContext(enriched) + require.True(t, ok) + require.Equal(t, runChatID, stepCtx.ChatID) +} + +func TestWrapModel_ReturnsDebugModel(t *testing.T) { + t.Parallel() + + model := &chattest.FakeModel{ProviderName: "provider", ModelName: "model"} + wrapped := WrapModel(model, &Service{}, RecorderOptions{}) + + require.NotSame(t, model, wrapped) + require.IsType(t, &debugModel{}, wrapped) + require.Implements(t, (*fantasy.LanguageModel)(nil), wrapped) + require.Equal(t, model.Provider(), wrapped.Provider()) + require.Equal(t, model.Model(), wrapped.Model()) +} diff --git a/coderd/x/chatd/chatdebug/redaction.go b/coderd/x/chatd/chatdebug/redaction.go new file mode 100644 index 0000000000000..784c5de65ca41 --- /dev/null +++ b/coderd/x/chatd/chatdebug/redaction.go @@ -0,0 +1,227 @@ +package chatdebug + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "net/http" + "strings" + + "golang.org/x/xerrors" +) + +// RedactedValue replaces sensitive values in debug payloads. +const RedactedValue = "[REDACTED]" + +var sensitiveHeaderNames = map[string]struct{}{ + "authorization": {}, + "x-api-key": {}, + "api-key": {}, + "proxy-authorization": {}, + "cookie": {}, + "set-cookie": {}, +} + +// sensitiveJSONKeyFragments triggers redaction for JSON keys containing +// these substrings. Notably, "token" is intentionally absent because it +// false-positively redacts LLM token-usage fields (input_tokens, +// output_tokens, prompt_tokens, completion_tokens, reasoning_tokens, +// cache_creation_input_tokens, cache_read_input_tokens, etc.). Auth- +// related token fields are caught by the exact-match set below. +var sensitiveJSONKeyFragments = []string{ + "secret", + "password", + "authorization", + "credential", +} + +// sensitiveJSONKeyExact matches auth-related token/key field names +// without false-positiving on LLM usage counters. Includes both +// snake_case originals and their camelCase-lowered equivalents +// (e.g. "accessToken" → "accesstoken") so that providers using +// either convention are caught. +var sensitiveJSONKeyExact = map[string]struct{}{ + "token": {}, + "access_token": {}, + "accesstoken": {}, + "refresh_token": {}, + "refreshtoken": {}, + "id_token": {}, + "idtoken": {}, + "api_token": {}, + "apitoken": {}, + "api_key": {}, + "apikey": {}, + "api-key": {}, + "x-api-key": {}, + "auth_token": {}, + "authtoken": {}, + "bearer_token": {}, + "bearertoken": {}, + "session_token": {}, + "sessiontoken": {}, + "security_token": {}, + "securitytoken": {}, + "private_key": {}, + "privatekey": {}, + "signing_key": {}, + "signingkey": {}, + "secret_key": {}, + "secretkey": {}, +} + +// RedactHeaders returns a flattened copy of h with sensitive values redacted. +func RedactHeaders(h http.Header) map[string]string { + if h == nil { + return nil + } + + redacted := make(map[string]string, len(h)) + for name, values := range h { + if isSensitiveName(name) { + redacted[name] = RedactedValue + continue + } + redacted[name] = strings.Join(values, ", ") + } + return redacted +} + +// RedactJSONSecrets redacts sensitive JSON values by key name. When +// the input is not valid JSON (truncated body, HTML error page, etc.) +// the raw bytes are replaced entirely with a diagnostic placeholder +// to avoid leaking credentials from malformed payloads. +func RedactJSONSecrets(data []byte) []byte { + if len(data) == 0 { + return data + } + + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.UseNumber() + + var value any + if err := decoder.Decode(&value); err != nil { + // Cannot parse: replace entirely to prevent credential leaks + // from non-JSON error responses (HTML pages, partial bodies). + return []byte(`{"error":"chatdebug: body is not valid JSON, redacted for safety"}`) + } + if err := consumeJSONEOF(decoder); err != nil { + return []byte(`{"error":"chatdebug: body contains extra JSON values, redacted for safety"}`) + } + + redacted, changed := redactJSONValue(value) + if !changed { + return data + } + + encoded, err := json.Marshal(redacted) + if err != nil { + return data + } + return encoded +} + +func consumeJSONEOF(decoder *json.Decoder) error { + var extra any + err := decoder.Decode(&extra) + if errors.Is(err, io.EOF) { + return nil + } + if err == nil { + return xerrors.New("chatdebug: extra JSON values") + } + return err +} + +var safeRateLimitHeaderNames = map[string]struct{}{ + "anthropic-ratelimit-requests-limit": {}, + "anthropic-ratelimit-requests-remaining": {}, + "anthropic-ratelimit-requests-reset": {}, + "anthropic-ratelimit-tokens-limit": {}, + "anthropic-ratelimit-tokens-remaining": {}, + "anthropic-ratelimit-tokens-reset": {}, + "x-ratelimit-limit-requests": {}, + "x-ratelimit-limit-tokens": {}, + "x-ratelimit-remaining-requests": {}, + "x-ratelimit-remaining-tokens": {}, + "x-ratelimit-reset-requests": {}, + "x-ratelimit-reset-tokens": {}, +} + +// isSensitiveName reports whether a name (header or query parameter) +// looks like a credential-carrying key. Exact-match headers are +// checked first, then the rate-limit allowlist, then substring +// patterns for API keys and auth tokens. +func isSensitiveName(name string) bool { + lowerName := strings.ToLower(name) + if _, ok := sensitiveHeaderNames[lowerName]; ok { + return true + } + if _, ok := safeRateLimitHeaderNames[lowerName]; ok { + return false + } + if strings.Contains(lowerName, "api-key") || + strings.Contains(lowerName, "api_key") || + strings.Contains(lowerName, "apikey") { + return true + } + // Catch any header containing "token" (e.g. Token, X-Token, + // X-Auth-Token). Safe rate-limit headers like + // x-ratelimit-remaining-tokens are already allowlisted above + // and will not reach this point. + if strings.Contains(lowerName, "token") { + return true + } + return strings.Contains(lowerName, "secret") || + strings.Contains(lowerName, "bearer") +} + +func isSensitiveJSONKey(key string) bool { + lowerKey := strings.ToLower(key) + if _, ok := sensitiveJSONKeyExact[lowerKey]; ok { + return true + } + for _, fragment := range sensitiveJSONKeyFragments { + if strings.Contains(lowerKey, fragment) { + return true + } + } + return false +} + +func redactJSONValue(value any) (any, bool) { + switch typed := value.(type) { + case map[string]any: + changed := false + for key, child := range typed { + if isSensitiveJSONKey(key) { + if current, ok := child.(string); ok && current == RedactedValue { + continue + } + typed[key] = RedactedValue + changed = true + continue + } + + redactedChild, childChanged := redactJSONValue(child) + if childChanged { + typed[key] = redactedChild + changed = true + } + } + return typed, changed + case []any: + changed := false + for i, child := range typed { + redactedChild, childChanged := redactJSONValue(child) + if childChanged { + typed[i] = redactedChild + changed = true + } + } + return typed, changed + default: + return value, false + } +} diff --git a/coderd/x/chatd/chatdebug/redaction_test.go b/coderd/x/chatd/chatdebug/redaction_test.go new file mode 100644 index 0000000000000..ac45d53262443 --- /dev/null +++ b/coderd/x/chatd/chatdebug/redaction_test.go @@ -0,0 +1,277 @@ +package chatdebug_test + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" +) + +func TestRedactHeaders(t *testing.T) { + t.Parallel() + + t.Run("nil input", func(t *testing.T) { + t.Parallel() + + require.Nil(t, chatdebug.RedactHeaders(nil)) + }) + + t.Run("empty header", func(t *testing.T) { + t.Parallel() + + redacted := chatdebug.RedactHeaders(http.Header{}) + require.NotNil(t, redacted) + require.Empty(t, redacted) + }) + + t.Run("authorization redacted and others preserved", func(t *testing.T) { + t.Parallel() + + headers := http.Header{ + "Authorization": {"Bearer secret-token"}, + "Accept": {"application/json"}, + } + + redacted := chatdebug.RedactHeaders(headers) + require.Equal(t, chatdebug.RedactedValue, redacted["Authorization"]) + require.Equal(t, "application/json", redacted["Accept"]) + }) + + t.Run("multi-value headers are flattened", func(t *testing.T) { + t.Parallel() + + headers := http.Header{ + "Accept": {"application/json", "text/plain"}, + } + + redacted := chatdebug.RedactHeaders(headers) + require.Equal(t, "application/json, text/plain", redacted["Accept"]) + }) + + t.Run("header name matching is case insensitive", func(t *testing.T) { + t.Parallel() + + lowerAuthorization := "authorization" + upperAuthorization := "AUTHORIZATION" + headers := http.Header{ + lowerAuthorization: {"lower"}, + upperAuthorization: {"upper"}, + } + + redacted := chatdebug.RedactHeaders(headers) + require.Equal(t, chatdebug.RedactedValue, redacted[lowerAuthorization]) + require.Equal(t, chatdebug.RedactedValue, redacted[upperAuthorization]) + }) + + t.Run("token and secret substrings are redacted", func(t *testing.T) { + t.Parallel() + + traceHeader := "X-Trace-ID" + headers := http.Header{ + "X-Auth-Token": {"abc"}, + "X-Custom-Secret": {"def"}, + "X-Bearer": {"ghi"}, + traceHeader: {"trace"}, + } + + redacted := chatdebug.RedactHeaders(headers) + require.Equal(t, chatdebug.RedactedValue, redacted["X-Auth-Token"]) + require.Equal(t, chatdebug.RedactedValue, redacted["X-Custom-Secret"]) + require.Equal(t, chatdebug.RedactedValue, redacted["X-Bearer"]) + require.Equal(t, "trace", redacted[traceHeader]) + }) + + t.Run("known safe rate limit headers containing token are not redacted", func(t *testing.T) { + t.Parallel() + + headers := http.Header{ + "Anthropic-Ratelimit-Tokens-Limit": {"1000000"}, + "Anthropic-Ratelimit-Tokens-Remaining": {"999000"}, + "Anthropic-Ratelimit-Tokens-Reset": {"2026-03-31T08:55:26Z"}, + "X-RateLimit-Limit-Tokens": {"120000"}, + "X-RateLimit-Remaining-Tokens": {"119500"}, + "X-RateLimit-Reset-Tokens": {"12ms"}, + } + + redacted := chatdebug.RedactHeaders(headers) + require.Equal(t, "1000000", redacted["Anthropic-Ratelimit-Tokens-Limit"]) + require.Equal(t, "999000", redacted["Anthropic-Ratelimit-Tokens-Remaining"]) + require.Equal(t, "2026-03-31T08:55:26Z", redacted["Anthropic-Ratelimit-Tokens-Reset"]) + require.Equal(t, "120000", redacted["X-RateLimit-Limit-Tokens"]) + require.Equal(t, "119500", redacted["X-RateLimit-Remaining-Tokens"]) + require.Equal(t, "12ms", redacted["X-RateLimit-Reset-Tokens"]) + }) + + t.Run("non-standard headers with api-key pattern are redacted", func(t *testing.T) { + t.Parallel() + + headers := http.Header{ + "X-Custom-Api-Key": {"secret-key"}, + "X-Custom-Secret": {"secret-val"}, + "X-Custom-Session-Token": {"session-id"}, + } + + redacted := chatdebug.RedactHeaders(headers) + require.Equal(t, chatdebug.RedactedValue, redacted["X-Custom-Api-Key"]) + require.Equal(t, chatdebug.RedactedValue, redacted["X-Custom-Secret"]) + require.Equal(t, chatdebug.RedactedValue, redacted["X-Custom-Session-Token"]) + }) + + t.Run("rate limit headers with token in name are preserved", func(t *testing.T) { + t.Parallel() + + // Rate-limit headers containing "token" should NOT be redacted + // because they carry usage/limit counts, not credentials. + headers := http.Header{ + "X-Ratelimit-Limit-Tokens": {"1000000"}, + "X-Ratelimit-Remaining-Tokens": {"999000"}, + } + + redacted := chatdebug.RedactHeaders(headers) + require.Equal(t, "1000000", redacted["X-Ratelimit-Limit-Tokens"]) + require.Equal(t, "999000", redacted["X-Ratelimit-Remaining-Tokens"]) + }) + + t.Run("original header is not modified", func(t *testing.T) { + t.Parallel() + + headers := http.Header{ + "Authorization": {"Bearer keep-me"}, + "X-Test": {"value"}, + } + + redacted := chatdebug.RedactHeaders(headers) + redacted["X-Test"] = "changed" + + require.Equal(t, []string{"Bearer keep-me"}, headers["Authorization"]) + require.Equal(t, []string{"value"}, headers["X-Test"]) + require.Equal(t, chatdebug.RedactedValue, redacted["Authorization"]) + }) + t.Run("api-key header variants are redacted", func(t *testing.T) { + t.Parallel() + + headers := http.Header{ + "X-Goog-Api-Key": {"secret"}, + "X-Api_Key": {"other-secret"}, + "X-Safe": {"ok"}, + } + + redacted := chatdebug.RedactHeaders(headers) + require.Equal(t, chatdebug.RedactedValue, redacted["X-Goog-Api-Key"]) + require.Equal(t, chatdebug.RedactedValue, redacted["X-Api_Key"]) + require.Equal(t, "ok", redacted["X-Safe"]) + }) + + t.Run("plain token headers are redacted", func(t *testing.T) { + t.Parallel() + + // Headers like "Token" or "X-Token" should be redacted + // even without auth/session/access qualifiers. + headers := http.Header{ + "Token": {"my-secret-token"}, + "X-Token": {"another-secret"}, + "X-Safe": {"ok"}, + } + + redacted := chatdebug.RedactHeaders(headers) + require.Equal(t, chatdebug.RedactedValue, redacted["Token"]) + require.Equal(t, chatdebug.RedactedValue, redacted["X-Token"]) + require.Equal(t, "ok", redacted["X-Safe"]) + }) +} + +func TestRedactJSONSecrets(t *testing.T) { + t.Parallel() + + t.Run("redacts top level secret fields", func(t *testing.T) { + t.Parallel() + + input := []byte(`{"api_key":"abc","token":"def","password":"ghi","safe":"ok"}`) + redacted := chatdebug.RedactJSONSecrets(input) + require.JSONEq(t, `{"api_key":"[REDACTED]","token":"[REDACTED]","password":"[REDACTED]","safe":"ok"}`, string(redacted)) + }) + + t.Run("redacts security_token exact key", func(t *testing.T) { + t.Parallel() + + input := []byte(`{"security_token":"s3cret","securityToken":"tok","safe":"ok"}`) + redacted := chatdebug.RedactJSONSecrets(input) + require.JSONEq(t, `{"security_token":"[REDACTED]","securityToken":"[REDACTED]","safe":"ok"}`, string(redacted)) + }) + + t.Run("preserves LLM token usage fields", func(t *testing.T) { + t.Parallel() + + input := []byte(`{"input_tokens":100,"output_tokens":50,"prompt_tokens":80,"completion_tokens":20,"reasoning_tokens":10,"cache_creation_input_tokens":5,"cache_read_input_tokens":3,"total_tokens":150,"max_tokens":4096,"max_output_tokens":2048}`) + redacted := chatdebug.RedactJSONSecrets(input) + // All usage/limit fields should be preserved, not redacted. + require.Equal(t, input, redacted) + }) + + t.Run("redacts nested objects", func(t *testing.T) { + t.Parallel() + + input := []byte(`{"outer":{"nested_secret":"abc","safe":1},"keep":true}`) + redacted := chatdebug.RedactJSONSecrets(input) + require.JSONEq(t, `{"outer":{"nested_secret":"[REDACTED]","safe":1},"keep":true}`, string(redacted)) + }) + + t.Run("redacts arrays of objects", func(t *testing.T) { + t.Parallel() + + input := []byte(`[{"token":"abc"},{"value":1,"credentials":{"access_key":"def"}}]`) + redacted := chatdebug.RedactJSONSecrets(input) + require.JSONEq(t, `[{"token":"[REDACTED]"},{"value":1,"credentials":"[REDACTED]"}]`, string(redacted)) + }) + + t.Run("concatenated JSON is replaced with diagnostic", func(t *testing.T) { + t.Parallel() + + input := []byte(`{"token":"abc"}{"safe":"ok"}`) + result := chatdebug.RedactJSONSecrets(input) + require.Contains(t, string(result), "extra JSON values") + }) + + t.Run("non JSON input is replaced with diagnostic", func(t *testing.T) { + t.Parallel() + + input := []byte("not json") + result := chatdebug.RedactJSONSecrets(input) + require.Contains(t, string(result), "not valid JSON") + }) + + t.Run("empty input is unchanged", func(t *testing.T) { + t.Parallel() + + input := []byte{} + require.Equal(t, input, chatdebug.RedactJSONSecrets(input)) + }) + + t.Run("JSON without sensitive keys is unchanged", func(t *testing.T) { + t.Parallel() + + input := []byte(`{"safe":"ok","nested":{"value":1}}`) + require.Equal(t, input, chatdebug.RedactJSONSecrets(input)) + }) + + t.Run("key matching is case insensitive", func(t *testing.T) { + t.Parallel() + + input := []byte(`{"API_KEY":"abc","Token":"def","PASSWORD":"ghi"}`) + redacted := chatdebug.RedactJSONSecrets(input) + require.JSONEq(t, `{"API_KEY":"[REDACTED]","Token":"[REDACTED]","PASSWORD":"[REDACTED]"}`, string(redacted)) + }) + + t.Run("camelCase token field names are redacted", func(t *testing.T) { + t.Parallel() + + // Providers may use camelCase (e.g. accessToken, refreshToken). + // These should be redacted even though they don't match the + // snake_case originals exactly. + input := []byte(`{"accessToken":"abc","refreshToken":"def","authToken":"ghi","input_tokens":100,"output_tokens":50}`) + redacted := chatdebug.RedactJSONSecrets(input) + require.JSONEq(t, `{"accessToken":"[REDACTED]","refreshToken":"[REDACTED]","authToken":"[REDACTED]","input_tokens":100,"output_tokens":50}`, string(redacted)) + }) +} diff --git a/coderd/x/chatd/chatdebug/reuse_step_test.go b/coderd/x/chatd/chatdebug/reuse_step_test.go new file mode 100644 index 0000000000000..90a06b7e2157b --- /dev/null +++ b/coderd/x/chatd/chatdebug/reuse_step_test.go @@ -0,0 +1,74 @@ +package chatdebug //nolint:testpackage // Uses unexported recorder helpers. + +import ( + "context" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/testutil" +) + +func TestBeginStepReuseStep(t *testing.T) { + t.Parallel() + + t.Run("reuses handle under ReuseStep", func(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + ownerID := uuid.New() + runID := uuid.New() + t.Cleanup(func() { CleanupStepCounter(runID) }) + + svc := NewService(nil, testutil.Logger(t), nil) + ctx := ContextWithRun(context.Background(), &RunContext{RunID: runID, ChatID: chatID}) + ctx = ReuseStep(ctx) + opts := RecorderOptions{ChatID: chatID, OwnerID: ownerID} + + firstHandle, firstEnriched := beginStep(ctx, svc, opts, OperationStream, nil) + secondHandle, secondEnriched := beginStep(ctx, svc, opts, OperationStream, nil) + + require.NotNil(t, firstHandle) + require.Same(t, firstHandle, secondHandle) + require.Same(t, firstHandle.stepCtx, secondHandle.stepCtx) + require.Same(t, firstHandle.sink, secondHandle.sink) + require.Equal(t, runID, firstHandle.stepCtx.RunID) + require.Equal(t, chatID, firstHandle.stepCtx.ChatID) + require.Equal(t, int32(1), firstHandle.stepCtx.StepNumber) + require.Equal(t, OperationStream, firstHandle.stepCtx.Operation) + require.NotEqual(t, uuid.Nil, firstHandle.stepCtx.StepID) + + firstStepCtx, ok := StepFromContext(firstEnriched) + require.True(t, ok) + secondStepCtx, ok := StepFromContext(secondEnriched) + require.True(t, ok) + require.Same(t, firstStepCtx, secondStepCtx) + require.Same(t, firstHandle.stepCtx, firstStepCtx) + require.Same(t, attemptSinkFromContext(firstEnriched), attemptSinkFromContext(secondEnriched)) + }) + + t.Run("creates new handles without ReuseStep", func(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + ownerID := uuid.New() + runID := uuid.New() + t.Cleanup(func() { CleanupStepCounter(runID) }) + + svc := NewService(nil, testutil.Logger(t), nil) + ctx := ContextWithRun(context.Background(), &RunContext{RunID: runID, ChatID: chatID}) + opts := RecorderOptions{ChatID: chatID, OwnerID: ownerID} + + firstHandle, _ := beginStep(ctx, svc, opts, OperationStream, nil) + secondHandle, _ := beginStep(ctx, svc, opts, OperationStream, nil) + + require.NotNil(t, firstHandle) + require.NotNil(t, secondHandle) + require.NotSame(t, firstHandle, secondHandle) + require.NotSame(t, firstHandle.sink, secondHandle.sink) + require.Equal(t, int32(1), firstHandle.stepCtx.StepNumber) + require.Equal(t, int32(2), secondHandle.stepCtx.StepNumber) + require.NotEqual(t, firstHandle.stepCtx.StepID, secondHandle.stepCtx.StepID) + }) +} diff --git a/coderd/x/chatd/chatdebug/stubs.go b/coderd/x/chatd/chatdebug/stubs.go index aee4b85ff7e28..72dc5246d0b02 100644 --- a/coderd/x/chatd/chatdebug/stubs.go +++ b/coderd/x/chatd/chatdebug/stubs.go @@ -6,6 +6,7 @@ import ( "strings" "sync" "sync/atomic" + "time" "unicode/utf8" "github.com/google/uuid" @@ -15,43 +16,76 @@ import ( "github.com/coder/coder/v2/coderd/database/pubsub" ) -// This branch-02 compatibility shim forward-declares recorder, service, -// and summary symbols that land in later stacked branches. Delete this -// file once recorder.go, service.go, and summary.go are available here. - -// RecorderOptions identifies the chat/model context for debug recording. -type RecorderOptions struct { - ChatID uuid.UUID - OwnerID uuid.UUID - Provider string - Model string -} +// This compatibility shim forward-declares service and summary symbols +// that land in later stacked branches. Delete this file once service.go +// and summary.go are available here. // Service is a placeholder for the later chat debug persistence service. -type Service struct{} +type Service struct { + log slog.Logger +} -// NewService constructs the branch-02 placeholder chat debug service. -func NewService(_ database.Store, _ slog.Logger, _ pubsub.Pubsub) *Service { - return &Service{} +// CreateStepParams identifies the data recorded when a debug step starts. +type CreateStepParams struct { + RunID uuid.UUID + ChatID uuid.UUID + StepNumber int32 + Operation Operation + Status Status + HistoryTipMessageID int64 + NormalizedRequest any } -type attemptSink struct{} +// UpdateStepParams identifies the data recorded when a debug step finishes. +type UpdateStepParams struct { + ID uuid.UUID + ChatID uuid.UUID + Status Status + NormalizedResponse any + Usage any + Attempts []Attempt + Error any + Metadata any + FinishedAt time.Time +} -type attemptSinkKey struct{} +// NewService constructs the placeholder chat debug service. +func NewService(_ database.Store, log slog.Logger, _ pubsub.Pubsub) *Service { + return &Service{log: log} +} -func withAttemptSink(ctx context.Context, sink *attemptSink) context.Context { - if sink == nil { - panic("chatdebug: nil attemptSink") - } - return context.WithValue(ctx, attemptSinkKey{}, sink) +// IsEnabled reports whether debug recording is enabled for a chat owner. +func (*Service) IsEnabled(context.Context, uuid.UUID, uuid.UUID) bool { + return true } -func attemptSinkFromContext(ctx context.Context) *attemptSink { - sink, _ := ctx.Value(attemptSinkKey{}).(*attemptSink) - return sink +// CreateStep synthesizes a debug step so recorder tests can exercise the +// wrapper without requiring the later persistence service implementation. +func (*Service) CreateStep( + _ context.Context, + params CreateStepParams, +) (database.ChatDebugStep, error) { + return database.ChatDebugStep{ + ID: uuid.New(), + RunID: params.RunID, + ChatID: params.ChatID, + StepNumber: params.StepNumber, + Operation: string(params.Operation), + Status: string(params.Status), + }, nil } -var stepCounters sync.Map // map[uuid.UUID]*atomic.Int32 +// UpdateStep accepts final step state once recording completes. +func (*Service) UpdateStep( + _ context.Context, + params UpdateStepParams, +) (database.ChatDebugStep, error) { + return database.ChatDebugStep{ + ID: params.ID, + ChatID: params.ChatID, + Status: string(params.Status), + }, nil +} // runRefCounts tracks how many live RunContext instances reference each // RunID. Cleanup of shared state (step counters) is deferred until the @@ -96,125 +130,6 @@ func releaseRunRef(runID uuid.UUID) { } } -func nextStepNumber(runID uuid.UUID) int32 { - val, _ := stepCounters.LoadOrStore(runID, &atomic.Int32{}) - counter, ok := val.(*atomic.Int32) - if !ok { - panic("chatdebug: invalid step counter type") - } - return counter.Add(1) -} - -// CleanupStepCounter removes per-run step counter and reference count -// state. This is used by tests and later stacked branches that have a -// real run lifecycle. -func CleanupStepCounter(runID uuid.UUID) { - stepCounters.Delete(runID) - runRefCounts.Delete(runID) -} - -type stepHandle struct { - stepCtx *StepContext - sink *attemptSink - mu sync.Mutex - status Status - response any - usage any - err any - metadata any -} - -func beginStep( - ctx context.Context, - svc *Service, - opts RecorderOptions, - op Operation, - _ any, -) (*stepHandle, context.Context) { - if svc == nil { - return nil, ctx - } - - rc, ok := RunFromContext(ctx) - if !ok || rc.RunID == uuid.Nil { - return nil, ctx - } - - if holder, reuseStep := reuseHolderFromContext(ctx); reuseStep { - holder.mu.Lock() - defer holder.mu.Unlock() - // Only reuse the cached handle if it belongs to the same run. - // A different RunContext means a new logical run, so we must - // create a fresh step to avoid cross-run attribution. - if holder.handle != nil && holder.handle.stepCtx.RunID == rc.RunID { - enriched := ContextWithStep(ctx, holder.handle.stepCtx) - enriched = withAttemptSink(enriched, holder.handle.sink) - return holder.handle, enriched - } - - handle, enriched := newStepHandle(ctx, rc, opts, op) - holder.handle = handle - return handle, enriched - } - - return newStepHandle(ctx, rc, opts, op) -} - -func newStepHandle( - ctx context.Context, - rc *RunContext, - opts RecorderOptions, - op Operation, -) (*stepHandle, context.Context) { - if rc == nil || rc.RunID == uuid.Nil { - return nil, ctx - } - - chatID := opts.ChatID - if chatID == uuid.Nil { - chatID = rc.ChatID - } - - handle := &stepHandle{ - stepCtx: &StepContext{ - StepID: uuid.New(), - RunID: rc.RunID, - ChatID: chatID, - StepNumber: nextStepNumber(rc.RunID), - Operation: op, - HistoryTipMessageID: rc.HistoryTipMessageID, - }, - sink: &attemptSink{}, - } - enriched := ContextWithStep(ctx, handle.stepCtx) - enriched = withAttemptSink(enriched, handle.sink) - return handle, enriched -} - -func (h *stepHandle) finish( - _ context.Context, - status Status, - response any, - usage any, - err any, - metadata any, -) { - if h == nil || h.stepCtx == nil { - return - } - // Guard with a mutex so concurrent callers (e.g. retried stream - // wrappers sharing a reused handle) don't race. Unlike sync.Once, - // later retries are allowed to overwrite earlier failure results so - // the step reflects the final outcome. - h.mu.Lock() - defer h.mu.Unlock() - h.status = status - h.response = response - h.usage = usage - h.err = err - h.metadata = metadata -} - // whitespaceRun matches one or more consecutive whitespace characters. var whitespaceRun = regexp.MustCompile(`\s+`) diff --git a/coderd/x/chatd/chatdebug/stubs_internal_test.go b/coderd/x/chatd/chatdebug/stubs_internal_test.go index b7bfa81c181c2..75d0aabdd444a 100644 --- a/coderd/x/chatd/chatdebug/stubs_internal_test.go +++ b/coderd/x/chatd/chatdebug/stubs_internal_test.go @@ -2,7 +2,6 @@ package chatdebug import ( "context" - "net/http" "testing" "unicode/utf8" @@ -19,15 +18,6 @@ func TestBeginStep_SkipsNilRunID(t *testing.T) { require.Equal(t, ctx, enriched) } -func TestNewStepHandle_SkipsNilRunID(t *testing.T) { - t.Parallel() - - ctx := context.Background() - handle, enriched := newStepHandle(ctx, &RunContext{ChatID: uuid.New()}, RecorderOptions{ChatID: uuid.New()}, OperationGenerate) - require.Nil(t, handle) - require.Equal(t, ctx, enriched) -} - func TestTruncateLabel(t *testing.T) { t.Parallel() @@ -58,25 +48,3 @@ func TestTruncateLabel(t *testing.T) { }) } } - -// RedactedValue replaces sensitive values in debug payloads. -const RedactedValue = "[REDACTED]" - -// RecordingTransport is the branch-02 placeholder HTTP recording transport. -type RecordingTransport struct { - Base http.RoundTripper -} - -var _ http.RoundTripper = (*RecordingTransport)(nil) - -func (t *RecordingTransport) RoundTrip(req *http.Request) (*http.Response, error) { - if req == nil { - panic("chatdebug: nil request") - } - - base := t.Base - if base == nil { - base = http.DefaultTransport - } - return base.RoundTrip(req) -} diff --git a/coderd/x/chatd/chatdebug/transport.go b/coderd/x/chatd/chatdebug/transport.go new file mode 100644 index 0000000000000..6741aa8366840 --- /dev/null +++ b/coderd/x/chatd/chatdebug/transport.go @@ -0,0 +1,382 @@ +package chatdebug + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "mime" + "net/http" + "net/url" + "strings" + "sync" + "time" +) + +// attemptStatusCompleted is the status recorded when a response body +// is fully read without transport-level errors. +const attemptStatusCompleted = "completed" + +// attemptStatusFailed is the status recorded when a transport error +// or body read error occurs. +const attemptStatusFailed = "failed" + +// maxRecordedRequestBodyBytes caps in-memory request capture when GetBody +// is available. +const maxRecordedRequestBodyBytes = 50_000 + +// maxRecordedResponseBodyBytes caps in-memory response capture. +const maxRecordedResponseBodyBytes = 50_000 + +// RecordingTransport captures HTTP request/response data for debug steps. +// When the request context carries an attemptSink, it records each round +// trip. Otherwise it delegates directly. +type RecordingTransport struct { + // Base is the underlying transport. nil defaults to http.DefaultTransport. + Base http.RoundTripper +} + +var _ http.RoundTripper = (*RecordingTransport)(nil) + +func (t *RecordingTransport) RoundTrip(req *http.Request) (*http.Response, error) { + if req == nil { + panic("chatdebug: nil request") + } + + base := t.Base + if base == nil { + base = http.DefaultTransport + } + + sink := attemptSinkFromContext(req.Context()) + if sink == nil { + return base.RoundTrip(req) + } + + requestHeaders := RedactHeaders(req.Header) + + // Capture method and URL/path from the request. + method := req.Method + reqURL := "" + reqPath := "" + if req.URL != nil { + reqURL = redactURL(req.URL) + reqPath = req.URL.Path + } + + requestBody, err := captureRequestBody(req) + if err != nil { + return nil, err + } + attemptNumber := sink.nextAttemptNumber() + + startedAt := time.Now() + resp, err := base.RoundTrip(req) + finishedAt := time.Now() + durationMs := finishedAt.Sub(startedAt).Milliseconds() + if err != nil { + sink.record(Attempt{ + Number: attemptNumber, + Status: attemptStatusFailed, + Method: method, + URL: reqURL, + Path: reqPath, + StartedAt: startedAt.UTC().Format(time.RFC3339Nano), + FinishedAt: finishedAt.UTC().Format(time.RFC3339Nano), + RequestHeaders: requestHeaders, + RequestBody: requestBody, + Error: err.Error(), + DurationMs: durationMs, + }) + return nil, err + } + + respHeaders := RedactHeaders(resp.Header) + resp.Body = &recordingBody{ + inner: resp.Body, + sink: sink, + startedAt: startedAt, + contentLength: resp.ContentLength, + base: Attempt{ + Number: attemptNumber, + Method: method, + URL: reqURL, + Path: reqPath, + RequestHeaders: requestHeaders, + RequestBody: requestBody, + ResponseStatus: resp.StatusCode, + ResponseHeaders: respHeaders, + DurationMs: durationMs, + }, + } + + return resp, nil +} + +func redactURL(u *url.URL) string { + if u == nil { + return "" + } + clone := *u + clone.User = nil + q := clone.Query() + for key, values := range q { + if isSensitiveName(key) || isSensitiveJSONKey(key) { + for i := range values { + values[i] = RedactedValue + } + q[key] = values + } + } + clone.RawQuery = q.Encode() + return clone.String() +} + +func captureRequestBody(req *http.Request) ([]byte, error) { + if req == nil || req.Body == nil { + return nil, nil + } + + if req.GetBody != nil { + clone, err := req.GetBody() + if err == nil { + defer clone.Close() + limited, err := io.ReadAll(io.LimitReader(clone, maxRecordedRequestBodyBytes+1)) + if err == nil { + if len(limited) > maxRecordedRequestBodyBytes { + return []byte("[TRUNCATED]"), nil + } + return RedactJSONSecrets(limited), nil + } + } + } + + // Without GetBody we cannot safely capture the request body without + // fully consuming a potentially large or streaming body before the + // request is sent. Skip capture in that case to keep debug logging + // lightweight and non-invasive. + return nil, nil +} + +type recordingBody struct { + inner io.ReadCloser + contentLength int64 + sink *attemptSink + base Attempt + startedAt time.Time + + mu sync.Mutex + buf bytes.Buffer + truncated bool + sawEOF bool + bytesRead int64 + + recordOnce sync.Once + closeOnce sync.Once +} + +func (r *recordingBody) Read(p []byte) (int, error) { + n, err := r.inner.Read(p) + + r.mu.Lock() + r.bytesRead += int64(n) + if n > 0 && !r.truncated { + remaining := maxRecordedResponseBodyBytes - r.buf.Len() + if remaining > 0 { + toWrite := n + if toWrite > remaining { + toWrite = remaining + r.truncated = true + } + _, _ = r.buf.Write(p[:toWrite]) + } else { + r.truncated = true + } + } + if errors.Is(err, io.EOF) { + r.sawEOF = true + } + r.mu.Unlock() + + if err != nil { + r.record(err) + } + return n, err +} + +func (r *recordingBody) Close() error { + r.mu.Lock() + sawEOF := r.sawEOF + bytesRead := r.bytesRead + contentLength := r.contentLength + truncated := r.truncated + responseBody := append([]byte(nil), r.buf.Bytes()...) + r.mu.Unlock() + + contentType := r.base.ResponseHeaders["Content-Type"] + shouldDrainUnknownLengthJSON := contentLength < 0 && + !sawEOF && + bytesRead > 0 && + !truncated && + isCompleteUnknownLengthJSONBody(contentType, responseBody) + + // Always close the inner reader first so that stalled chunked + // bodies cannot block drainToEOF indefinitely. Once inner is + // closed, reads return immediately with an error or EOF. + var closeErr error + r.closeOnce.Do(func() { + closeErr = r.inner.Close() + }) + if closeErr != nil { + r.record(closeErr) + return closeErr + } + + // Drain remaining bytes that may already be buffered inside the + // HTTP transport after close. Because inner is closed, this + // finishes immediately rather than blocking on the network. + if shouldDrainUnknownLengthJSON { + // Best-effort drain; ignore errors since inner is closed. + _ = r.drainToEOF() + } + + r.mu.Lock() + sawEOF = r.sawEOF + bytesRead = r.bytesRead + contentLength = r.contentLength + truncated = r.truncated + responseBody = append([]byte(nil), r.buf.Bytes()...) + r.mu.Unlock() + + switch { + // Only check JSON completeness when the recording buffer is + // not truncated. A truncated buffer is an incomplete prefix + // of the body, so the completeness check would false-positive. + case sawEOF && !truncated && contentLength < 0 && isJSONLikeContentType(contentType) && !isCompleteUnknownLengthJSONBody(contentType, responseBody): + r.record(io.ErrUnexpectedEOF) + case sawEOF: + r.record(io.EOF) + case responseHasNoBody(r.base.Method, r.base.ResponseStatus): + r.record(nil) + case contentLength >= 0 && bytesRead >= contentLength: + r.record(nil) + case contentLength < 0 && !truncated && isCompleteUnknownLengthJSONBody(contentType, responseBody): + r.record(nil) + default: + r.record(io.ErrUnexpectedEOF) + } + return nil +} + +func responseHasNoBody(method string, statusCode int) bool { + if method == http.MethodHead { + return true + } + return statusCode == http.StatusNoContent || + statusCode == http.StatusNotModified || + (statusCode >= 100 && statusCode < 200) +} + +func isJSONLikeContentType(contentType string) bool { + mediaType, _, err := mime.ParseMediaType(contentType) + if err != nil { + mediaType = strings.TrimSpace(strings.Split(contentType, ";")[0]) + } + return mediaType == "application/json" || strings.HasSuffix(mediaType, "+json") +} + +// maxDrainBytes caps how many trailing bytes drainToEOF will consume. +// This prevents Close() from blocking indefinitely on a misbehaving +// or extremely large chunked body. +const maxDrainBytes = 64 * 1024 // 64 KB + +func (r *recordingBody) drainToEOF() error { + buf := make([]byte, 4*1024) + var drained int64 + for { + n, err := r.inner.Read(buf) + + r.mu.Lock() + r.bytesRead += int64(n) + drained += int64(n) + if n > 0 && !r.truncated { + remaining := maxRecordedResponseBodyBytes - r.buf.Len() + if remaining > 0 { + toWrite := n + if toWrite > remaining { + toWrite = remaining + r.truncated = true + } + _, _ = r.buf.Write(buf[:toWrite]) + } else { + r.truncated = true + } + } + if errors.Is(err, io.EOF) { + r.sawEOF = true + } + r.mu.Unlock() + + if err != nil { + if errors.Is(err, io.EOF) { + return nil + } + return err + } + + // Safety valve: stop draining after maxDrainBytes to prevent + // Close() from blocking indefinitely on a chunked body. + if drained >= maxDrainBytes { + return io.ErrUnexpectedEOF + } + } +} + +func isCompleteUnknownLengthJSONBody(contentType string, body []byte) bool { + if !isJSONLikeContentType(contentType) { + return false + } + + trimmed := bytes.TrimSpace(body) + if len(trimmed) == 0 { + return false + } + + decoder := json.NewDecoder(bytes.NewReader(trimmed)) + var value any + if err := decoder.Decode(&value); err != nil { + return false + } + var extra any + return errors.Is(decoder.Decode(&extra), io.EOF) +} + +func (r *recordingBody) record(err error) { + r.recordOnce.Do(func() { + finishedAt := time.Now() + + r.mu.Lock() + truncated := r.truncated + responseBody := append([]byte(nil), r.buf.Bytes()...) + base := r.base + startedAt := r.startedAt + r.mu.Unlock() + + if truncated { + base.ResponseBody = []byte("[TRUNCATED]") + } else { + base.ResponseBody = RedactJSONSecrets(responseBody) + } + base.StartedAt = startedAt.UTC().Format(time.RFC3339Nano) + base.FinishedAt = finishedAt.UTC().Format(time.RFC3339Nano) + // Recompute duration to include body read time. + base.DurationMs = finishedAt.Sub(startedAt).Milliseconds() + if err != nil && !errors.Is(err, io.EOF) { + base.Error = err.Error() + base.Status = attemptStatusFailed + } else { + base.Status = attemptStatusCompleted + } + r.sink.record(base) + }) +} diff --git a/coderd/x/chatd/chatdebug/transport_test.go b/coderd/x/chatd/chatdebug/transport_test.go new file mode 100644 index 0000000000000..b1da8e24f2787 --- /dev/null +++ b/coderd/x/chatd/chatdebug/transport_test.go @@ -0,0 +1,737 @@ +package chatdebug //nolint:testpackage // Uses unexported recorder helpers. + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" +) + +func newTestSinkContext(t *testing.T) (context.Context, *attemptSink) { + t.Helper() + + sink := &attemptSink{} + return withAttemptSink(context.Background(), sink), sink +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +type scriptedReadCloser struct { + chunks [][]byte + index int + offset int // byte offset within current chunk +} + +func (r *scriptedReadCloser) Read(p []byte) (int, error) { + if r.index >= len(r.chunks) { + return 0, io.EOF + } + chunk := r.chunks[r.index] + remaining := chunk[r.offset:] + n := copy(p, remaining) + r.offset += n + if r.offset >= len(chunk) { + r.index++ + r.offset = 0 + } + return n, nil +} + +func (*scriptedReadCloser) Close() error { + return nil +} + +func TestRecordingTransport_NoSink(t *testing.T) { + t.Parallel() + + gotMethod := make(chan string, 1) + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + gotMethod <- req.Method + _, _ = rw.Write([]byte("ok")) + })) + defer server.Close() + + client := &http.Client{ + Transport: &RecordingTransport{Base: server.Client().Transport}, + } + + req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, server.URL, nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, "ok", string(body)) + require.Equal(t, http.MethodGet, <-gotMethod) +} + +func TestRecordingTransport_CaptureRequest(t *testing.T) { + t.Parallel() + + const requestBody = `{"message":"hello","api_key":"super-secret"}` + + type receivedRequest struct { + authorization string + body []byte + } + gotRequest := make(chan receivedRequest, 1) + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + body, err := io.ReadAll(req.Body) + require.NoError(t, err) + gotRequest <- receivedRequest{ + authorization: req.Header.Get("Authorization"), + body: body, + } + _, _ = rw.Write([]byte(`{"ok":true}`)) + })) + defer server.Close() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{Base: server.Client().Transport}, + } + + req, err := http.NewRequestWithContext( + ctx, + http.MethodPost, + server.URL, + strings.NewReader(requestBody), + ) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer top-secret") + req.Header.Set("Content-Type", "application/json") + + resp, err := client.Do(req) + require.NoError(t, err) + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, 1, attempts[0].Number) + require.Equal(t, RedactedValue, attempts[0].RequestHeaders["Authorization"]) + require.Equal(t, "application/json", attempts[0].RequestHeaders["Content-Type"]) + require.JSONEq(t, `{"message":"hello","api_key":"[REDACTED]"}`, string(attempts[0].RequestBody)) + + received := <-gotRequest + require.JSONEq(t, requestBody, string(received.body)) + require.Equal(t, "Bearer top-secret", received.authorization) +} + +func TestRecordingTransport_RedactsSensitiveQueryParameters(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + _, _ = rw.Write([]byte(`ok`)) + })) + defer server.Close() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{Transport: &RecordingTransport{Base: server.Client().Transport}} + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL+`?api_key=secret&safe=ok`, nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Contains(t, attempts[0].URL, "api_key=%5BREDACTED%5D") + require.Contains(t, attempts[0].URL, "safe=ok") +} + +func TestRecordingTransport_TruncatesLargeRequestBodies(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + _, _ = io.Copy(io.Discard, req.Body) + _, _ = rw.Write([]byte(`ok`)) + })) + defer server.Close() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{Transport: &RecordingTransport{Base: server.Client().Transport}} + + large := strings.Repeat("x", maxRecordedRequestBodyBytes+1024) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, server.URL, strings.NewReader(large)) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, []byte("[TRUNCATED]"), attempts[0].RequestBody) +} + +func TestRecordingTransport_StripsURLUserinfo(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + _, _ = rw.Write([]byte(`ok`)) + })) + defer server.Close() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{Transport: &RecordingTransport{Base: server.Client().Transport}} + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.Replace(server.URL, "http://", "http://user:secret@", 1)+`?api_key=secret`, nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.NotContains(t, attempts[0].URL, "user:secret") + require.Contains(t, attempts[0].URL, "api_key=%5BREDACTED%5D") +} + +func TestRecordingTransport_SkipsNonReplayableRequestBodyCapture(t *testing.T) { + t.Parallel() + + const requestBody = `{"message":"hello"}` + gotRequest := make(chan []byte, 1) + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + body, err := io.ReadAll(req.Body) + require.NoError(t, err) + gotRequest <- body + _, _ = rw.Write([]byte(`ok`)) + })) + defer server.Close() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{Transport: &RecordingTransport{Base: server.Client().Transport}} + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, server.URL, io.NopCloser(strings.NewReader(requestBody))) + require.NoError(t, err) + req.GetBody = nil + + resp, err := client.Do(req) + require.NoError(t, err) + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + require.JSONEq(t, requestBody, string(<-gotRequest)) + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Nil(t, attempts[0].RequestBody) +} + +func TestRecordingTransport_CaptureResponse(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + rw.Header().Set("X-API-Key", "response-secret") + rw.Header().Set("X-Trace-ID", "trace-123") + rw.WriteHeader(http.StatusCreated) + _, _ = rw.Write([]byte(`{"token":"response-secret","safe":"ok"}`)) + })) + defer server.Close() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{Base: server.Client().Transport}, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + require.JSONEq(t, `{"token":"response-secret","safe":"ok"}`, string(body)) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, http.StatusCreated, attempts[0].ResponseStatus) + require.Equal(t, RedactedValue, attempts[0].ResponseHeaders["X-Api-Key"]) + require.Equal(t, "trace-123", attempts[0].ResponseHeaders["X-Trace-Id"]) + require.JSONEq(t, `{"token":"[REDACTED]","safe":"ok"}`, string(attempts[0].ResponseBody)) +} + +func TestRecordingTransport_CaptureResponseOnEOFWithoutClose(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + rw.Header().Set("Content-Type", "application/json") + rw.Header().Set("X-API-Key", "response-secret") + rw.WriteHeader(http.StatusAccepted) + _, _ = rw.Write([]byte(`{"token":"response-secret","safe":"ok"}`)) + })) + defer server.Close() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{Base: server.Client().Transport}, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.JSONEq(t, `{"token":"response-secret","safe":"ok"}`, string(body)) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, http.StatusAccepted, attempts[0].ResponseStatus) + require.Equal(t, "application/json", attempts[0].ResponseHeaders["Content-Type"]) + require.Equal(t, RedactedValue, attempts[0].ResponseHeaders["X-Api-Key"]) + require.JSONEq(t, `{"token":"[REDACTED]","safe":"ok"}`, string(attempts[0].ResponseBody)) + require.NoError(t, resp.Body.Close()) +} + +func TestRecordingTransport_StreamingBody(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + flusher, ok := rw.(http.Flusher) + require.True(t, ok) + + rw.Header().Set("Content-Type", "application/json") + _, _ = rw.Write([]byte(`{"safe":"stream",`)) + flusher.Flush() + _, _ = rw.Write([]byte(`"token":"chunk-secret"}`)) + flusher.Flush() + })) + defer server.Close() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{Base: server.Client().Transport}, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + buf := make([]byte, 5) + var body strings.Builder + for { + n, readErr := resp.Body.Read(buf) + if n > 0 { + _, writeErr := body.Write(buf[:n]) + require.NoError(t, writeErr) + } + if errors.Is(readErr, io.EOF) { + break + } + require.NoError(t, readErr) + } + require.NoError(t, resp.Body.Close()) + require.JSONEq(t, `{"safe":"stream","token":"chunk-secret"}`, body.String()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.JSONEq(t, `{"safe":"stream","token":"[REDACTED]"}`, string(attempts[0].ResponseBody)) +} + +func TestRecordingTransport_CloseAfterDecoderConsumesContentLengthSucceeds(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + rw.Header().Set("Content-Type", "application/json") + _, _ = rw.Write([]byte(`{"token":"response-secret","safe":"ok"}`)) + })) + defer server.Close() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{Transport: &RecordingTransport{Base: server.Client().Transport}} + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + var decoded map[string]string + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + require.Equal(t, "ok", decoded["safe"]) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + require.Empty(t, attempts[0].Error) +} + +func TestRecordingTransport_CloseAfterDecoderConsumesUnknownLengthJSONSucceeds(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test response exercises unknown-length close semantics. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &scriptedReadCloser{chunks: [][]byte{[]byte(`{"token":"response-secret","safe":"ok"}`)}}, + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + var decoded map[string]string + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + require.Equal(t, "ok", decoded["safe"]) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + require.Empty(t, attempts[0].Error) +} + +func TestRecordingTransport_CloseAfterDecoderConsumesUnknownLengthJSONWithTrailingDocumentMarksFailed(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test response exercises unknown-length close semantics. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &scriptedReadCloser{chunks: [][]byte{[]byte("{\"token\":\"response-secret\",\"safe\":\"ok\"}{\"token\":\"second\"}")}}, + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + var decoded map[string]string + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + require.Equal(t, "ok", decoded["safe"]) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusFailed, attempts[0].Status) + require.Equal(t, io.ErrUnexpectedEOF.Error(), attempts[0].Error) +} + +func TestRecordingTransport_CloseAfterDecoderConsumesUnknownLengthNDJSONMarksFailed(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test response exercises unknown-length close semantics. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/x-ndjson"}}, + Body: &scriptedReadCloser{chunks: [][]byte{[]byte("{\"token\":\"response-secret\",\"safe\":\"ok\"}\n{\"token\":\"second\"}\n")}}, + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + var decoded map[string]string + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + require.Equal(t, "ok", decoded["safe"]) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusFailed, attempts[0].Status) + require.Equal(t, io.ErrUnexpectedEOF.Error(), attempts[0].Error) +} + +func TestRecordingTransport_CloseAfterDecoderDrainsUnknownLengthSucceeds(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test response exercises unknown-length close semantics. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &scriptedReadCloser{chunks: [][]byte{[]byte(`{"token":"response-secret","safe":"ok"}`)}}, + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + var decoded map[string]string + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + require.Equal(t, "ok", decoded["safe"]) + _, err = io.Copy(io.Discard, resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + require.Empty(t, attempts[0].Error) +} + +func TestRecordingTransport_CloseWithoutReadingHeadResponseSucceeds(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test response exercises no-body close semantics. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &scriptedReadCloser{chunks: [][]byte{[]byte(`{"ignored":true}`)}}, + ContentLength: 13, + Request: req, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodHead, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + require.Empty(t, attempts[0].Error) +} + +func TestRecordingTransport_CloseWithoutReadingUnknownLengthMarksFailed(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test response exercises unknown-length close semantics. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &scriptedReadCloser{chunks: [][]byte{[]byte(`{"token":"response-secret","safe":"ok"}`)}}, + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusFailed, attempts[0].Status) + require.Equal(t, io.ErrUnexpectedEOF.Error(), attempts[0].Error) +} + +func TestRecordingTransport_PrematureCloseUnknownLengthMarksFailed(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test response exercises unknown-length close semantics. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &scriptedReadCloser{chunks: [][]byte{[]byte(`{"token":"response-secret","safe":"ok"}`)}}, + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + buf := make([]byte, 5) + _, err = resp.Body.Read(buf) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusFailed, attempts[0].Status) + require.Equal(t, io.ErrUnexpectedEOF.Error(), attempts[0].Error) +} + +func TestRecordingTransport_PrematureCloseMarksFailed(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + _, _ = rw.Write([]byte(`{"token":"response-secret","safe":"ok"}`)) + })) + defer server.Close() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{Transport: &RecordingTransport{Base: server.Client().Transport}} + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + buf := make([]byte, 5) + _, err = resp.Body.Read(buf) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusFailed, attempts[0].Status) +} + +func TestRecordingTransport_TruncatesLargeResponses(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + _, _ = rw.Write([]byte(strings.Repeat("x", maxRecordedResponseBodyBytes+1024))) + })) + defer server.Close() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{Transport: &RecordingTransport{Base: server.Client().Transport}} + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, []byte("[TRUNCATED]"), attempts[0].ResponseBody) +} + +func TestRecordingTransport_TransportError(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return nil, xerrors.New("transport exploded") + }), + }, + } + + req, err := http.NewRequestWithContext( + ctx, + http.MethodPost, + "http://example.invalid", + strings.NewReader(`{"password":"secret","safe":"ok"}`), + ) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer top-secret") + + resp, err := client.Do(req) + if resp != nil { + defer resp.Body.Close() + } + require.Nil(t, resp) + require.EqualError(t, err, "Post \"http://example.invalid\": transport exploded") + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, 1, attempts[0].Number) + require.Equal(t, RedactedValue, attempts[0].RequestHeaders["Authorization"]) + require.JSONEq(t, `{"password":"[REDACTED]","safe":"ok"}`, string(attempts[0].RequestBody)) + require.Zero(t, attempts[0].ResponseStatus) + require.Equal(t, "transport exploded", attempts[0].Error) + require.GreaterOrEqual(t, attempts[0].DurationMs, int64(0)) +} + +func TestRecordingTransport_NilBase(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + _, _ = rw.Write([]byte("ok")) + })) + defer server.Close() + + client := &http.Client{Transport: &RecordingTransport{}} + + req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, server.URL, nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, "ok", string(body)) +} From bdbf53be0ea70fbaec7af50f4498e2d805359e89 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Mon, 13 Apr 2026 19:27:20 +0200 Subject: [PATCH 2/3] fix(coderd/x/chatd/chatdebug): address remaining review feedback on PR #23915 - Skip JSON redaction for non-JSON content types (SSE, text/plain) to preserve debug content instead of replacing with diagnostic placeholder. Falls back to redaction for unknown/missing types. - Sanitize transport error strings by redacting embedded URLs that may contain credentials (userinfo, query parameters). - Add test cleanup to TestNextStepNumber_Concurrent for consistency. - Add tests for SSE read-to-EOF, SSE closed-early, text/plain preservation, and URL-containing transport error sanitization. Change-Id: I4f59a18e1fb10240fdd8fbb176ea953e08a9cd2a Signed-off-by: Thomas Kosiewski --- coderd/x/chatd/chatdebug/recorder_test.go | 2 + coderd/x/chatd/chatdebug/transport.go | 34 ++++- coderd/x/chatd/chatdebug/transport_test.go | 139 +++++++++++++++++++++ 3 files changed, 172 insertions(+), 3 deletions(-) diff --git a/coderd/x/chatd/chatdebug/recorder_test.go b/coderd/x/chatd/chatdebug/recorder_test.go index a2c3c88846a3d..1072db8515c3d 100644 --- a/coderd/x/chatd/chatdebug/recorder_test.go +++ b/coderd/x/chatd/chatdebug/recorder_test.go @@ -82,6 +82,8 @@ func TestNextStepNumber_Concurrent(t *testing.T) { const n = 256 runID := uuid.New() + t.Cleanup(func() { CleanupStepCounter(runID) }) + results := make([]int, n) var wg sync.WaitGroup wg.Add(n) diff --git a/coderd/x/chatd/chatdebug/transport.go b/coderd/x/chatd/chatdebug/transport.go index 6741aa8366840..c07f063dd0173 100644 --- a/coderd/x/chatd/chatdebug/transport.go +++ b/coderd/x/chatd/chatdebug/transport.go @@ -8,6 +8,7 @@ import ( "mime" "net/http" "net/url" + "regexp" "strings" "sync" "time" @@ -85,7 +86,7 @@ func (t *RecordingTransport) RoundTrip(req *http.Request) (*http.Response, error FinishedAt: finishedAt.UTC().Format(time.RFC3339Nano), RequestHeaders: requestHeaders, RequestBody: requestBody, - Error: err.Error(), + Error: sanitizeErrorString(err.Error()), DurationMs: durationMs, }) return nil, err @@ -113,6 +114,24 @@ func (t *RecordingTransport) RoundTrip(req *http.Request) (*http.Response, error return resp, nil } +// urlInErrorPattern matches URL-like substrings that transports or +// retry middleware may embed in error messages. Credentials can +// appear in userinfo or query parameters. +var urlInErrorPattern = regexp.MustCompile(`https?://[^\s"']+`) + +// sanitizeErrorString redacts URL-like substrings that may contain +// credentials (userinfo, query parameters) from transport error +// messages before they are persisted in debug attempts. +func sanitizeErrorString(errMsg string) string { + return urlInErrorPattern.ReplaceAllStringFunc(errMsg, func(rawURL string) string { + parsed, err := url.Parse(rawURL) + if err != nil { + return "[REDACTED_URL]" + } + return redactURL(parsed) + }) +} + func redactURL(u *url.URL) string { if u == nil { return "" @@ -362,17 +381,26 @@ func (r *recordingBody) record(err error) { startedAt := r.startedAt r.mu.Unlock() + contentType := base.ResponseHeaders["Content-Type"] if truncated { base.ResponseBody = []byte("[TRUNCATED]") - } else { + } else if contentType == "" || isJSONLikeContentType(contentType) { + // Redact JSON secrets when the content type is JSON-like + // or absent (unknown). For unknown types, RedactJSONSecrets + // fails closed by replacing non-JSON payloads with a + // diagnostic message. base.ResponseBody = RedactJSONSecrets(responseBody) + } else { + // Non-JSON content types (SSE, text/plain, HTML, etc.) + // are preserved as-is to avoid losing debug content. + base.ResponseBody = responseBody } base.StartedAt = startedAt.UTC().Format(time.RFC3339Nano) base.FinishedAt = finishedAt.UTC().Format(time.RFC3339Nano) // Recompute duration to include body read time. base.DurationMs = finishedAt.Sub(startedAt).Milliseconds() if err != nil && !errors.Is(err, io.EOF) { - base.Error = err.Error() + base.Error = sanitizeErrorString(err.Error()) base.Status = attemptStatusFailed } else { base.Status = attemptStatusCompleted diff --git a/coderd/x/chatd/chatdebug/transport_test.go b/coderd/x/chatd/chatdebug/transport_test.go index b1da8e24f2787..0cdb5a347b383 100644 --- a/coderd/x/chatd/chatdebug/transport_test.go +++ b/coderd/x/chatd/chatdebug/transport_test.go @@ -249,6 +249,7 @@ func TestRecordingTransport_CaptureResponse(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + rw.Header().Set("Content-Type", "application/json") rw.Header().Set("X-API-Key", "response-secret") rw.Header().Set("X-Trace-ID", "trace-123") rw.WriteHeader(http.StatusCreated) @@ -275,6 +276,7 @@ func TestRecordingTransport_CaptureResponse(t *testing.T) { attempts := sink.snapshot() require.Len(t, attempts, 1) require.Equal(t, http.StatusCreated, attempts[0].ResponseStatus) + require.Equal(t, "application/json", attempts[0].ResponseHeaders["Content-Type"]) require.Equal(t, RedactedValue, attempts[0].ResponseHeaders["X-Api-Key"]) require.Equal(t, "trace-123", attempts[0].ResponseHeaders["X-Trace-Id"]) require.JSONEq(t, `{"token":"[REDACTED]","safe":"ok"}`, string(attempts[0].ResponseBody)) @@ -714,6 +716,34 @@ func TestRecordingTransport_TransportError(t *testing.T) { require.GreaterOrEqual(t, attempts[0].DurationMs, int64(0)) } +func TestRecordingTransport_TransportErrorSanitizesURLCredentials(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return nil, xerrors.New("connection to http://admin:s3cret@api.example.com/v1?api_key=sk-1234 refused") + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + if resp != nil { + defer resp.Body.Close() + } + require.Error(t, err) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.NotContains(t, attempts[0].Error, "s3cret") + require.NotContains(t, attempts[0].Error, "sk-1234") + require.Contains(t, attempts[0].Error, "api_key=%5BREDACTED%5D") +} + func TestRecordingTransport_NilBase(t *testing.T) { t.Parallel() @@ -735,3 +765,112 @@ func TestRecordingTransport_NilBase(t *testing.T) { require.NoError(t, err) require.Equal(t, "ok", string(body)) } + +func TestRecordingTransport_SSEReadToEOFMarksCompleted(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + ssePayload := "data: {\"token\":\"secret\"}\n\ndata: [DONE]\n\n" + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test SSE content type. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(ssePayload)), + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + require.Equal(t, ssePayload, string(body)) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + require.Empty(t, attempts[0].Error) + // SSE bodies should be preserved as-is, not replaced with + // a redaction diagnostic. + require.Equal(t, ssePayload, string(attempts[0].ResponseBody)) +} + +func TestRecordingTransport_SSEClosedEarlyMarksFailed(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + ssePayload := "data: {\"token\":\"secret\"}\n\ndata: [DONE]\n\n" + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test SSE content type. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: &scriptedReadCloser{chunks: [][]byte{[]byte(ssePayload)}}, + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + // Read only a few bytes then close early. + buf := make([]byte, 5) + _, err = resp.Body.Read(buf) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusFailed, attempts[0].Status) + require.Equal(t, io.ErrUnexpectedEOF.Error(), attempts[0].Error) +} + +func TestRecordingTransport_TextPlainPreservedNotRedacted(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + textPayload := "This is plain text, not JSON." + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test text/plain content type. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/plain"}}, + Body: io.NopCloser(strings.NewReader(textPayload)), + ContentLength: int64(len(textPayload)), + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + // Non-JSON bodies should be preserved as-is, not replaced + // with a redaction diagnostic. + require.Equal(t, textPayload, string(attempts[0].ResponseBody)) +} From 90a4df56266e0d10bf5abe9364049a35894aa1e0 Mon Sep 17 00:00:00 2001 From: Thomas Kosiewski Date: Mon, 13 Apr 2026 19:47:07 +0200 Subject: [PATCH 3/3] fix(coderd/x/chatd/chatdebug): use case-insensitive Content-Type lookup Store content type from resp.Header.Get() (case-insensitive) in recordingBody instead of doing case-sensitive map lookup on the redacted headers. Prevents non-canonical header casing from bypassing JSON redaction. Change-Id: I3c9721be18536fe421eee26982f005e2a922bede Signed-off-by: Thomas Kosiewski --- coderd/x/chatd/chatdebug/recorder_test.go | 20 +- coderd/x/chatd/chatdebug/redaction.go | 28 ++ coderd/x/chatd/chatdebug/redaction_test.go | 44 ++ coderd/x/chatd/chatdebug/transport.go | 77 ++-- coderd/x/chatd/chatdebug/transport_test.go | 499 ++++++++++++++++++++- 5 files changed, 624 insertions(+), 44 deletions(-) diff --git a/coderd/x/chatd/chatdebug/recorder_test.go b/coderd/x/chatd/chatdebug/recorder_test.go index 1072db8515c3d..290757735450a 100644 --- a/coderd/x/chatd/chatdebug/recorder_test.go +++ b/coderd/x/chatd/chatdebug/recorder_test.go @@ -2,7 +2,7 @@ package chatdebug //nolint:testpackage // Uses unexported recorder helpers. import ( "context" - "sort" + "slices" "sync" "testing" @@ -20,13 +20,11 @@ func TestAttemptSink_ThreadSafe(t *testing.T) { sink := &attemptSink{} var wg sync.WaitGroup - wg.Add(n) for i := range n { - go func() { - defer wg.Done() + wg.Go(func() { sink.record(Attempt{Number: i + 1, ResponseStatus: 200 + i}) - }() + }) } wg.Wait() @@ -40,8 +38,8 @@ func TestAttemptSink_ThreadSafe(t *testing.T) { numbers = append(numbers, attempt.Number) statuses = append(statuses, attempt.ResponseStatus) } - sort.Ints(numbers) - sort.Ints(statuses) + slices.Sort(numbers) + slices.Sort(statuses) for i := range n { require.Equal(t, i+1, numbers[i]) @@ -86,18 +84,16 @@ func TestNextStepNumber_Concurrent(t *testing.T) { results := make([]int, n) var wg sync.WaitGroup - wg.Add(n) for i := range n { - go func() { - defer wg.Done() + wg.Go(func() { results[i] = int(nextStepNumber(runID)) - }() + }) } wg.Wait() - sort.Ints(results) + slices.Sort(results) for i := range n { require.Equal(t, i+1, results[i]) } diff --git a/coderd/x/chatd/chatdebug/redaction.go b/coderd/x/chatd/chatdebug/redaction.go index 784c5de65ca41..6ac8d6d71b661 100644 --- a/coderd/x/chatd/chatdebug/redaction.go +++ b/coderd/x/chatd/chatdebug/redaction.go @@ -122,6 +122,34 @@ func RedactJSONSecrets(data []byte) []byte { return encoded } +// RedactNDJSONSecrets redacts sensitive values in newline-delimited +// JSON (NDJSON) payloads. Each non-empty line is treated as an +// independent JSON document and redacted individually. Lines that +// fail to parse are replaced with a diagnostic placeholder. +func RedactNDJSONSecrets(data []byte) []byte { + if len(data) == 0 { + return data + } + + lines := bytes.Split(data, []byte("\n")) + changed := false + for i, line := range lines { + trimmed := bytes.TrimSpace(line) + if len(trimmed) == 0 { + continue + } + redacted := RedactJSONSecrets(trimmed) + if !bytes.Equal(redacted, trimmed) { + lines[i] = redacted + changed = true + } + } + if !changed { + return data + } + return bytes.Join(lines, []byte("\n")) +} + func consumeJSONEOF(decoder *json.Decoder) error { var extra any err := decoder.Decode(&extra) diff --git a/coderd/x/chatd/chatdebug/redaction_test.go b/coderd/x/chatd/chatdebug/redaction_test.go index ac45d53262443..c60d72a80242b 100644 --- a/coderd/x/chatd/chatdebug/redaction_test.go +++ b/coderd/x/chatd/chatdebug/redaction_test.go @@ -2,6 +2,7 @@ package chatdebug_test import ( "net/http" + "strings" "testing" "github.com/stretchr/testify/require" @@ -275,3 +276,46 @@ func TestRedactJSONSecrets(t *testing.T) { require.JSONEq(t, `{"accessToken":"[REDACTED]","refreshToken":"[REDACTED]","authToken":"[REDACTED]","input_tokens":100,"output_tokens":50}`, string(redacted)) }) } + +func TestRedactNDJSONSecrets(t *testing.T) { + t.Parallel() + + t.Run("empty input", func(t *testing.T) { + t.Parallel() + require.Empty(t, chatdebug.RedactNDJSONSecrets(nil)) + require.Empty(t, chatdebug.RedactNDJSONSecrets([]byte{})) + }) + + t.Run("redacts secrets in each line", func(t *testing.T) { + t.Parallel() + input := []byte("{\"api_key\":\"sk-123\",\"safe\":\"ok\"}\n{\"token\":\"tok-456\",\"data\":\"value\"}\n") + redacted := chatdebug.RedactNDJSONSecrets(input) + lines := strings.Split(string(redacted), "\n") + require.JSONEq(t, `{"api_key":"[REDACTED]","safe":"ok"}`, lines[0]) + require.JSONEq(t, `{"token":"[REDACTED]","data":"value"}`, lines[1]) + }) + + t.Run("preserves lines without secrets", func(t *testing.T) { + t.Parallel() + input := []byte("{\"safe\":\"ok\"}\n{\"data\":\"value\"}\n") + redacted := chatdebug.RedactNDJSONSecrets(input) + require.Equal(t, string(input), string(redacted)) + }) + + t.Run("handles malformed lines with fail-closed", func(t *testing.T) { + t.Parallel() + input := []byte("{\"safe\":\"ok\"}\nnot-json\n{\"token\":\"secret\"}\n") + redacted := chatdebug.RedactNDJSONSecrets(input) + lines := strings.Split(string(redacted), "\n") + require.JSONEq(t, `{"safe":"ok"}`, lines[0]) + require.Contains(t, lines[1], "not valid JSON") + require.JSONEq(t, `{"token":"[REDACTED]"}`, lines[2]) + }) + + t.Run("handles single line without trailing newline", func(t *testing.T) { + t.Parallel() + input := []byte(`{"api_key":"secret","value":"ok"}`) + redacted := chatdebug.RedactNDJSONSecrets(input) + require.JSONEq(t, `{"api_key":"[REDACTED]","value":"ok"}`, string(redacted)) + }) +} diff --git a/coderd/x/chatd/chatdebug/transport.go b/coderd/x/chatd/chatdebug/transport.go index c07f063dd0173..20a05f344e19a 100644 --- a/coderd/x/chatd/chatdebug/transport.go +++ b/coderd/x/chatd/chatdebug/transport.go @@ -98,6 +98,7 @@ func (t *RecordingTransport) RoundTrip(req *http.Request) (*http.Response, error sink: sink, startedAt: startedAt, contentLength: resp.ContentLength, + contentType: resp.Header.Get("Content-Type"), base: Attempt{ Number: attemptNumber, Method: method, @@ -180,6 +181,7 @@ func captureRequestBody(req *http.Request) ([]byte, error) { type recordingBody struct { inner io.ReadCloser contentLength int64 + contentType string // from resp.Header.Get (case-insensitive) sink *attemptSink base Attempt startedAt time.Time @@ -194,10 +196,9 @@ type recordingBody struct { closeOnce sync.Once } -func (r *recordingBody) Read(p []byte) (int, error) { - n, err := r.inner.Read(p) - - r.mu.Lock() +// accumulateReadLocked updates the buffer, byte counters, and +// truncation/EOF flags after a read. The caller must hold r.mu. +func (r *recordingBody) accumulateReadLocked(data []byte, n int, err error) { r.bytesRead += int64(n) if n > 0 && !r.truncated { remaining := maxRecordedResponseBodyBytes - r.buf.Len() @@ -207,7 +208,7 @@ func (r *recordingBody) Read(p []byte) (int, error) { toWrite = remaining r.truncated = true } - _, _ = r.buf.Write(p[:toWrite]) + _, _ = r.buf.Write(data[:toWrite]) } else { r.truncated = true } @@ -215,9 +216,20 @@ func (r *recordingBody) Read(p []byte) (int, error) { if errors.Is(err, io.EOF) { r.sawEOF = true } +} + +func (r *recordingBody) Read(p []byte) (int, error) { + n, err := r.inner.Read(p) + + r.mu.Lock() + r.accumulateReadLocked(p, n, err) r.mu.Unlock() - if err != nil { + // Only record non-EOF errors immediately. io.EOF is deferred + // to Close() which runs more sophisticated validation (JSON + // completeness checks, content-length verification, etc.). + // Recording EOF here would preempt Close() via recordOnce. + if err != nil && !errors.Is(err, io.EOF) { r.record(err) } return n, err @@ -232,7 +244,7 @@ func (r *recordingBody) Close() error { responseBody := append([]byte(nil), r.buf.Bytes()...) r.mu.Unlock() - contentType := r.base.ResponseHeaders["Content-Type"] + contentType := r.contentType shouldDrainUnknownLengthJSON := contentLength < 0 && !sawEOF && bytesRead > 0 && @@ -281,6 +293,12 @@ func (r *recordingBody) Close() error { r.record(nil) case contentLength < 0 && !truncated && isCompleteUnknownLengthJSONBody(contentType, responseBody): r.record(nil) + // Truncated unknown-length bodies: the caller consumed the + // response successfully but the recording buffer exceeded + // maxRecordedResponseBodyBytes. This is not a transport + // failure - mark as completed with the truncated capture. + case contentLength < 0 && truncated: + r.record(nil) default: r.record(io.ErrUnexpectedEOF) } @@ -296,14 +314,26 @@ func responseHasNoBody(method string, statusCode int) bool { (statusCode >= 100 && statusCode < 200) } -func isJSONLikeContentType(contentType string) bool { +// parseMediaType extracts the media type from a Content-Type header +// value, falling back to splitting on ";" when mime.ParseMediaType +// fails. +func parseMediaType(contentType string) string { mediaType, _, err := mime.ParseMediaType(contentType) if err != nil { - mediaType = strings.TrimSpace(strings.Split(contentType, ";")[0]) + mediaType = strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0])) } + return mediaType +} + +func isJSONLikeContentType(contentType string) bool { + mediaType := parseMediaType(contentType) return mediaType == "application/json" || strings.HasSuffix(mediaType, "+json") } +func isNDJSONContentType(contentType string) bool { + return parseMediaType(contentType) == "application/x-ndjson" +} + // maxDrainBytes caps how many trailing bytes drainToEOF will consume. // This prevents Close() from blocking indefinitely on a misbehaving // or extremely large chunked body. @@ -316,24 +346,8 @@ func (r *recordingBody) drainToEOF() error { n, err := r.inner.Read(buf) r.mu.Lock() - r.bytesRead += int64(n) + r.accumulateReadLocked(buf, n, err) drained += int64(n) - if n > 0 && !r.truncated { - remaining := maxRecordedResponseBodyBytes - r.buf.Len() - if remaining > 0 { - toWrite := n - if toWrite > remaining { - toWrite = remaining - r.truncated = true - } - _, _ = r.buf.Write(buf[:toWrite]) - } else { - r.truncated = true - } - } - if errors.Is(err, io.EOF) { - r.sawEOF = true - } r.mu.Unlock() if err != nil { @@ -381,16 +395,19 @@ func (r *recordingBody) record(err error) { startedAt := r.startedAt r.mu.Unlock() - contentType := base.ResponseHeaders["Content-Type"] - if truncated { + contentType := r.contentType + switch { + case truncated: base.ResponseBody = []byte("[TRUNCATED]") - } else if contentType == "" || isJSONLikeContentType(contentType) { + case isNDJSONContentType(contentType): + base.ResponseBody = RedactNDJSONSecrets(responseBody) + case contentType == "" || isJSONLikeContentType(contentType): // Redact JSON secrets when the content type is JSON-like // or absent (unknown). For unknown types, RedactJSONSecrets // fails closed by replacing non-JSON payloads with a // diagnostic message. base.ResponseBody = RedactJSONSecrets(responseBody) - } else { + default: // Non-JSON content types (SSE, text/plain, HTML, etc.) // are preserved as-is to avoid losing debug content. base.ResponseBody = responseBody diff --git a/coderd/x/chatd/chatdebug/transport_test.go b/coderd/x/chatd/chatdebug/transport_test.go index 0cdb5a347b383..06f2a29997218 100644 --- a/coderd/x/chatd/chatdebug/transport_test.go +++ b/coderd/x/chatd/chatdebug/transport_test.go @@ -124,6 +124,7 @@ func TestRecordingTransport_CaptureRequest(t *testing.T) { attempts := sink.snapshot() require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) require.Equal(t, 1, attempts[0].Number) require.Equal(t, RedactedValue, attempts[0].RequestHeaders["Authorization"]) require.Equal(t, "application/json", attempts[0].RequestHeaders["Content-Type"]) @@ -156,6 +157,7 @@ func TestRecordingTransport_RedactsSensitiveQueryParameters(t *testing.T) { attempts := sink.snapshot() require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) require.Contains(t, attempts[0].URL, "api_key=%5BREDACTED%5D") require.Contains(t, attempts[0].URL, "safe=ok") } @@ -184,6 +186,7 @@ func TestRecordingTransport_TruncatesLargeRequestBodies(t *testing.T) { attempts := sink.snapshot() require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) require.Equal(t, []byte("[TRUNCATED]"), attempts[0].RequestBody) } @@ -209,6 +212,7 @@ func TestRecordingTransport_StripsURLUserinfo(t *testing.T) { attempts := sink.snapshot() require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) require.NotContains(t, attempts[0].URL, "user:secret") require.Contains(t, attempts[0].URL, "api_key=%5BREDACTED%5D") } @@ -242,6 +246,7 @@ func TestRecordingTransport_SkipsNonReplayableRequestBodyCapture(t *testing.T) { require.JSONEq(t, requestBody, string(<-gotRequest)) attempts := sink.snapshot() require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) require.Nil(t, attempts[0].RequestBody) } @@ -275,6 +280,7 @@ func TestRecordingTransport_CaptureResponse(t *testing.T) { attempts := sink.snapshot() require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) require.Equal(t, http.StatusCreated, attempts[0].ResponseStatus) require.Equal(t, "application/json", attempts[0].ResponseHeaders["Content-Type"]) require.Equal(t, RedactedValue, attempts[0].ResponseHeaders["X-Api-Key"]) @@ -282,7 +288,11 @@ func TestRecordingTransport_CaptureResponse(t *testing.T) { require.JSONEq(t, `{"token":"[REDACTED]","safe":"ok"}`, string(attempts[0].ResponseBody)) } -func TestRecordingTransport_CaptureResponseOnEOFWithoutClose(t *testing.T) { +// TestRecordingTransport_CaptureResponseRecordsOnClose verifies that +// EOF recording is deferred to Close() rather than firing in Read(). +// This ensures Close()'s validation logic (JSON integrity, content- +// length checks) always runs. +func TestRecordingTransport_CaptureResponseRecordsOnClose(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { @@ -308,13 +318,19 @@ func TestRecordingTransport_CaptureResponseOnEOFWithoutClose(t *testing.T) { require.NoError(t, err) require.JSONEq(t, `{"token":"response-secret","safe":"ok"}`, string(body)) + // Before Close(), the attempt should not yet be recorded + // because EOF recording is deferred to Close(). + require.Empty(t, sink.snapshot(), "attempt should not be recorded before Close()") + + require.NoError(t, resp.Body.Close()) + attempts := sink.snapshot() require.Len(t, attempts, 1) require.Equal(t, http.StatusAccepted, attempts[0].ResponseStatus) require.Equal(t, "application/json", attempts[0].ResponseHeaders["Content-Type"]) require.Equal(t, RedactedValue, attempts[0].ResponseHeaders["X-Api-Key"]) require.JSONEq(t, `{"token":"[REDACTED]","safe":"ok"}`, string(attempts[0].ResponseBody)) - require.NoError(t, resp.Body.Close()) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) } func TestRecordingTransport_StreamingBody(t *testing.T) { @@ -361,6 +377,7 @@ func TestRecordingTransport_StreamingBody(t *testing.T) { attempts := sink.snapshot() require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) require.JSONEq(t, `{"safe":"stream","token":"[REDACTED]"}`, string(attempts[0].ResponseBody)) } @@ -651,6 +668,7 @@ func TestRecordingTransport_PrematureCloseMarksFailed(t *testing.T) { attempts := sink.snapshot() require.Len(t, attempts, 1) require.Equal(t, attemptStatusFailed, attempts[0].Status) + require.NotEmpty(t, attempts[0].Error, "failure-path attempt should record an Error") } func TestRecordingTransport_TruncatesLargeResponses(t *testing.T) { @@ -675,6 +693,7 @@ func TestRecordingTransport_TruncatesLargeResponses(t *testing.T) { attempts := sink.snapshot() require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) require.Equal(t, []byte("[TRUNCATED]"), attempts[0].ResponseBody) } @@ -708,6 +727,7 @@ func TestRecordingTransport_TransportError(t *testing.T) { attempts := sink.snapshot() require.Len(t, attempts, 1) + require.Equal(t, attemptStatusFailed, attempts[0].Status) require.Equal(t, 1, attempts[0].Number) require.Equal(t, RedactedValue, attempts[0].RequestHeaders["Authorization"]) require.JSONEq(t, `{"password":"[REDACTED]","safe":"ok"}`, string(attempts[0].RequestBody)) @@ -739,6 +759,7 @@ func TestRecordingTransport_TransportErrorSanitizesURLCredentials(t *testing.T) attempts := sink.snapshot() require.Len(t, attempts, 1) + require.Equal(t, attemptStatusFailed, attempts[0].Status) require.NotContains(t, attempts[0].Error, "s3cret") require.NotContains(t, attempts[0].Error, "sk-1234") require.Contains(t, attempts[0].Error, "api_key=%5BREDACTED%5D") @@ -874,3 +895,477 @@ func TestRecordingTransport_TextPlainPreservedNotRedacted(t *testing.T) { // with a redaction diagnostic. require.Equal(t, textPayload, string(attempts[0].ResponseBody)) } + +// TestRecordingTransport_NDJSONRedacted verifies that NDJSON response +// bodies have secrets redacted on a per-line basis rather than being +// treated as non-JSON and preserved raw. +func TestRecordingTransport_NDJSONRedacted(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + ndjsonPayload := "{\"api_key\":\"sk-123\",\"safe\":\"ok\"}\n{\"token\":\"tok-456\",\"data\":\"value\"}\n" + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test NDJSON content type. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/x-ndjson"}}, + Body: io.NopCloser(strings.NewReader(ndjsonPayload)), + ContentLength: int64(len(ndjsonPayload)), + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + // Caller sees original unredacted payload. + require.Equal(t, ndjsonPayload, string(body)) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + // Recorded body should have secrets redacted per-line. + lines := strings.Split(string(attempts[0].ResponseBody), "\n") + require.JSONEq(t, `{"api_key":"[REDACTED]","safe":"ok"}`, lines[0]) + require.JSONEq(t, `{"token":"[REDACTED]","data":"value"}`, lines[1]) +} + +// TestRecordingTransport_PlusJSONSuffixRedacted verifies that +// content types with a +json suffix (e.g. application/vnd.api+json) +// are treated as JSON-like and have secrets redacted in recorded +// response bodies. +func TestRecordingTransport_PlusJSONSuffixRedacted(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + jsonPayload := `{"token":"secret","safe":"ok"}` + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test +json suffix content type. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/vnd.api+json"}}, + Body: io.NopCloser(strings.NewReader(jsonPayload)), + ContentLength: int64(len(jsonPayload)), + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + // Caller sees original unredacted payload. + require.Equal(t, jsonPayload, string(body)) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + // Token must be redacted in the recorded body. + require.JSONEq(t, `{"token":"[REDACTED]","safe":"ok"}`, string(attempts[0].ResponseBody)) +} + +// TestRecordingTransport_UnrecognizedContentTypeDefaultsToJSONRedaction +// verifies that an unrecognized content-type header (e.g. non-canonical +// lowercase key not found by http.Header.Get) defaults to JSON +// redaction rather than falling into the raw-body preservation path. +func TestRecordingTransport_UnrecognizedContentTypeDefaultsToJSONRedaction(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + // Use lowercase header key to simulate non-canonical transport. + return &http.Response{ //nolint:exhaustruct // Test lowercase content-type. + StatusCode: http.StatusOK, + Header: http.Header{"content-type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"token":"secret","safe":"ok"}`)), + ContentLength: int64(len(`{"token":"secret","safe":"ok"}`)), + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + // The token should be redacted, not preserved raw or replaced + // with the fail-closed diagnostic. + require.JSONEq(t, `{"token":"[REDACTED]","safe":"ok"}`, string(attempts[0].ResponseBody)) +} + +// TestRecordingTransport_NonJSONBodyFailClosedRedaction verifies that +// when the Content-Type is empty (or JSON-like) but the response body +// is not valid JSON, RedactJSONSecrets' fail-closed behavior replaces +// the body with a diagnostic message rather than preserving the raw +// content which could contain credentials. +func TestRecordingTransport_NonJSONBodyFailClosedRedaction(t *testing.T) { + t.Parallel() + + htmlBody := `502 Bad Gateway` + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + // Empty Content-Type triggers the JSON-or-unknown + // branch in record(), which calls RedactJSONSecrets. + return &http.Response{ //nolint:exhaustruct // Test fail-closed redaction. + StatusCode: http.StatusBadGateway, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader(htmlBody)), + ContentLength: int64(len(htmlBody)), + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + // The caller sees the original body. + require.Equal(t, htmlBody, string(body)) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + // The recorded body must be the fail-closed diagnostic, not the + // raw HTML which could contain tokens or session data. + require.JSONEq(t, + `{"error":"chatdebug: body is not valid JSON, redacted for safety"}`, + string(attempts[0].ResponseBody)) +} + +// TestRecordingTransport_TruncatedUnknownLengthMarksCompleted verifies +// that an unknown-length (chunked) response that exceeds the recording +// buffer is marked as completed, not failed. The caller consumed the +// body successfully; we just couldn't buffer all of it. +func TestRecordingTransport_TruncatedUnknownLengthMarksCompleted(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + largeBody := strings.Repeat("x", maxRecordedResponseBodyBytes+1024) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test unknown-length body. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/octet-stream"}}, + Body: io.NopCloser(strings.NewReader(largeBody)), + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Len(t, body, maxRecordedResponseBodyBytes+1024) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + require.Empty(t, attempts[0].Error) + require.Equal(t, []byte("[TRUNCATED]"), attempts[0].ResponseBody) +} + +// errorAfterReadCloser returns data for the first N reads, then an error. +type errorAfterReadCloser struct { + data []byte + offset int + errAt int // byte offset at which to return the error + err error +} + +func (r *errorAfterReadCloser) Read(p []byte) (int, error) { + if r.offset >= r.errAt { + return 0, r.err + } + remaining := r.data[r.offset:] + if len(remaining) > len(p) { + remaining = remaining[:len(p)] + } + if r.offset+len(remaining) > r.errAt { + remaining = remaining[:r.errAt-r.offset] + } + n := copy(p, remaining) + r.offset += n + if r.offset >= r.errAt { + return n, r.err + } + return n, nil +} + +func (*errorAfterReadCloser) Close() error { + return nil +} + +// TestRecordingTransport_MidStreamReadError verifies that a non-EOF +// read error during body consumption is recorded immediately with +// "failed" status and the correct error message. +func TestRecordingTransport_MidStreamReadError(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test mid-stream error. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &errorAfterReadCloser{data: []byte(`{"key":"value"}`), errAt: 10, err: io.ErrUnexpectedEOF}, + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + _, err = io.ReadAll(resp.Body) + require.ErrorIs(t, err, io.ErrUnexpectedEOF) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusFailed, attempts[0].Status) + require.Equal(t, io.ErrUnexpectedEOF.Error(), attempts[0].Error) +} + +// trackingReadCloser wraps a reader and counts total bytes delivered +// via Read. Close always succeeds. +type trackingReadCloser struct { + inner io.Reader + bytesRead int64 + closed bool +} + +func (r *trackingReadCloser) Read(p []byte) (int, error) { + n, err := r.inner.Read(p) + r.bytesRead += int64(n) + return n, err +} + +func (r *trackingReadCloser) Close() error { + r.closed = true + return nil +} + +// failingCloseReader reads normally but returns an error on Close. +type failingCloseReader struct { + inner io.Reader + closeErr error +} + +func (r *failingCloseReader) Read(p []byte) (int, error) { + return r.inner.Read(p) +} + +func (r *failingCloseReader) Close() error { + return r.closeErr +} + +// TestRecordingTransport_MaxDrainBytesRespected verifies that +// drainToEOF stops after maxDrainBytes, preventing unbounded reads. +// The test uses a tracking reader to assert the byte cap. +func TestRecordingTransport_MaxDrainBytesRespected(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + + // Build a body where json.Decoder consumes the first JSON document + // but leaves trailing whitespace larger than maxDrainBytes. The + // drain path should stop after maxDrainBytes, not read everything. + jsonDoc := `{"safe":"ok"}` + // Trailing whitespace much larger than maxDrainBytes. The drain + // should consume at most maxDrainBytes of it. + trailing := strings.Repeat(" ", maxDrainBytes*2) + fullBody := jsonDoc + trailing + + tracker := &trackingReadCloser{inner: strings.NewReader(fullBody)} + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test maxDrainBytes. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: tracker, + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + var decoded map[string]string + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + require.Equal(t, "ok", decoded["safe"]) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + + // The key assertion: total bytes read through the tracker should + // be bounded. The json.Decoder reads the JSON doc (~13 bytes), + // then drainToEOF reads at most maxDrainBytes more. Without the + // cap, the full body (maxDrainBytes*2 + 13) would be consumed. + maxExpected := int64(len(jsonDoc)) + int64(maxDrainBytes) + 4096 // small buffer overhead + require.Less(t, tracker.bytesRead, int64(len(fullBody)), + "drain should NOT have consumed the entire body") + require.LessOrEqual(t, tracker.bytesRead, maxExpected, + "total bytes read should be bounded by maxDrainBytes") + require.True(t, tracker.closed, "inner body should be closed") +} + +// TestRecordingTransport_InnerCloseError verifies that an error from +// the inner body's Close() is recorded as a failed attempt and +// returned to the caller. +func TestRecordingTransport_InnerCloseError(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + closeErr := xerrors.New("connection reset by peer") + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test close error. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &failingCloseReader{inner: strings.NewReader(`{"ok":true}`), closeErr: closeErr}, + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + err = resp.Body.Close() + require.Error(t, err) + require.Contains(t, err.Error(), "connection reset by peer") + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusFailed, attempts[0].Status) + require.Contains(t, attempts[0].Error, "connection reset by peer") +} + +// TestRecordingTransport_204NoContentSucceeds verifies that a 204 No +// Content response is marked completed when closed without reading. +func TestRecordingTransport_204NoContentSucceeds(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test 204 no-body. + StatusCode: http.StatusNoContent, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader("")), + ContentLength: 0, + Request: req, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodDelete, "http://example.invalid/resource", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + require.Empty(t, attempts[0].Error) +} + +// TestRecordingTransport_304NotModifiedSucceeds verifies that a 304 +// Not Modified response is marked completed when closed without +// reading, even when Content-Length is non-zero. +func TestRecordingTransport_304NotModifiedSucceeds(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test 304 no-body. + StatusCode: http.StatusNotModified, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader("")), + ContentLength: 42, + Request: req, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid/resource", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + require.Empty(t, attempts[0].Error) +}