fix(gemini): route Vertex token exchange through account proxy
This commit is contained in:
@@ -16,6 +16,8 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
@@ -173,7 +175,7 @@ func getVertexServiceAccountAccessToken(ctx context.Context, cache GeminiTokenCa
|
||||
}
|
||||
}
|
||||
|
||||
accessToken, ttl, err := exchangeVertexServiceAccountToken(ctx, key)
|
||||
accessToken, ttl, err := exchangeVertexServiceAccountToken(ctx, key, vertexServiceAccountProxyURL(account))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
@@ -183,7 +185,32 @@ func getVertexServiceAccountAccessToken(ctx context.Context, cache GeminiTokenCa
|
||||
return accessToken, nil
|
||||
}
|
||||
|
||||
func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAccountKey) (string, time.Duration, error) {
|
||||
func vertexServiceAccountProxyURL(account *Account) string {
|
||||
if account == nil || account.ProxyID == nil || account.Proxy == nil {
|
||||
return ""
|
||||
}
|
||||
return account.Proxy.URL()
|
||||
}
|
||||
|
||||
func newVertexServiceAccountHTTPClient(proxyURL string) (*http.Client, error) {
|
||||
proxyURL = strings.TrimSpace(proxyURL)
|
||||
if proxyURL == "" {
|
||||
return &http.Client{Timeout: 15 * time.Second}, nil
|
||||
}
|
||||
|
||||
_, parsedProxy, err := proxyurl.Parse(proxyURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
transport := http.DefaultTransport.(*http.Transport).Clone()
|
||||
transport.Proxy = nil
|
||||
if err := proxyutil.ConfigureTransportProxy(transport, parsedProxy); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &http.Client{Timeout: 15 * time.Second, Transport: transport}, nil
|
||||
}
|
||||
|
||||
func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAccountKey, proxyURL string) (string, time.Duration, error) {
|
||||
now := time.Now()
|
||||
claims := jwt.MapClaims{
|
||||
"iss": key.ClientEmail,
|
||||
@@ -215,7 +242,10 @@ func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAc
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
|
||||
client := &http.Client{Timeout: 15 * time.Second}
|
||||
client, err := newVertexServiceAccountHTTPClient(proxyURL)
|
||||
if err != nil {
|
||||
return "", 0, fmt.Errorf("configure service account token proxy: %w", err)
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", 0, fmt.Errorf("service account token request failed: %w", err)
|
||||
|
||||
Reference in New Issue
Block a user