398 lines
14 KiB
Go
398 lines
14 KiB
Go
package auth
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net/http"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/home"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
|
|
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
|
)
|
|
|
|
func TestPickNextViaHomeDoesNotReusePinnedWebsocketAuthWithoutSelection(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.RegisterExecutor(schedulerTestExecutor{})
|
|
|
|
auth := &Auth{
|
|
ID: "home-auth-1",
|
|
Provider: "test",
|
|
Status: StatusActive,
|
|
Attributes: map[string]string{
|
|
"websockets": "true",
|
|
homeUpstreamModelAttributeKey: "upstream-model",
|
|
},
|
|
Metadata: map[string]any{"email": "home@example.com"},
|
|
}
|
|
auth.EnsureIndex()
|
|
manager.rememberHomeRuntimeAuth("session-1", auth)
|
|
cachedAuth, ok := manager.GetExecutionSessionAuthByID("session-1", "home-auth-1")
|
|
if !ok || cachedAuth == nil || !authWebsocketsEnabled(cachedAuth) {
|
|
t.Fatalf("GetExecutionSessionAuthByID() did not expose remembered websocket home auth: auth=%#v ok=%v", cachedAuth, ok)
|
|
}
|
|
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
opts := cliproxyexecutor.Options{
|
|
Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "session-1",
|
|
cliproxyexecutor.PinnedAuthMetadataKey: "home-auth-1",
|
|
},
|
|
Headers: http.Header{"Authorization": {"Bearer client-key"}},
|
|
}
|
|
|
|
got, executor, provider, errPick := manager.pickNextViaHome(ctx, "gpt-5.4", opts, nil)
|
|
if errPick == nil {
|
|
t.Fatal("pickNextViaHome() unexpectedly reused an auth without a Home selection")
|
|
}
|
|
if got != nil || executor != nil || provider != "" {
|
|
t.Fatalf("pickNextViaHome() returned unbound execution target: auth=%#v executor=%#v provider=%q", got, executor, provider)
|
|
}
|
|
}
|
|
|
|
func TestPickNextViaHomeRejectsSessionScopedAuthCache(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.RegisterExecutor(schedulerTestExecutor{})
|
|
|
|
manager.rememberHomeRuntimeAuth("session-1", &Auth{
|
|
ID: "home-auth-1",
|
|
Provider: "test",
|
|
Status: StatusActive,
|
|
Attributes: map[string]string{
|
|
"websockets": "true",
|
|
homeUpstreamModelAttributeKey: "upstream-model-a",
|
|
},
|
|
})
|
|
manager.rememberHomeRuntimeAuth("session-2", &Auth{
|
|
ID: "home-auth-1",
|
|
Provider: "test",
|
|
Status: StatusActive,
|
|
Attributes: map[string]string{
|
|
"websockets": "true",
|
|
homeUpstreamModelAttributeKey: "upstream-model-b",
|
|
},
|
|
})
|
|
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
optsSession1 := cliproxyexecutor.Options{
|
|
Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "session-1",
|
|
cliproxyexecutor.PinnedAuthMetadataKey: "home-auth-1",
|
|
},
|
|
}
|
|
optsSession2 := cliproxyexecutor.Options{
|
|
Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "session-2",
|
|
cliproxyexecutor.PinnedAuthMetadataKey: "home-auth-1",
|
|
},
|
|
}
|
|
|
|
if _, _, _, errSession1 := manager.pickNextViaHome(ctx, "gpt-5.4", optsSession1, nil); errSession1 == nil {
|
|
t.Fatal("pickNextViaHome(session-1) unexpectedly reused a session auth cache")
|
|
}
|
|
if _, _, _, errSession2 := manager.pickNextViaHome(ctx, "gpt-5.4", optsSession2, nil); errSession2 == nil {
|
|
t.Fatal("pickNextViaHome(session-2) unexpectedly reused a session auth cache")
|
|
}
|
|
}
|
|
|
|
func TestPickNextViaHomeDoesNotReuseTriedPinnedWebsocketAuth(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.RegisterExecutor(schedulerTestExecutor{})
|
|
|
|
auth := &Auth{
|
|
ID: "home-auth-1",
|
|
Provider: "test",
|
|
Status: StatusActive,
|
|
Attributes: map[string]string{
|
|
"websockets": "true",
|
|
},
|
|
}
|
|
manager.rememberHomeRuntimeAuth("session-1", auth)
|
|
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
opts := cliproxyexecutor.Options{
|
|
Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "session-1",
|
|
cliproxyexecutor.PinnedAuthMetadataKey: "home-auth-1",
|
|
},
|
|
}
|
|
tried := map[string]struct{}{"home-auth-1": {}}
|
|
|
|
got, executor, provider, errPick := manager.pickNextViaHome(ctx, "gpt-5.4", opts, tried)
|
|
if errPick == nil {
|
|
t.Fatal("pickNextViaHome() error is nil, want home unavailable error")
|
|
}
|
|
var authErr *Error
|
|
if !errors.As(errPick, &authErr) || authErr.Code != "home_unavailable" {
|
|
t.Fatalf("pickNextViaHome() error = %v, want home_unavailable", errPick)
|
|
}
|
|
if got != nil || executor != nil || provider != "" {
|
|
t.Fatalf("pickNextViaHome() reused tried auth: auth=%#v executor=%#v provider=%q", got, executor, provider)
|
|
}
|
|
}
|
|
|
|
func TestPickNextViaHomeDoesNotReusePinnedWebsocketAuthAfterFirstHomeAttempt(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.RegisterExecutor(schedulerTestExecutor{})
|
|
|
|
auth := &Auth{
|
|
ID: "home-auth-1",
|
|
Provider: "test",
|
|
Status: StatusActive,
|
|
Attributes: map[string]string{
|
|
"websockets": "true",
|
|
},
|
|
}
|
|
manager.rememberHomeRuntimeAuth("session-1", auth)
|
|
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
opts := withHomeAuthCount(cliproxyexecutor.Options{
|
|
Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "session-1",
|
|
cliproxyexecutor.PinnedAuthMetadataKey: "home-auth-1",
|
|
},
|
|
}, 2)
|
|
|
|
got, executor, provider, errPick := manager.pickNextViaHome(ctx, "gpt-5.4", opts, nil)
|
|
if errPick == nil {
|
|
t.Fatal("pickNextViaHome() error is nil, want home unavailable error")
|
|
}
|
|
var authErr *Error
|
|
if !errors.As(errPick, &authErr) || authErr.Code != "home_unavailable" {
|
|
t.Fatalf("pickNextViaHome() error = %v, want home_unavailable", errPick)
|
|
}
|
|
if got != nil || executor != nil || provider != "" {
|
|
t.Fatalf("pickNextViaHome() reused auth after first home attempt: auth=%#v executor=%#v provider=%q", got, executor, provider)
|
|
}
|
|
}
|
|
|
|
func TestPickNextViaHomeDoesNotReusePinnedNonWebsocketAuth(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.RegisterExecutor(schedulerTestExecutor{})
|
|
|
|
manager.mu.Lock()
|
|
manager.homeRuntimeAuths["session-1"] = map[string]*Auth{
|
|
"home-auth-1": &Auth{
|
|
ID: "home-auth-1",
|
|
Provider: "test",
|
|
Status: StatusActive,
|
|
},
|
|
}
|
|
manager.mu.Unlock()
|
|
|
|
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
|
opts := cliproxyexecutor.Options{
|
|
Metadata: map[string]any{
|
|
cliproxyexecutor.ExecutionSessionMetadataKey: "session-1",
|
|
cliproxyexecutor.PinnedAuthMetadataKey: "home-auth-1",
|
|
},
|
|
Headers: http.Header{"Authorization": {"Bearer client-key"}},
|
|
}
|
|
|
|
got, executor, provider, errPick := manager.pickNextViaHome(ctx, "gpt-5.4", opts, nil)
|
|
if errPick == nil {
|
|
t.Fatal("pickNextViaHome() error is nil, want home unavailable error")
|
|
}
|
|
var authErr *Error
|
|
if !errors.As(errPick, &authErr) || authErr.Code != "home_unavailable" {
|
|
t.Fatalf("pickNextViaHome() error = %v, want home_unavailable", errPick)
|
|
}
|
|
if got != nil || executor != nil || provider != "" {
|
|
t.Fatalf("pickNextViaHome() reused non-websocket auth: auth=%#v executor=%#v provider=%q", got, executor, provider)
|
|
}
|
|
}
|
|
|
|
type homeAuthTransportErrorDispatcher struct {
|
|
err error
|
|
aborts atomic.Int32
|
|
onAbort func()
|
|
}
|
|
|
|
func (d *homeAuthTransportErrorDispatcher) HeartbeatOK() bool {
|
|
return true
|
|
}
|
|
|
|
func (d *homeAuthTransportErrorDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) {
|
|
return nil, d.err
|
|
}
|
|
|
|
func (d *homeAuthTransportErrorDispatcher) AbortAmbiguousDispatch() {
|
|
d.aborts.Add(1)
|
|
if d.onAbort != nil {
|
|
d.onAbort()
|
|
}
|
|
}
|
|
|
|
func TestPickNextViaHomeClassifiesTransportErrorsAsHomeUnavailable(t *testing.T) {
|
|
dispatcher := &homeAuthTransportErrorDispatcher{err: errors.New("read tcp 127.0.0.1:46704->127.0.0.1:8327: i/o timeout")}
|
|
oldCurrentHomeDispatcher := currentHomeDispatcher
|
|
currentHomeDispatcher = func() homeAuthDispatcher {
|
|
return dispatcher
|
|
}
|
|
t.Cleanup(func() {
|
|
currentHomeDispatcher = oldCurrentHomeDispatcher
|
|
})
|
|
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.SetHomeExecutionRegistry(executionregistry.New())
|
|
|
|
_, _, _, errPick := manager.pickNextViaHome(context.Background(), "gpt-5.4", cliproxyexecutor.Options{}, nil)
|
|
if errPick == nil {
|
|
t.Fatal("pickNextViaHome() error is nil, want home unavailable error")
|
|
}
|
|
var authErr *Error
|
|
if !errors.As(errPick, &authErr) {
|
|
t.Fatalf("pickNextViaHome() error = %T, want *Error", errPick)
|
|
}
|
|
if authErr.Code != "home_unavailable" {
|
|
t.Fatalf("pickNextViaHome() error code = %q, want home_unavailable (%v)", authErr.Code, errPick)
|
|
}
|
|
if authErr.StatusCode() != http.StatusServiceUnavailable {
|
|
t.Fatalf("pickNextViaHome() status = %d, want %d", authErr.StatusCode(), http.StatusServiceUnavailable)
|
|
}
|
|
if !authErr.Retryable {
|
|
t.Fatal("pickNextViaHome() retryable = false, want true")
|
|
}
|
|
}
|
|
|
|
func TestPickNextViaHomeAbortsBeforeEndingPendingDispatch(t *testing.T) {
|
|
registry := executionregistry.New()
|
|
abortSawPending := make(chan bool, 1)
|
|
dispatcher := &homeAuthTransportErrorDispatcher{
|
|
err: home.NewAmbiguousDispatchError(errors.New("response connection closed")),
|
|
onAbort: func() {
|
|
cancelledCtx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
abortSawPending <- errors.Is(registry.Drain(cancelledCtx), context.Canceled)
|
|
},
|
|
}
|
|
oldCurrentHomeDispatcher := currentHomeDispatcher
|
|
currentHomeDispatcher = func() homeAuthDispatcher {
|
|
return dispatcher
|
|
}
|
|
t.Cleanup(func() {
|
|
currentHomeDispatcher = oldCurrentHomeDispatcher
|
|
})
|
|
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.SetHomeExecutionRegistry(registry)
|
|
|
|
_, _, _, errPick := manager.pickNextViaHome(context.Background(), "gpt-5.4", cliproxyexecutor.Options{}, nil)
|
|
if errPick == nil {
|
|
t.Fatal("pickNextViaHome() error = nil, want home unavailable")
|
|
}
|
|
if sawPending := <-abortSawPending; !sawPending {
|
|
t.Fatal("AbortAmbiguousDispatch() observed an already-ended pending dispatch")
|
|
}
|
|
}
|
|
|
|
func TestPickNextViaHomeDoesNotAbortDeterministicDispatchFailure(t *testing.T) {
|
|
dispatcher := &homeAuthTransportErrorDispatcher{err: home.ErrNotConnected}
|
|
oldCurrentHomeDispatcher := currentHomeDispatcher
|
|
currentHomeDispatcher = func() homeAuthDispatcher {
|
|
return dispatcher
|
|
}
|
|
t.Cleanup(func() {
|
|
currentHomeDispatcher = oldCurrentHomeDispatcher
|
|
})
|
|
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.SetHomeExecutionRegistry(executionregistry.New())
|
|
|
|
_, _, _, errPick := manager.pickNextViaHome(context.Background(), "gpt-5.4", cliproxyexecutor.Options{}, nil)
|
|
if errPick == nil {
|
|
t.Fatal("pickNextViaHome() error = nil, want home unavailable")
|
|
}
|
|
if got := dispatcher.aborts.Load(); got != 0 {
|
|
t.Fatalf("AbortAmbiguousDispatch() calls = %d, want 0 for deterministic failure", got)
|
|
}
|
|
}
|
|
|
|
func TestPickNextViaHomeAbortsAmbiguousTransport(t *testing.T) {
|
|
dispatcher := &homeAuthTransportErrorDispatcher{err: home.NewAmbiguousDispatchError(errors.New("response connection closed"))}
|
|
oldCurrentHomeDispatcher := currentHomeDispatcher
|
|
currentHomeDispatcher = func() homeAuthDispatcher {
|
|
return dispatcher
|
|
}
|
|
t.Cleanup(func() {
|
|
currentHomeDispatcher = oldCurrentHomeDispatcher
|
|
})
|
|
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
registry := executionregistry.New()
|
|
manager.SetHomeExecutionRegistry(registry)
|
|
|
|
_, _, _, errPick := manager.pickNextViaHome(context.Background(), "gpt-5.4", cliproxyexecutor.Options{}, nil)
|
|
if errPick == nil {
|
|
t.Fatal("pickNextViaHome() error = nil, want home unavailable")
|
|
}
|
|
if got := dispatcher.aborts.Load(); got != 1 {
|
|
t.Fatalf("AbortAmbiguousDispatch() calls = %d, want 1", got)
|
|
}
|
|
|
|
drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancelDrain()
|
|
if errDrain := registry.Drain(drainCtx); errDrain != nil {
|
|
t.Fatalf("Drain() error = %v, ambiguous pending dispatch was not ended", errDrain)
|
|
}
|
|
}
|
|
|
|
func TestHomeRuntimeAuthsClearWhenHomeDisabled(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}})
|
|
manager.rememberHomeRuntimeAuth("session-1", &Auth{
|
|
ID: "home-auth-1",
|
|
Provider: "test",
|
|
Attributes: map[string]string{
|
|
"websockets": "true",
|
|
},
|
|
})
|
|
|
|
if _, ok := manager.GetExecutionSessionAuthByID("session-1", "home-auth-1"); !ok {
|
|
t.Fatal("expected remembered home auth before disabling home")
|
|
}
|
|
|
|
manager.SetConfig(&internalconfig.Config{})
|
|
if _, ok := manager.GetExecutionSessionAuthByID("session-1", "home-auth-1"); ok {
|
|
t.Fatal("remembered home auth was not cleared when home was disabled")
|
|
}
|
|
}
|
|
|
|
func TestCloseExecutionSessionClearsHomeRuntimeAuthForSession(t *testing.T) {
|
|
manager := NewManager(nil, nil, nil)
|
|
auth := &Auth{
|
|
ID: "home-auth-1",
|
|
Provider: "test",
|
|
Attributes: map[string]string{
|
|
"websockets": "true",
|
|
},
|
|
}
|
|
|
|
manager.rememberHomeRuntimeAuth("session-1", auth)
|
|
manager.rememberHomeRuntimeAuth("session-2", auth)
|
|
|
|
manager.CloseExecutionSession("session-1")
|
|
if _, ok := manager.GetExecutionSessionAuthByID("session-1", "home-auth-1"); ok {
|
|
t.Fatal("home auth for closed session was not cleared")
|
|
}
|
|
if _, ok := manager.GetExecutionSessionAuthByID("session-2", "home-auth-1"); !ok {
|
|
t.Fatal("home auth for another session was cleared")
|
|
}
|
|
|
|
manager.CloseExecutionSession("session-2")
|
|
if _, ok := manager.GetExecutionSessionAuthByID("session-2", "home-auth-1"); ok {
|
|
t.Fatal("home auth was not cleared when its last session closed")
|
|
}
|
|
}
|