google_cloud_ai/
lib.rs

1use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
2use serde::{Deserialize, Serialize};
3
4#[path = "google/protos.rs"]
5pub mod google_cloud_ai;
6
7#[derive(Debug, Serialize, Deserialize)]
8struct Token {
9    access_token: String,
10}
11
12#[derive(Debug, Serialize, Deserialize)]
13struct Claims {
14    iss: String,
15    scope: String,
16    aud: String,
17    exp: usize,
18    iat: usize,
19}
20
21pub async fn get_token() -> String {
22    let client = reqwest::Client::new();
23
24    let response = client.get("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token")
25    .header("Metadata-Flavor", "Google")
26    .send()
27    .await
28    .unwrap();
29
30    let result = response.json::<Token>().await.unwrap();
31
32    return result.access_token;
33}
34
35pub async fn get_service_account_token(email: String, rsa_pem: String) -> String {
36    let timestamp = chrono::Utc::now();
37
38    let iat_value = timestamp.timestamp().try_into().unwrap();
39    let exp_value = (timestamp.timestamp() + 3600).try_into().unwrap();
40
41    let claims = Claims {
42        iss: email,
43        scope: String::from("https://www.googleapis.com/auth/datastore"),
44        aud: String::from("https://oauth2.googleapis.com/token"),
45        exp: exp_value,
46        iat: iat_value,
47    };
48
49    let mut headers = Header::new(Algorithm::RS256);
50
51    headers.typ = Some("JWT".to_string());
52
53    let token = encode(
54        &headers,
55        &claims,
56        &EncodingKey::from_rsa_pem(rsa_pem.as_ref()).unwrap(),
57    )
58    .unwrap();
59
60    let client = reqwest::Client::new();
61
62    let resp = client
63        .post("https://oauth2.googleapis.com/token")
64        .header("Content-Type", "application/x-www-form-urlencoded")
65        .body(format!(
66            "grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer&assertion={}",
67            token
68        ))
69        .send()
70        .await
71        .unwrap();
72
73    let result = resp.json::<Token>().await.unwrap();
74
75    return result.access_token;
76}