Skip to main content

reqsign_google/provide_credential/
vm_metadata.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use log::debug;
19use serde::Deserialize;
20use std::time::Duration;
21
22use crate::credential::{Credential, Token};
23use reqsign_core::time::Timestamp;
24use reqsign_core::{Context, ProvideCredential, Result};
25
26/// VM metadata token response.
27#[derive(Deserialize)]
28struct VmMetadataTokenResponse {
29    access_token: String,
30    expires_in: u64,
31}
32
33/// VmMetadataCredentialProvider loads tokens from Google Compute Engine VM metadata service.
34#[derive(Debug, Clone, Default)]
35pub struct VmMetadataCredentialProvider {
36    scope: Option<String>,
37    endpoint: Option<String>,
38    service_account: Option<String>,
39}
40
41impl VmMetadataCredentialProvider {
42    /// Create a new VmMetadataCredentialProvider.
43    pub fn new() -> Self {
44        Self::default()
45    }
46
47    /// Set the OAuth2 scope.
48    pub fn with_scope(mut self, scope: impl Into<String>) -> Self {
49        self.scope = Some(scope.into());
50        self
51    }
52
53    /// Set the metadata endpoint.
54    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
55        self.endpoint = Some(endpoint.into());
56        self
57    }
58
59    /// Set the service account used to retrieve a token from VM metadata service.
60    ///
61    /// Defaults to `default` if not configured. A configured value other than `default` is also
62    /// preserved as the token credential's signer email for query signing.
63    pub fn with_service_account(mut self, service_account: impl Into<String>) -> Self {
64        self.service_account = Some(service_account.into());
65        self
66    }
67}
68impl ProvideCredential for VmMetadataCredentialProvider {
69    type Credential = Credential;
70
71    async fn provide_credential(&self, ctx: &Context) -> Result<Option<Self::Credential>> {
72        // Get scope from instance, environment, or use default
73        let scope = self
74            .scope
75            .clone()
76            .or_else(|| ctx.env_var(crate::constants::GOOGLE_SCOPE))
77            .unwrap_or_else(|| crate::constants::DEFAULT_SCOPE.to_string());
78
79        let service_account = self.service_account.as_deref().unwrap_or("default");
80
81        debug!("loading token from VM metadata service for account: {service_account}");
82
83        // Allow overriding metadata host for testing
84        let metadata_host = self
85            .endpoint
86            .clone()
87            .or_else(|| ctx.env_var("GCE_METADATA_HOST"))
88            .unwrap_or_else(|| "metadata.google.internal".to_string());
89
90        let url = format!(
91            "http://{metadata_host}/computeMetadata/v1/instance/service-accounts/{service_account}/token?scopes={scope}"
92        );
93
94        let req = http::Request::builder()
95            .method(http::Method::GET)
96            .uri(&url)
97            .header("Metadata-Flavor", "Google")
98            .body(Vec::<u8>::new().into())
99            .map_err(|e| {
100                reqsign_core::Error::unexpected("failed to build HTTP request").with_source(e)
101            })?;
102
103        let resp = ctx.http_send(req).await?;
104
105        if resp.status() != http::StatusCode::OK {
106            // VM metadata service might not be available (e.g., not running on GCE)
107            debug!("VM metadata service not available or returned error");
108            return Ok(None);
109        }
110
111        let token_resp: VmMetadataTokenResponse =
112            serde_json::from_slice(resp.body()).map_err(|e| {
113                reqsign_core::Error::unexpected("failed to parse VM metadata response")
114                    .with_source(e)
115            })?;
116
117        let expires_at = Timestamp::now() + Duration::from_secs(token_resp.expires_in);
118        let credential = Credential::with_token(Token {
119            access_token: token_resp.access_token,
120            expires_at: Some(expires_at),
121        });
122        Ok(Some(match self.service_account.as_deref() {
123            Some(service_account) if service_account != "default" => {
124                credential.with_signer_email(service_account)
125            }
126            _ => credential,
127        }))
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134    use bytes::Bytes;
135    use reqsign_core::HttpSend;
136    use std::sync::{Arc, Mutex};
137
138    #[derive(Clone, Debug, Default)]
139    struct MockHttpSend {
140        uris: Arc<Mutex<Vec<String>>>,
141    }
142
143    impl HttpSend for MockHttpSend {
144        async fn http_send(&self, req: http::Request<Bytes>) -> Result<http::Response<Bytes>> {
145            self.uris.lock().unwrap().push(req.uri().to_string());
146
147            Ok(http::Response::builder()
148                .status(http::StatusCode::OK)
149                .body(
150                    include_bytes!("../../tests/fixtures/vm_metadata_token_response.json")
151                        .as_slice()
152                        .into(),
153                )
154                .expect("response must build"))
155        }
156    }
157
158    #[tokio::test]
159    async fn test_vm_metadata_uses_default_service_account() -> Result<()> {
160        let http = MockHttpSend::default();
161        let ctx = Context::new().with_http_send(http.clone());
162
163        let provider = VmMetadataCredentialProvider::new().with_endpoint("127.0.0.1:8080");
164        let cred = provider
165            .provide_credential(&ctx)
166            .await?
167            .expect("credential must exist");
168
169        assert!(cred.has_token());
170        assert_eq!(
171            cred.token.as_ref().map(|token| token.access_token.as_str()),
172            Some("REDACTED")
173        );
174        assert!(cred.signer_email.is_none());
175        assert_eq!(
176            http.uris.lock().unwrap().as_slice(),
177            &["http://127.0.0.1:8080/computeMetadata/v1/instance/service-accounts/default/token?scopes=https://www.googleapis.com/auth/cloud-platform".to_string()]
178        );
179
180        Ok(())
181    }
182
183    #[tokio::test]
184    async fn test_vm_metadata_uses_configured_service_account() -> Result<()> {
185        let http = MockHttpSend::default();
186        let ctx = Context::new().with_http_send(http.clone());
187
188        let provider = VmMetadataCredentialProvider::new()
189            .with_endpoint("127.0.0.1:8080")
190            .with_service_account("custom@test-project.iam.gserviceaccount.com");
191        let cred = provider
192            .provide_credential(&ctx)
193            .await?
194            .expect("credential must exist");
195
196        assert!(cred.has_token());
197        assert_eq!(
198            cred.token.as_ref().map(|token| token.access_token.as_str()),
199            Some("REDACTED")
200        );
201        assert_eq!(
202            cred.signer_email.as_deref(),
203            Some("custom@test-project.iam.gserviceaccount.com")
204        );
205        assert_eq!(
206            http.uris.lock().unwrap().as_slice(),
207            &["http://127.0.0.1:8080/computeMetadata/v1/instance/service-accounts/custom@test-project.iam.gserviceaccount.com/token?scopes=https://www.googleapis.com/auth/cloud-platform".to_string()]
208        );
209
210        Ok(())
211    }
212}