fix(gemini): route Vertex token exchange through account proxy

This commit is contained in:
ZeroDeng
2026-05-08 14:03:15 +08:00
parent fda1ed459d
commit 2a17c0b229
2 changed files with 109 additions and 3 deletions
@@ -1,8 +1,17 @@
package service
import (
"context"
"crypto/rand"
"crypto/rsa"
"crypto/x509"
"encoding/json"
"encoding/pem"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
@@ -75,3 +84,70 @@ func TestParseVertexServiceAccountKey(t *testing.T) {
require.Equal(t, vertexDefaultTokenURL, key.TokenURI)
require.True(t, strings.Contains(key.PrivateKey, "BEGIN PRIVATE KEY"))
}
func TestVertexServiceAccountProxyURL(t *testing.T) {
proxyID := int64(7)
account := &Account{
ProxyID: &proxyID,
Proxy: &Proxy{
Protocol: "http",
Host: "proxy.example.com",
Port: 8080,
},
}
require.Equal(t, "http://proxy.example.com:8080", vertexServiceAccountProxyURL(account))
require.Empty(t, vertexServiceAccountProxyURL(&Account{Proxy: account.Proxy}))
require.Empty(t, vertexServiceAccountProxyURL(&Account{ProxyID: &proxyID}))
}
func TestExchangeVertexServiceAccountTokenUsesProxy(t *testing.T) {
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
require.NoError(t, err)
pemBytes := pem.EncodeToMemory(&pem.Block{
Type: "RSA PRIVATE KEY",
Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
})
seenProxyRequest := make(chan string, 1)
proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
seenProxyRequest <- r.URL.String()
require.Equal(t, "oauth2.googleapis.com", r.URL.Host)
require.Equal(t, "/token", r.URL.Path)
require.NoError(t, r.ParseForm())
require.Equal(t, "urn:ietf:params:oauth:grant-type:jwt-bearer", r.PostForm.Get("grant_type"))
w.Header().Set("Content-Type", "application/json")
require.NoError(t, json.NewEncoder(w).Encode(vertexTokenResponse{
AccessToken: "proxied-token",
TokenType: "Bearer",
ExpiresIn: 3600,
}))
}))
defer proxy.Close()
key := &vertexServiceAccountKey{
ClientEmail: "svc@example.iam.gserviceaccount.com",
PrivateKey: string(pemBytes),
TokenURI: "http://oauth2.googleapis.com/token",
}
token, ttl, err := exchangeVertexServiceAccountToken(contextWithTestTimeout(t), key, proxy.URL)
require.NoError(t, err)
require.Equal(t, "proxied-token", token)
require.Equal(t, 55*time.Minute, ttl)
select {
case got := <-seenProxyRequest:
require.Equal(t, "http://oauth2.googleapis.com/token", got)
default:
t.Fatal("token request did not reach proxy")
}
}
func contextWithTestTimeout(t *testing.T) context.Context {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
t.Cleanup(cancel)
return ctx
}