fix(gemini): route Vertex token exchange through account proxy
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user