fix(vertex): audit fixes for Vertex Service Account feature (#1977)
- Security: force token_uri to Google default, preventing SSRF via crafted service account JSON - Dedup: extract shared getVertexServiceAccountAccessToken() to eliminate ~35 lines of duplication between ClaudeTokenProvider and GeminiTokenProvider - Fix: apply model mapping + Vertex model ID normalization in forward_as_responses and forward_as_chat_completions paths - Fix: exclude service_account from AI Studio endpoint selection (Vertex cannot serve generativelanguage.googleapis.com) - Feature: add model restriction/mapping UI for service_account in EditAccountModal - Dedup: extract VERTEX_LOCATION_OPTIONS to shared constants - i18n: replace all hardcoded Chinese strings in Vertex UI with translation keys
This commit is contained in:
@@ -9,6 +9,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"regexp"
|
||||
@@ -23,6 +24,7 @@ const (
|
||||
vertexDefaultTokenURL = "https://oauth2.googleapis.com/token"
|
||||
vertexCloudPlatformScope = "https://www.googleapis.com/auth/cloud-platform"
|
||||
vertexServiceAccountCacheSkew = 5 * time.Minute
|
||||
vertexLockWaitTime = 200 * time.Millisecond
|
||||
vertexAnthropicVersion = "vertex-2023-10-16"
|
||||
)
|
||||
|
||||
@@ -123,9 +125,8 @@ func parseVertexServiceAccountJSON(raw []byte) (*vertexServiceAccountKey, error)
|
||||
if strings.TrimSpace(key.ProjectID) == "" {
|
||||
return nil, errors.New("service account json missing project_id")
|
||||
}
|
||||
if strings.TrimSpace(key.TokenURI) == "" {
|
||||
key.TokenURI = vertexDefaultTokenURL
|
||||
}
|
||||
// Always use the well-known Google token endpoint to prevent SSRF via crafted token_uri.
|
||||
key.TokenURI = vertexDefaultTokenURL
|
||||
return &key, nil
|
||||
}
|
||||
|
||||
@@ -141,6 +142,47 @@ func vertexServiceAccountCacheKey(account *Account, key *vertexServiceAccountKey
|
||||
return "vertex:service_account:" + fingerprint
|
||||
}
|
||||
|
||||
// getVertexServiceAccountAccessToken obtains an access token for a Vertex service account,
|
||||
// using the shared cache and distributed lock to avoid redundant exchanges.
|
||||
func getVertexServiceAccountAccessToken(ctx context.Context, cache GeminiTokenCache, account *Account) (string, error) {
|
||||
key, err := parseVertexServiceAccountKey(account)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
cacheKey := vertexServiceAccountCacheKey(account, key)
|
||||
|
||||
if cache != nil {
|
||||
if token, err := cache.GetAccessToken(ctx, cacheKey); err == nil && strings.TrimSpace(token) != "" {
|
||||
return token, nil
|
||||
}
|
||||
}
|
||||
|
||||
locked := false
|
||||
if cache != nil {
|
||||
var lockErr error
|
||||
locked, lockErr = cache.AcquireRefreshLock(ctx, cacheKey, 30*time.Second)
|
||||
if lockErr == nil && locked {
|
||||
defer func() { _ = cache.ReleaseRefreshLock(ctx, cacheKey) }()
|
||||
} else if lockErr != nil {
|
||||
slog.Warn("vertex_service_account_token_lock_failed", "account_id", account.ID, "error", lockErr)
|
||||
} else {
|
||||
time.Sleep(vertexLockWaitTime)
|
||||
if token, err := cache.GetAccessToken(ctx, cacheKey); err == nil && strings.TrimSpace(token) != "" {
|
||||
return token, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
accessToken, ttl, err := exchangeVertexServiceAccountToken(ctx, key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if cache != nil {
|
||||
_ = cache.SetAccessToken(ctx, cacheKey, accessToken, ttl)
|
||||
}
|
||||
return accessToken, nil
|
||||
}
|
||||
|
||||
func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAccountKey) (string, time.Duration, error) {
|
||||
now := time.Now()
|
||||
claims := jwt.MapClaims{
|
||||
|
||||
Reference in New Issue
Block a user