Skip to main content

agent_core/core/auth/
google_vertex.rs

1//! Provider-local Google Vertex AI installed-app OAuth primitives.
2use serde::Deserialize;
3use std::fmt;
4use url::Url;
5
6pub const PROVIDER: &str = "google-vertex";
7pub const AUTHORIZE_URL: &str = "https://accounts.google.com/o/oauth2/v2/auth";
8pub const TOKEN_URL: &str = "https://oauth2.googleapis.com/token";
9pub const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
10
11#[derive(Clone, PartialEq, Eq)]
12pub struct VertexRegistration {
13    client_id: String,
14}
15impl VertexRegistration {
16    pub fn new(client_id: Option<&str>) -> Result<Self, VertexError> {
17        let client_id = client_id
18            .map(str::trim)
19            .filter(|v| !v.is_empty())
20            .ok_or(VertexError::RegistrationRequired)?;
21        if !client_id.ends_with(".apps.googleusercontent.com") || client_id.contains('/') {
22            return Err(VertexError::InvalidRegistration);
23        }
24        Ok(Self {
25            client_id: client_id.into(),
26        })
27    }
28    pub fn client_id(&self) -> &str {
29        &self.client_id
30    }
31}
32impl fmt::Debug for VertexRegistration {
33    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34        f.debug_struct("VertexRegistration")
35            .field("client_id", &"[configured]")
36            .finish()
37    }
38}
39
40#[derive(Debug, Clone, PartialEq, Eq)]
41pub struct VertexContext {
42    project: String,
43    location: String,
44}
45impl VertexContext {
46    pub fn new(project: &str, location: &str) -> Result<Self, VertexError> {
47        if !valid_project(project) {
48            return Err(VertexError::InvalidProject);
49        }
50        if !valid_location(location) {
51            return Err(VertexError::InvalidLocation);
52        }
53        Ok(Self {
54            project: project.into(),
55            location: location.into(),
56        })
57    }
58    pub fn project(&self) -> &str {
59        &self.project
60    }
61    pub fn location(&self) -> &str {
62        &self.location
63    }
64    pub fn host(&self) -> String {
65        format!("{}-aiplatform.googleapis.com", self.location)
66    }
67}
68fn valid_project(v: &str) -> bool {
69    (6..=63).contains(&v.len())
70        && v.starts_with(|c: char| c.is_ascii_lowercase())
71        && v.ends_with(|c: char| c.is_ascii_lowercase() || c.is_ascii_digit())
72        && v.bytes()
73            .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == b'-')
74}
75fn valid_location(v: &str) -> bool {
76    (2..=32).contains(&v.len())
77        && v.bytes()
78            .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == b'-')
79        && !v.starts_with('-')
80        && !v.ends_with('-')
81}
82
83pub struct VertexCredentials {
84    access: String,
85    refresh: String,
86    expires: u64,
87}
88impl VertexCredentials {
89    pub fn new(access: impl Into<String>, refresh: impl Into<String>, expires: u64) -> Self {
90        Self {
91            access: access.into(),
92            refresh: refresh.into(),
93            expires,
94        }
95    }
96    pub fn access_token(&self) -> &str {
97        &self.access
98    }
99    pub fn refresh_token(&self) -> &str {
100        &self.refresh
101    }
102    pub fn expires_at_ms(&self) -> u64 {
103        self.expires
104    }
105}
106impl fmt::Debug for VertexCredentials {
107    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
108        f.debug_struct("VertexCredentials")
109            .field("access", &"[REDACTED]")
110            .field("refresh", &"[REDACTED]")
111            .field("expires", &self.expires)
112            .finish()
113    }
114}
115
116#[derive(Debug, thiserror::Error, PartialEq, Eq)]
117pub enum VertexError {
118    #[error("google-vertex: Synaps Google Desktop OAuth client registration required; configure the broker-owned client ID")]
119    RegistrationRequired,
120    #[error("google-vertex: invalid Desktop OAuth client registration")]
121    InvalidRegistration,
122    #[error("google-vertex: invalid project")]
123    InvalidProject,
124    #[error("google-vertex: invalid location")]
125    InvalidLocation,
126    #[error("google-vertex: untrusted OAuth endpoint")]
127    UntrustedEndpoint,
128    #[error("google-vertex: invalid token response")]
129    InvalidTokenResponse,
130}
131impl VertexError {
132    pub fn code(&self) -> &'static str {
133        match self {
134            Self::RegistrationRequired => "registration_required",
135            Self::InvalidRegistration => "invalid_registration",
136            Self::InvalidProject => "invalid_project",
137            Self::InvalidLocation => "invalid_location",
138            Self::UntrustedEndpoint => "untrusted_endpoint",
139            Self::InvalidTokenResponse => "invalid_token_response",
140        }
141    }
142}
143
144pub fn build_authorize_url(
145    reg: &VertexRegistration,
146    challenge: &str,
147    state: &str,
148    redirect: &str,
149) -> Result<Url, VertexError> {
150    let redirect = Url::parse(redirect).map_err(|_| VertexError::UntrustedEndpoint)?;
151    if redirect.scheme() != "http"
152        || redirect.host_str() != Some("127.0.0.1")
153        || redirect.port().is_none()
154        || redirect.path() != "/oauth2callback"
155    {
156        return Err(VertexError::UntrustedEndpoint);
157    }
158    let mut url = Url::parse(AUTHORIZE_URL).unwrap();
159    url.query_pairs_mut()
160        .append_pair("client_id", reg.client_id())
161        .append_pair("response_type", "code")
162        .append_pair("redirect_uri", redirect.as_str())
163        .append_pair("scope", CLOUD_PLATFORM_SCOPE)
164        .append_pair("access_type", "offline")
165        .append_pair("prompt", "consent")
166        .append_pair("code_challenge", challenge)
167        .append_pair("code_challenge_method", "S256")
168        .append_pair("state", state);
169    Ok(url)
170}
171
172pub fn validate_oauth_endpoint(value: &str, test_loopback: bool) -> Result<Url, VertexError> {
173    let url = Url::parse(value).map_err(|_| VertexError::UntrustedEndpoint)?;
174    let production =
175        (url.as_str() == TOKEN_URL || url.as_str() == AUTHORIZE_URL) && url.scheme() == "https";
176    let seam = test_loopback && url.scheme() == "http" && url.host_str() == Some("127.0.0.1");
177    if production || seam {
178        Ok(url)
179    } else {
180        Err(VertexError::UntrustedEndpoint)
181    }
182}
183
184#[derive(Deserialize)]
185struct TokenWire {
186    access_token: String,
187    #[serde(default)]
188    refresh_token: Option<String>,
189    expires_in: u64,
190}
191pub fn parse_token_response(
192    body: &[u8],
193    previous: Option<&VertexCredentials>,
194    now_ms: u64,
195) -> Result<VertexCredentials, VertexError> {
196    if body.len() > 64 * 1024 {
197        return Err(VertexError::InvalidTokenResponse);
198    }
199    let wire: TokenWire =
200        serde_json::from_slice(body).map_err(|_| VertexError::InvalidTokenResponse)?;
201    if wire.access_token.is_empty() {
202        return Err(VertexError::InvalidTokenResponse);
203    }
204    let refresh = wire
205        .refresh_token
206        .filter(|v| !v.is_empty())
207        .or_else(|| previous.map(|v| v.refresh.clone()))
208        .ok_or(VertexError::InvalidTokenResponse)?;
209    Ok(VertexCredentials::new(
210        wire.access_token,
211        refresh,
212        now_ms.saturating_add(wire.expires_in.saturating_mul(1000)),
213    ))
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219    #[test]
220    fn identity_and_registration_are_distinct_from_gemini_cli() {
221        assert_eq!(PROVIDER, "google-vertex");
222        let error = VertexRegistration::new(None).unwrap_err();
223        assert_eq!(error.code(), "registration_required");
224        assert!(!error.to_string().contains("681255809395"));
225    }
226    #[test]
227    fn validates_context_before_network() {
228        assert!(VertexContext::new("my-project-123", "us-central1").is_ok());
229        for project in ["", "UPPER", "-bad", "bad_", "a"] {
230            assert!(
231                VertexContext::new(project, "us-central1").is_err(),
232                "{project}"
233            );
234        }
235        for location in ["", "US-CENTRAL1", "../evil", "us_central1"] {
236            assert!(
237                VertexContext::new("my-project-123", location).is_err(),
238                "{location}"
239            );
240        }
241    }
242    #[test]
243    fn authorization_is_pkce_offline_and_exact_scope() {
244        let registration =
245            VertexRegistration::new(Some("synaps.apps.googleusercontent.com")).unwrap();
246        let url = build_authorize_url(
247            &registration,
248            "challenge",
249            "state",
250            "http://127.0.0.1:3210/oauth2callback",
251        )
252        .unwrap();
253        let query: std::collections::HashMap<_, _> = url.query_pairs().into_owned().collect();
254        assert_eq!(query["scope"], CLOUD_PLATFORM_SCOPE);
255        assert_eq!(query["code_challenge_method"], "S256");
256        assert_eq!(query["access_type"], "offline");
257        assert_eq!(query["state"], "state");
258        assert!(!query.contains_key("client_secret"));
259    }
260    #[test]
261    fn refresh_rotation_preserves_omitted_refresh_token() {
262        let old = VertexCredentials::new("access-old", "refresh-old", 1);
263        let next = parse_token_response(
264            br#"{"access_token":"access-new","expires_in":3600}"#,
265            Some(&old),
266            10,
267        )
268        .unwrap();
269        assert_eq!(next.refresh_token(), "refresh-old");
270        assert_eq!(next.access_token(), "access-new");
271        let rotated = parse_token_response(
272            br#"{"access_token":"a","refresh_token":"refresh-new","expires_in":2}"#,
273            Some(&old),
274            10,
275        )
276        .unwrap();
277        assert_eq!(rotated.refresh_token(), "refresh-new");
278    }
279    #[test]
280    fn only_exact_oauth_endpoints_and_loopback_test_seam_pass() {
281        assert!(validate_oauth_endpoint(TOKEN_URL, false).is_ok());
282        assert!(
283            validate_oauth_endpoint("https://oauth2.googleapis.com.evil.test/token", false)
284                .is_err()
285        );
286        assert!(validate_oauth_endpoint("http://127.0.0.1:9000/token", true).is_ok());
287        assert!(validate_oauth_endpoint("http://10.0.0.1/token", true).is_err());
288    }
289}