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.
62    pub fn with_service_account(mut self, service_account: impl Into<String>) -> Self {
63        self.service_account = Some(service_account.into());
64        self
65    }
66}
67impl ProvideCredential for VmMetadataCredentialProvider {
68    type Credential = Credential;
69
70    async fn provide_credential(&self, ctx: &Context) -> Result<Option<Self::Credential>> {
71        // Get scope from instance, environment, or use default
72        let scope = self
73            .scope
74            .clone()
75            .or_else(|| ctx.env_var(crate::constants::GOOGLE_SCOPE))
76            .unwrap_or_else(|| crate::constants::DEFAULT_SCOPE.to_string());
77
78        let service_account = self.service_account.as_deref().unwrap_or("default");
79
80        debug!("loading token from VM metadata service for account: {service_account}");
81
82        // Allow overriding metadata host for testing
83        let metadata_host = self
84            .endpoint
85            .clone()
86            .or_else(|| ctx.env_var("GCE_METADATA_HOST"))
87            .unwrap_or_else(|| "metadata.google.internal".to_string());
88
89        let url = format!(
90            "http://{metadata_host}/computeMetadata/v1/instance/service-accounts/{service_account}/token?scopes={scope}"
91        );
92
93        let req = http::Request::builder()
94            .method(http::Method::GET)
95            .uri(&url)
96            .header("Metadata-Flavor", "Google")
97            .body(Vec::<u8>::new().into())
98            .map_err(|e| {
99                reqsign_core::Error::unexpected("failed to build HTTP request").with_source(e)
100            })?;
101
102        let resp = ctx.http_send(req).await?;
103
104        if resp.status() != http::StatusCode::OK {
105            // VM metadata service might not be available (e.g., not running on GCE)
106            debug!("VM metadata service not available or returned error");
107            return Ok(None);
108        }
109
110        let token_resp: VmMetadataTokenResponse =
111            serde_json::from_slice(resp.body()).map_err(|e| {
112                reqsign_core::Error::unexpected("failed to parse VM metadata response")
113                    .with_source(e)
114            })?;
115
116        let expires_at = Timestamp::now() + Duration::from_secs(token_resp.expires_in);
117        Ok(Some(Credential::with_token(Token {
118            access_token: token_resp.access_token,
119            expires_at: Some(expires_at),
120        })))
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127    use bytes::Bytes;
128    use reqsign_core::HttpSend;
129    use std::sync::{Arc, Mutex};
130
131    #[derive(Clone, Debug, Default)]
132    struct MockHttpSend {
133        uris: Arc<Mutex<Vec<String>>>,
134    }
135
136    impl HttpSend for MockHttpSend {
137        async fn http_send(&self, req: http::Request<Bytes>) -> Result<http::Response<Bytes>> {
138            self.uris.lock().unwrap().push(req.uri().to_string());
139
140            Ok(http::Response::builder()
141                .status(http::StatusCode::OK)
142                .body(
143                    br#"{"access_token":"test-access-token","expires_in":3600}"#
144                        .as_slice()
145                        .into(),
146                )
147                .expect("response must build"))
148        }
149    }
150
151    #[tokio::test]
152    async fn test_vm_metadata_uses_default_service_account() -> Result<()> {
153        let http = MockHttpSend::default();
154        let ctx = Context::new().with_http_send(http.clone());
155
156        let provider = VmMetadataCredentialProvider::new().with_endpoint("127.0.0.1:8080");
157        let cred = provider
158            .provide_credential(&ctx)
159            .await?
160            .expect("credential must exist");
161
162        assert!(cred.has_token());
163        assert_eq!(
164            http.uris.lock().unwrap().as_slice(),
165            &["http://127.0.0.1:8080/computeMetadata/v1/instance/service-accounts/default/token?scopes=https://www.googleapis.com/auth/cloud-platform".to_string()]
166        );
167
168        Ok(())
169    }
170
171    #[tokio::test]
172    async fn test_vm_metadata_uses_configured_service_account() -> Result<()> {
173        let http = MockHttpSend::default();
174        let ctx = Context::new().with_http_send(http.clone());
175
176        let provider = VmMetadataCredentialProvider::new()
177            .with_endpoint("127.0.0.1:8080")
178            .with_service_account("custom@test-project.iam.gserviceaccount.com");
179        let cred = provider
180            .provide_credential(&ctx)
181            .await?
182            .expect("credential must exist");
183
184        assert!(cred.has_token());
185        assert_eq!(
186            http.uris.lock().unwrap().as_slice(),
187            &["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()]
188        );
189
190        Ok(())
191    }
192}