Skip to content

Commit ada8e29

Browse files
luispaterluispater
committed
feat(api): enhance proxy resolution for API key-based auth
Added comprehensive support for resolving proxy URLs from configuration based on API key and provider attributes. Introduced new helper functions and extended the test suite to validate fallback mechanisms and compatibility cases.
1 parent 4ba1053 commit ada8e29

2 files changed

Lines changed: 222 additions & 0 deletions

File tree

‎internal/api/handlers/management/api_tools.go‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"time"
1212

1313
"github.com/gin-gonic/gin"
14+
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
1415
"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/geminicli"
1516
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
1617
"github.com/router-for-me/CLIProxyAPI/v6/sdk/proxyutil"
@@ -636,6 +637,11 @@ func (h *Handler) apiCallTransport(auth *coreauth.Auth) http.RoundTripper {
636637
if proxyStr := strings.TrimSpace(auth.ProxyURL); proxyStr != "" {
637638
proxyCandidates = append(proxyCandidates, proxyStr)
638639
}
640+
if h != nil && h.cfg != nil {
641+
if proxyStr := strings.TrimSpace(proxyURLFromAPIKeyConfig(h.cfg, auth)); proxyStr != "" {
642+
proxyCandidates = append(proxyCandidates, proxyStr)
643+
}
644+
}
639645
}
640646
if h != nil && h.cfg != nil {
641647
if proxyStr := strings.TrimSpace(h.cfg.ProxyURL); proxyStr != "" {
@@ -658,6 +664,123 @@ func (h *Handler) apiCallTransport(auth *coreauth.Auth) http.RoundTripper {
658664
return clone
659665
}
660666

667+
type apiKeyConfigEntry interface {
668+
GetAPIKey() string
669+
GetBaseURL() string
670+
}
671+
672+
func resolveAPIKeyConfig[T apiKeyConfigEntry](entries []T, auth *coreauth.Auth) *T {
673+
if auth == nil || len(entries) == 0 {
674+
return nil
675+
}
676+
attrKey, attrBase := "", ""
677+
if auth.Attributes != nil {
678+
attrKey = strings.TrimSpace(auth.Attributes["api_key"])
679+
attrBase = strings.TrimSpace(auth.Attributes["base_url"])
680+
}
681+
for i := range entries {
682+
entry := &entries[i]
683+
cfgKey := strings.TrimSpace((*entry).GetAPIKey())
684+
cfgBase := strings.TrimSpace((*entry).GetBaseURL())
685+
if attrKey != "" && attrBase != "" {
686+
if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) {
687+
return entry
688+
}
689+
continue
690+
}
691+
if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
692+
if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
693+
return entry
694+
}
695+
}
696+
if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
697+
return entry
698+
}
699+
}
700+
if attrKey != "" {
701+
for i := range entries {
702+
entry := &entries[i]
703+
if strings.EqualFold(strings.TrimSpace((*entry).GetAPIKey()), attrKey) {
704+
return entry
705+
}
706+
}
707+
}
708+
return nil
709+
}
710+
711+
func proxyURLFromAPIKeyConfig(cfg *config.Config, auth *coreauth.Auth) string {
712+
if cfg == nil || auth == nil {
713+
return ""
714+
}
715+
authKind, authAccount := auth.AccountInfo()
716+
if !strings.EqualFold(strings.TrimSpace(authKind), "api_key") {
717+
return ""
718+
}
719+
720+
attrs := auth.Attributes
721+
compatName := ""
722+
providerKey := ""
723+
if len(attrs) > 0 {
724+
compatName = strings.TrimSpace(attrs["compat_name"])
725+
providerKey = strings.TrimSpace(attrs["provider_key"])
726+
}
727+
if compatName != "" || strings.EqualFold(strings.TrimSpace(auth.Provider), "openai-compatibility") {
728+
return resolveOpenAICompatAPIKeyProxyURL(cfg, auth, strings.TrimSpace(authAccount), providerKey, compatName)
729+
}
730+
731+
switch strings.ToLower(strings.TrimSpace(auth.Provider)) {
732+
case "gemini":
733+
if entry := resolveAPIKeyConfig(cfg.GeminiKey, auth); entry != nil {
734+
return strings.TrimSpace(entry.ProxyURL)
735+
}
736+
case "claude":
737+
if entry := resolveAPIKeyConfig(cfg.ClaudeKey, auth); entry != nil {
738+
return strings.TrimSpace(entry.ProxyURL)
739+
}
740+
case "codex":
741+
if entry := resolveAPIKeyConfig(cfg.CodexKey, auth); entry != nil {
742+
return strings.TrimSpace(entry.ProxyURL)
743+
}
744+
}
745+
return ""
746+
}
747+
748+
func resolveOpenAICompatAPIKeyProxyURL(cfg *config.Config, auth *coreauth.Auth, apiKey, providerKey, compatName string) string {
749+
if cfg == nil || auth == nil {
750+
return ""
751+
}
752+
apiKey = strings.TrimSpace(apiKey)
753+
if apiKey == "" {
754+
return ""
755+
}
756+
candidates := make([]string, 0, 3)
757+
if v := strings.TrimSpace(compatName); v != "" {
758+
candidates = append(candidates, v)
759+
}
760+
if v := strings.TrimSpace(providerKey); v != "" {
761+
candidates = append(candidates, v)
762+
}
763+
if v := strings.TrimSpace(auth.Provider); v != "" {
764+
candidates = append(candidates, v)
765+
}
766+
767+
for i := range cfg.OpenAICompatibility {
768+
compat := &cfg.OpenAICompatibility[i]
769+
for _, candidate := range candidates {
770+
if candidate != "" && strings.EqualFold(strings.TrimSpace(candidate), compat.Name) {
771+
for j := range compat.APIKeyEntries {
772+
entry := &compat.APIKeyEntries[j]
773+
if strings.EqualFold(strings.TrimSpace(entry.APIKey), apiKey) {
774+
return strings.TrimSpace(entry.ProxyURL)
775+
}
776+
}
777+
return ""
778+
}
779+
}
780+
}
781+
return ""
782+
}
783+
661784
func buildProxyTransport(proxyStr string) *http.Transport {
662785
transport, _, errBuild := proxyutil.BuildHTTPTransport(proxyStr)
663786
if errBuild != nil {

‎internal/api/handlers/management/api_tools_test.go‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,105 @@ func TestAPICallTransportInvalidAuthFallsBackToGlobalProxy(t *testing.T) {
5858
}
5959
}
6060

61+
func TestAPICallTransportAPIKeyAuthFallsBackToConfigProxyURL(t *testing.T) {
62+
t.Parallel()
63+
64+
h := &Handler{
65+
cfg: &config.Config{
66+
SDKConfig: sdkconfig.SDKConfig{ProxyURL: "http://global-proxy.example.com:8080"},
67+
GeminiKey: []config.GeminiKey{{
68+
APIKey: "gemini-key",
69+
ProxyURL: "http://gemini-proxy.example.com:8080",
70+
}},
71+
ClaudeKey: []config.ClaudeKey{{
72+
APIKey: "claude-key",
73+
ProxyURL: "http://claude-proxy.example.com:8080",
74+
}},
75+
CodexKey: []config.CodexKey{{
76+
APIKey: "codex-key",
77+
ProxyURL: "http://codex-proxy.example.com:8080",
78+
}},
79+
OpenAICompatibility: []config.OpenAICompatibility{{
80+
Name: "bohe",
81+
BaseURL: "https://bohe.example.com",
82+
APIKeyEntries: []config.OpenAICompatibilityAPIKey{{
83+
APIKey: "compat-key",
84+
ProxyURL: "http://compat-proxy.example.com:8080",
85+
}},
86+
}},
87+
},
88+
}
89+
90+
cases := []struct {
91+
name string
92+
auth *coreauth.Auth
93+
wantProxy string
94+
}{
95+
{
96+
name: "gemini",
97+
auth: &coreauth.Auth{
98+
Provider: "gemini",
99+
Attributes: map[string]string{"api_key": "gemini-key"},
100+
},
101+
wantProxy: "http://gemini-proxy.example.com:8080",
102+
},
103+
{
104+
name: "claude",
105+
auth: &coreauth.Auth{
106+
Provider: "claude",
107+
Attributes: map[string]string{"api_key": "claude-key"},
108+
},
109+
wantProxy: "http://claude-proxy.example.com:8080",
110+
},
111+
{
112+
name: "codex",
113+
auth: &coreauth.Auth{
114+
Provider: "codex",
115+
Attributes: map[string]string{"api_key": "codex-key"},
116+
},
117+
wantProxy: "http://codex-proxy.example.com:8080",
118+
},
119+
{
120+
name: "openai-compatibility",
121+
auth: &coreauth.Auth{
122+
Provider: "bohe",
123+
Attributes: map[string]string{
124+
"api_key": "compat-key",
125+
"compat_name": "bohe",
126+
"provider_key": "bohe",
127+
},
128+
},
129+
wantProxy: "http://compat-proxy.example.com:8080",
130+
},
131+
}
132+
133+
for _, tc := range cases {
134+
tc := tc
135+
t.Run(tc.name, func(t *testing.T) {
136+
t.Parallel()
137+
138+
transport := h.apiCallTransport(tc.auth)
139+
httpTransport, ok := transport.(*http.Transport)
140+
if !ok {
141+
t.Fatalf("transport type = %T, want *http.Transport", transport)
142+
}
143+
144+
req, errRequest := http.NewRequest(http.MethodGet, "https://example.com", nil)
145+
if errRequest != nil {
146+
t.Fatalf("http.NewRequest returned error: %v", errRequest)
147+
}
148+
149+
proxyURL, errProxy := httpTransport.Proxy(req)
150+
if errProxy != nil {
151+
t.Fatalf("httpTransport.Proxy returned error: %v", errProxy)
152+
}
153+
if proxyURL == nil || proxyURL.String() != tc.wantProxy {
154+
t.Fatalf("proxy URL = %v, want %s", proxyURL, tc.wantProxy)
155+
}
156+
})
157+
}
158+
}
159+
61160
func TestAuthByIndexDistinguishesSharedAPIKeysAcrossProviders(t *testing.T) {
62161
t.Parallel()
63162

0 commit comments

Comments
 (0)