Bound authentication response recognition

internal/domain/smoke_diagnostics.go

@@ -112,12 +112,13 @@ type SmokeTestAttempt struct {
 type SmokeTestReason string
 
 const (
-	SmokeTestReasonNone                               SmokeTestReason = ""
-	SmokeTestReasonAuthenticationNotApplicable        SmokeTestReason = "authentication_not_applicable"
-	SmokeTestReasonProbeNotConfigured                 SmokeTestReason = "probe_not_configured"
-	SmokeTestReasonHTTPDiagnosticsUnavailable         SmokeTestReason = "http_diagnostics_unavailable"
-	SmokeTestReasonRequiredCredentialMissing          SmokeTestReason = "required_credential_missing"
-	SmokeTestReasonAuthenticationResponseUnrecognized SmokeTestReason = "authentication_response_unrecognized"
+	SmokeTestReasonNone                                SmokeTestReason = ""
+	SmokeTestReasonAuthenticationNotApplicable         SmokeTestReason = "authentication_not_applicable"
+	SmokeTestReasonProbeNotConfigured                  SmokeTestReason = "probe_not_configured"
+	SmokeTestReasonHTTPDiagnosticsUnavailable          SmokeTestReason = "http_diagnostics_unavailable"
+	SmokeTestReasonRequiredCredentialMissing           SmokeTestReason = "required_credential_missing"
+	SmokeTestReasonAuthenticationResponseUnrecognized  SmokeTestReason = "authentication_response_unrecognized"
+	SmokeTestReasonAuthenticationResponseLimitExceeded SmokeTestReason = "authentication_response_limit_exceeded"
 )
 
 // CloneBackend returns a copy of a backend, including its mutable settings.

internal/usecase/smoke_authentication.go

@@ -17,6 +17,8 @@ type authenticationProbe struct {
 	client smokeHTTPGetter
 }
 
+const maxAuthenticationResponseBytes = 1 << 20
+
 type authenticationProbeOutcome struct {
 	supported           bool
 	applicable          *bool
@@ -104,7 +106,17 @@ func (p authenticationProbe) runOpenRouterKey(
 				outcome.failure = domain.ExecutionFailureMalformedResponse
 				return outcome, nil
 			}
-			recognized, bodyErr := recognizedOpenRouterKeyResponse(response.Body)
+			contents, bodyErr := io.ReadAll(io.LimitReader(response.Body, maxAuthenticationResponseBytes+1))
+			if len(contents) > maxAuthenticationResponseBytes {
+				outcome.reason = domain.SmokeTestReasonAuthenticationResponseLimitExceeded
+				outcome.failure = domain.ExecutionFailureOther
+				return outcome, nil
+			}
+			if bodyErr != nil {
+				outcome.failure = httpProbeFailure(ctx, bodyErr)
+				return outcome, nil
+			}
+			recognized, bodyErr := recognizedOpenRouterKeyResponse(bytes.NewReader(contents))
 			if bodyErr != nil {
 				outcome.failure = httpProbeFailure(ctx, bodyErr)
 				return outcome, nil

internal/usecase/smoke_authentication_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"strings"
 	"testing"
@@ -11,6 +12,57 @@ import (
 	"go.maximumdirect.net/promptkit/internal/domain"
 )
 
+type measuredAuthenticationBody struct {
+	reader io.Reader
+	read   int
+	closed bool
+}
+
+func (b *measuredAuthenticationBody) Read(p []byte) (int, error) {
+	n, err := b.reader.Read(p)
+	b.read += n
+	return n, err
+}
+func (b *measuredAuthenticationBody) Close() error { b.closed = true; return nil }
+
+func TestAuthenticationResponseRecognitionBound(t *testing.T) {
+	valid := `{"data":{"label":"member","is_free_tier":false,"usage":0}}`
+	for _, test := range []struct {
+		name     string
+		size     int
+		declared int64
+		status   int
+		overflow bool
+	}{
+		{"exact limit", maxAuthenticationResponseBytes, -1, 200, false},
+		{"overflow unknown length", maxAuthenticationResponseBytes + 10, -1, 200, true},
+		{"overflow understated length", maxAuthenticationResponseBytes + 10, 1, 200, true},
+		{"unauthorized does not read", maxAuthenticationResponseBytes + 10, -1, 401, false},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			body := &measuredAuthenticationBody{reader: strings.NewReader(valid + strings.Repeat(" ", test.size-len(valid)))}
+			client := reachabilityClient(t, func(*http.Request) (*http.Response, error) {
+				return &http.Response{StatusCode: test.status, Header: make(http.Header), Body: body, ContentLength: test.declared}, nil
+			})
+			outcome, err := (authenticationProbe{client: client}).run(context.Background(), backendWithAuthenticationKind("openrouter_key"), domain.SmokeTestTarget{Endpoint: "https://example.invalid/key"})
+			if err != nil || !body.closed {
+				t.Fatalf("outcome = %+v, err = %v, closed = %v", outcome, err, body.closed)
+			}
+			if test.overflow {
+				if outcome.reason != domain.SmokeTestReasonAuthenticationResponseLimitExceeded || outcome.failure != domain.ExecutionFailureOther || outcome.credentialsAccepted != nil || body.read != maxAuthenticationResponseBytes+1 {
+					t.Fatalf("overflow = %+v, read = %d", outcome, body.read)
+				}
+			} else if test.status == 401 {
+				if body.read != 0 || outcome.credentialsAccepted == nil || *outcome.credentialsAccepted {
+					t.Fatalf("unauthorized = %+v, read = %d", outcome, body.read)
+				}
+			} else if outcome.credentialsAccepted == nil || !*outcome.credentialsAccepted {
+				t.Fatalf("exact limit = %+v", outcome)
+			}
+		})
+	}
+}
+
 func TestAuthenticationProbeApplicabilityAndHTTPAvailability(t *testing.T) {
 	client := reachabilityClient(t, func(*http.Request) (*http.Response, error) {
 		t.Fatal("authentication applicability unexpectedly dispatched")