[INT-715] Retry transient failures when verifying OpenAI keys (#5117)
* Retry transient failures when verifying OpenAI keys The OpenAI verification endpoint can be slow to respond under load, and a single 5s client timeout was recorded as an indeterminate verification result with no retry. Switch the default client to the common retryable client with a 10s per-attempt timeout and up to 2 retries with backoff on timeouts, connection errors, and 429/5xx responses. * Check error when closing response body in test * Add test for exhausted retry budget
This commit is contained in:
@@ -25,7 +25,14 @@ type Scanner struct {
|
||||
var _ detectors.Detector = (*Scanner)(nil)
|
||||
|
||||
var (
|
||||
defaultClient = common.SaneHttpClient()
|
||||
// The OpenAI API can be slow to respond under load, so use a longer
|
||||
// per-attempt timeout than the default 5s and retry transient failures
|
||||
// (timeouts, connection errors, 429/5xx) so a single slow response does
|
||||
// not record an indeterminate verification result.
|
||||
defaultClient = common.RetryableHTTPClient(
|
||||
common.WithTimeout(10*time.Second),
|
||||
common.WithMaxRetries(2),
|
||||
)
|
||||
|
||||
// The magic string T3BlbkFJ is the base64-encoded string: OpenAI
|
||||
// Matches: legacy keys (sk-{alnum}T3BlbkFJ...), project keys (sk-proj-...),
|
||||
|
||||
@@ -2,12 +2,76 @@ package openai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
|
||||
"github.com/trufflesecurity/trufflehog/v3/pkg/engine/ahocorasick"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The default client must retry transient failures so a single slow or
|
||||
// failed OpenAI API response does not record an indeterminate verification
|
||||
// result (CSM-2131).
|
||||
func TestOpenAI_DefaultClientRetriesTransientErrors(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if requests.Add(1) < 3 {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, ts.URL, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
res, err := defaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
defer func() { _ = res.Body.Close() }()
|
||||
|
||||
if res.StatusCode != http.StatusOK {
|
||||
t.Errorf("expected retries to reach a 200 response, got %d", res.StatusCode)
|
||||
}
|
||||
if got := requests.Load(); got != 3 {
|
||||
t.Errorf("expected 3 attempts (initial + 2 retries), got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
// When the API never recovers, the client must give up after the configured
|
||||
// retry budget (initial attempt + 2 retries) and surface the failure rather
|
||||
// than retrying indefinitely.
|
||||
func TestOpenAI_DefaultClientGivesUpAfterRetryBudget(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requests.Add(1)
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, ts.URL, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
res, err := defaultClient.Do(req)
|
||||
if res != nil {
|
||||
defer func() { _ = res.Body.Close() }()
|
||||
}
|
||||
|
||||
if err == nil && res.StatusCode != http.StatusInternalServerError {
|
||||
t.Errorf("expected exhausted retries to surface the failure, got status %d with no error", res.StatusCode)
|
||||
}
|
||||
if got := requests.Load(); got != 3 {
|
||||
t.Errorf("expected 3 attempts (initial + 2 retries), got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAI_DoesNotMatchAdminKeys(t *testing.T) {
|
||||
d := Scanner{}
|
||||
adminKey := `OPENAI_ADMIN_KEY = "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"`
|
||||
|
||||
Reference in New Issue
Block a user