[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 _ detectors.Detector = (*Scanner)(nil)
|
||||||
|
|
||||||
var (
|
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
|
// The magic string T3BlbkFJ is the base64-encoded string: OpenAI
|
||||||
// Matches: legacy keys (sk-{alnum}T3BlbkFJ...), project keys (sk-proj-...),
|
// Matches: legacy keys (sk-{alnum}T3BlbkFJ...), project keys (sk-proj-...),
|
||||||
|
|||||||
@@ -2,12 +2,76 @@ package openai
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
|
||||||
"github.com/google/go-cmp/cmp"
|
"github.com/google/go-cmp/cmp"
|
||||||
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
|
"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
|
||||||
"github.com/trufflesecurity/trufflehog/v3/pkg/engine/ahocorasick"
|
"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) {
|
func TestOpenAI_DoesNotMatchAdminKeys(t *testing.T) {
|
||||||
d := Scanner{}
|
d := Scanner{}
|
||||||
adminKey := `OPENAI_ADMIN_KEY = "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"`
|
adminKey := `OPENAI_ADMIN_KEY = "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"`
|
||||||
|
|||||||
Reference in New Issue
Block a user