Move authentication probes into the outbound adapter
This commit is contained in:
105
internal/llm/openrouter_authentication.go
Normal file
105
internal/llm/openrouter_authentication.go
Normal file
@@ -0,0 +1,105 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
"go.maximumdirect.net/promptkit/internal/domain"
|
||||
)
|
||||
|
||||
const maxAuthenticationResponseBytes = 1 << 20
|
||||
|
||||
// AuthenticationEvidence describes the observable result of a key probe.
|
||||
type AuthenticationEvidence struct {
|
||||
HTTP HTTPExecutionEvidence
|
||||
Accepted *bool
|
||||
Overflow bool
|
||||
Unrecognized bool
|
||||
}
|
||||
|
||||
// OpenRouterKeyEndpoint derives the default key endpoint from a base endpoint.
|
||||
func OpenRouterKeyEndpoint(base string) (string, error) {
|
||||
return url.JoinPath(base, "key")
|
||||
}
|
||||
|
||||
// ProbeOpenRouterKey executes and interprets an OpenRouter key request.
|
||||
func (c *OpenAICompatibleClient) ProbeOpenRouterKey(ctx context.Context, endpoint string, timeoutSeconds int, target domain.ExecutionTarget) (AuthenticationEvidence, error) {
|
||||
parsed, err := url.Parse(endpoint)
|
||||
if err != nil {
|
||||
return AuthenticationEvidence{}, err
|
||||
}
|
||||
parsed.User = nil
|
||||
response, evidence, requestErr := c.GetWithAuthenticationEvidence(ctx, parsed.String(), timeoutSeconds, target)
|
||||
result := AuthenticationEvidence{HTTP: evidence}
|
||||
if response != nil && response.Body != nil {
|
||||
defer response.Body.Close()
|
||||
}
|
||||
if !evidence.ResponseReceived && response == nil {
|
||||
return result, requestErr
|
||||
}
|
||||
status := evidence.HTTPStatus
|
||||
if status == nil && response != nil {
|
||||
status = &response.StatusCode
|
||||
}
|
||||
if status == nil {
|
||||
result.Unrecognized = true
|
||||
return result, nil
|
||||
}
|
||||
switch *status {
|
||||
case http.StatusUnauthorized:
|
||||
accepted := false
|
||||
result.Accepted = &accepted
|
||||
case http.StatusOK:
|
||||
if response == nil || response.Body == nil {
|
||||
result.Unrecognized = true
|
||||
return result, nil
|
||||
}
|
||||
contents, err := io.ReadAll(io.LimitReader(response.Body, maxAuthenticationResponseBytes+1))
|
||||
if len(contents) > maxAuthenticationResponseBytes {
|
||||
result.Overflow = true
|
||||
return result, nil
|
||||
}
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
if recognizedOpenRouterKeyResponse(contents) {
|
||||
accepted := true
|
||||
result.Accepted = &accepted
|
||||
} else {
|
||||
result.Unrecognized = true
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func recognizedOpenRouterKeyResponse(contents []byte) bool {
|
||||
var envelope map[string]json.RawMessage
|
||||
if err := json.Unmarshal(contents, &envelope); err != nil || envelope == nil {
|
||||
return false
|
||||
}
|
||||
var data map[string]json.RawMessage
|
||||
if err := json.Unmarshal(envelope["data"], &data); err != nil || data == nil {
|
||||
return false
|
||||
}
|
||||
label := bytes.TrimSpace(data["label"])
|
||||
var labelValue string
|
||||
if len(label) == 0 || label[0] != '"' || json.Unmarshal(label, &labelValue) != nil {
|
||||
return false
|
||||
}
|
||||
freeTier := bytes.TrimSpace(data["is_free_tier"])
|
||||
if !bytes.Equal(freeTier, []byte("true")) && !bytes.Equal(freeTier, []byte("false")) {
|
||||
return false
|
||||
}
|
||||
decoder := json.NewDecoder(bytes.NewReader(data["usage"]))
|
||||
decoder.UseNumber()
|
||||
var number any
|
||||
if decoder.Decode(&number) != nil {
|
||||
return false
|
||||
}
|
||||
_, ok := number.(json.Number)
|
||||
return ok
|
||||
}
|
||||
146
internal/llm/openrouter_authentication_test.go
Normal file
146
internal/llm/openrouter_authentication_test.go
Normal file
@@ -0,0 +1,146 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"go.maximumdirect.net/promptkit/internal/domain"
|
||||
)
|
||||
|
||||
func TestProbeOpenRouterKey(t *testing.T) {
|
||||
valid := `{"data":{"label":"x","is_free_tier":false,"usage":0}}`
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
status int
|
||||
body string
|
||||
key string
|
||||
accepted *bool
|
||||
unrecognized, overflow bool
|
||||
}{
|
||||
{name: "recognized", status: 200, body: valid, key: "secret", accepted: boolValue(true)},
|
||||
{name: "exact limit", status: 200, body: valid + strings.Repeat(" ", maxAuthenticationResponseBytes-len(valid)), accepted: boolValue(true)},
|
||||
{name: "unknown", status: 200, body: `{"data":{"label":"x","is_free_tier":false,"usage":"0"}}`, unrecognized: true},
|
||||
{name: "trailing value", status: 200, body: valid + ` {}`, unrecognized: true},
|
||||
{name: "overflow", status: 200, body: strings.Repeat("x", maxAuthenticationResponseBytes+1), overflow: true},
|
||||
{name: "unauthorized anonymous", status: 401, accepted: boolValue(false)},
|
||||
{name: "other status", status: 503},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var authorization string
|
||||
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
if req.URL.User != nil {
|
||||
t.Fatal("request retained URL credentials")
|
||||
}
|
||||
authorization = req.Header.Get("Authorization")
|
||||
return &http.Response{StatusCode: tc.status, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(tc.body))}, nil
|
||||
})}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := client.ProbeOpenRouterKey(context.Background(), "https://user:password@example.invalid/key", 0, domain.ExecutionTarget{APIKey: tc.key})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.HTTP.Endpoint != "https://example.invalid/key" || got.HTTP.HTTPStatus == nil || *got.HTTP.HTTPStatus != tc.status || got.Unrecognized != tc.unrecognized || got.Overflow != tc.overflow {
|
||||
t.Fatalf("evidence = %+v", got)
|
||||
}
|
||||
if (got.Accepted == nil) != (tc.accepted == nil) || got.Accepted != nil && *got.Accepted != *tc.accepted {
|
||||
t.Fatalf("acceptance = %v", got.Accepted)
|
||||
}
|
||||
if (authorization != "") != (tc.key != "") || got.HTTP.CredentialsSupplied != (tc.key != "") {
|
||||
t.Fatalf("credential evidence = %+v", got.HTTP)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func boolValue(value bool) *bool { return &value }
|
||||
|
||||
type authenticationFailingBody struct{ err error }
|
||||
|
||||
func (b authenticationFailingBody) Read([]byte) (int, error) { return 0, b.err }
|
||||
func (b authenticationFailingBody) Close() error { return nil }
|
||||
|
||||
func TestProbeOpenRouterKeyRetainsBodyFailure(t *testing.T) {
|
||||
want := errors.New("interrupted")
|
||||
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
return &http.Response{StatusCode: 200, Header: make(http.Header), Body: authenticationFailingBody{err: want}}, nil
|
||||
})}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := client.ProbeOpenRouterKey(context.Background(), "https://example.invalid/key", 0, domain.ExecutionTarget{})
|
||||
if !errors.Is(err, want) || !got.HTTP.ResponseReceived {
|
||||
t.Fatalf("evidence = %+v, error = %v", got, err)
|
||||
}
|
||||
}
|
||||
|
||||
type observedAuthenticationBody struct {
|
||||
reader io.Reader
|
||||
read int
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (b *observedAuthenticationBody) Read(p []byte) (int, error) {
|
||||
n, err := b.reader.Read(p)
|
||||
b.read += n
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (b *observedAuthenticationBody) Close() error {
|
||||
b.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestProbeOpenRouterKeyBoundsAndClosesBodies(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
status int
|
||||
wantRead int
|
||||
}{
|
||||
{"overflow", http.StatusOK, maxAuthenticationResponseBytes + 1},
|
||||
{"unauthorized", http.StatusUnauthorized, 0},
|
||||
{"other status", http.StatusServiceUnavailable, 0},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
body := &observedAuthenticationBody{reader: strings.NewReader(strings.Repeat("x", 2*maxAuthenticationResponseBytes))}
|
||||
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{HTTPClient: &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
return &http.Response{StatusCode: tc.status, ContentLength: 1, Header: make(http.Header), Body: body}, nil
|
||||
})}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := client.ProbeOpenRouterKey(context.Background(), "https://example.invalid/key", 0, domain.ExecutionTarget{})
|
||||
if err != nil || body.read != tc.wantRead || !body.closed || got.Overflow != (tc.status == http.StatusOK) {
|
||||
t.Fatalf("evidence = %+v, read = %d, closed = %v, error = %v", got, body.read, body.closed, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenRouterKeyEndpoint(t *testing.T) {
|
||||
got, err := OpenRouterKeyEndpoint("https://example.invalid/api/v1/")
|
||||
if err != nil || got != "https://example.invalid/api/v1/key" {
|
||||
t.Fatalf("endpoint = %q, error = %v", got, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbeOpenRouterKeyPreservesCancellation(t *testing.T) {
|
||||
client, err := NewOpenAICompatibleClient(OpenAICompatibleConfig{HTTPClient: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
<-req.Context().Done()
|
||||
return nil, req.Context().Err()
|
||||
})}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
got, err := client.ProbeOpenRouterKey(ctx, "https://example.invalid/key", 0, domain.ExecutionTarget{})
|
||||
if !errors.Is(err, context.Canceled) || got.Accepted != nil || got.HTTP.ResponseReceived {
|
||||
t.Fatalf("evidence = %+v, error = %v", got, err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user