Adding service account impersonation (waiting on #76)

msdrigg · msdrigg · commit fc39daccb7fa · 2023-08-18T11:20:08.000-04:00
This PR adds a new `ServiceAccount` format that takes credentials from `source_credentials: ServiceAccount` and then makes a request to get a service account token using those credentials.

This also adds the ability to parse the token format created by `gcloud auth application-default login --impersonate-service-account &lt;service account&gt;`
diff --git a/src/error.rs b/src/error.rs
@@ -96,6 +96,10 @@ pub enum Error {
     #[error("Failed to parse output of GCloud")]
     GCloudParseError,
 
+    /// Currently, nested service account impersonation is not supported
+    #[error("Nested impersonation is not supported")]
+    NestedImpersonation,
+
     /// Represents all other cases of `std::io::Error`.
     #[error(transparent)]
     IOError(#[from] std::io::Error),
diff --git a/src/flexible_credential_source.rs b/src/flexible_credential_source.rs
@@ -1,12 +1,13 @@
 use std::path::{Path, PathBuf};
 
-use serde::{Deserialize, Serialize};
+use serde::Deserialize;
 use tokio::fs;
 
 use crate::{
     authentication_manager::ServiceAccount,
     custom_service_account::ApplicationCredentials,
     default_authorized_user::{ConfigDefaultCredentials, UserCredentials},
+    service_account_impersonation::ImpersonatedServiceAccount,
     types::HyperClient,
     CustomServiceAccount, Error,
 };
@@ -15,7 +16,7 @@ use crate::{
 // https://github.com/golang/oauth2/blob/a835fc4358f6852f50c4c5c33fddcd1adade5b0a/google/google.go#L158
 // Currently not implementing external account credentials
 // Currently not implementing impersonating service accounts (coming soon !)
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Deserialize, Debug)]
 #[serde(tag = "type", rename_all = "snake_case")]
 pub(crate) enum FlexibleCredentialSource {
     // This credential parses the `key.json` file created when running
@@ -24,6 +25,9 @@ pub(crate) enum FlexibleCredentialSource {
     // This credential parses the `~/.config/gcloud/application_default_credentials.json` file
     // created when running `gcloud auth application-default login`
     AuthorizedUser(UserCredentials),
+    // This credential parses the `~/.config/gcloud/application_default_credentials.json` file
+    // created when running `gcloud auth application-default login --impersonate-service-account <service account>`
+    ImpersonatedServiceAccount(ImpersonatedServiceAccountCredentials),
 }
 
 impl FlexibleCredentialSource {
@@ -62,6 +66,30 @@ impl FlexibleCredentialSource {
                     ConfigDefaultCredentials::from_user_credentials(creds, client).await?;
                 Ok(Box::new(service_account))
             }
+            FlexibleCredentialSource::ImpersonatedServiceAccount(creds) => {
+                let source_creds: Box<dyn ServiceAccount> = match *creds.source_credentials {
+                    FlexibleCredentialSource::AuthorizedUser(creds) => {
+                        let service_account =
+                            ConfigDefaultCredentials::from_user_credentials(creds, client).await?;
+                        Box::new(service_account)
+                    }
+                    FlexibleCredentialSource::ServiceAccount(creds) => {
+                        let service_account = CustomServiceAccount::new(creds)?;
+                        Box::new(service_account)
+                    }
+                    FlexibleCredentialSource::ImpersonatedServiceAccount(_) => {
+                        return Err(Error::NestedImpersonation)
+                    }
+                };
+
+                let service_account = ImpersonatedServiceAccount::new(
+                    source_creds,
+                    creds.service_account_impersonation_url,
+                    creds.delegates,
+                );
+
+                Ok(Box::new(service_account))
+            }
         }
     }
 
@@ -76,6 +104,17 @@ impl FlexibleCredentialSource {
     }
 }
 
+// This credential uses the `source_credentials` to get a token
+// and then uses that token to get a token impersonating the service
+// account specified by `service_account_impersonation_url`.
+// refresh logic https://github.com/golang/oauth2/blob/a835fc4358f6852f50c4c5c33fddcd1adade5b0a/google/internal/externalaccount/impersonate.go#L57
+#[derive(Deserialize, Debug)]
+pub(crate) struct ImpersonatedServiceAccountCredentials {
+    service_account_impersonation_url: String,
+    source_credentials: Box<FlexibleCredentialSource>,
+    delegates: Vec<String>,
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{flexible_credential_source::FlexibleCredentialSource, types};
@@ -192,4 +231,67 @@ mod tests {
             );
         }
     }
+
+    #[tokio::test]
+    async fn test_parse_impersonating_service_account() {
+        let impersonate_from_user_creds = r#"{
+            "delegates": [],
+            "service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/test_account@test_project.iam.gserviceaccount.com:generateAccessToken",
+            "source_credentials": {
+                "client_id": "***id***.apps.googleusercontent.com",
+                "client_secret": "***secret***",
+                "refresh_token": "***refresh***",
+                "type": "authorized_user",
+                "quota_project_id": "test_project"
+            },
+            "type": "impersonated_service_account"
+        }"#;
+
+        let cred_source: FlexibleCredentialSource =
+            serde_json::from_str(impersonate_from_user_creds).expect("Valid creds to parse");
+
+        assert!(matches!(
+            cred_source,
+            FlexibleCredentialSource::ImpersonatedServiceAccount(_)
+        ));
+
+        let impersonate_from_service_key = r#"{
+            "delegates": [],
+            "service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/test_account@test_project.iam.gserviceaccount.com:generateAccessToken",
+            "source_credentials": {
+                "private_key_id": "268f54e43a1af97cfc71731688434f45aca15c8b",
+                "private_key": "-----BEGIN PRIVATE KEY-----\nMIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC5M5y3WwsRk8NX\npF9fKaZukNspot9Ecmk1PAkupcHLKVhalwPxU4sMNWXgM9H2LTWSvvyOT//rDQpn\n3SGYri/lMhzb4lI8h10E7k6zyFQUPujxkXFBkMOzhIDUgtiiht0WvIw6M8nbaPqI\nxn/aYmPsFhvJfKCthYAt2UUz+D3enI9QjCuhic8iSMnvKT8m0QkOG2eALYGUaLF1\ngRkbV4BiBUGZfXfNEBdux3Wf4kNUau32LA0XotomlvNvf1oH77v5Hc1R/KMMIk5F\nJWVBuAr4jwkN9hwtOozpJ/52wSpddxsZuj+0nP1a3f0UyvrmMnuwszardPK39BoH\nJ+5+HZM3AgMBAAECggEADrHZrXK73hkrVrjkGFjlq8Ayo4sYzAWH84Ff+SONzODq\n8cUpuuw2DDHwc2mpLy9HIO2mfGQ8mhneyX7yO3sWscjYIVpDzCmxZ8LA2+L5SOH0\n+bXglqM14/iPgE0hg0PQJw2u0q9pRM9/kXquilVkOEdIzSPmW95L3Vdv9j+sKQ2A\nOL23l4dsaG4+i1lWRBKiGsLh1kB9FRnm4BzcOxd3WGooy7L1/jo9BoYRss1YABls\nmmyZ9f7r28zjclhpOBkE3OXX0zNbp4yIu1O1Bt9X2p87EOuYqlFA5eEvDbiTPZbk\n6wKEX3BPUkeIo8OaGvsGhHCWx0lv/sDPw/UofycOgQKBgQD4BD059aXEV13Byc5D\nh8LQSejjeM/Vx+YeCFI66biaIOvUs+unyxkH+qxXTuW6AgOgcvrJo93xkyAZ9SeR\nc6Vj9g5mZ5vqSJz5Hg8h8iZBAYtf40qWq0pHcmUIm2Z9LvrG5ZFHU5EEcCtLyBVS\nAv+pLLLf3OsAkJuuqTAgygBbOwKBgQC/KcBa9sUg2u9qIpq020UOW/n4KFWhSJ8h\ngXqqmjOnPqmDc5AnYg1ZdYdqSSgdiK8lJpRL/S2UjYUQp3H+56z0eK/b1iKM51n+\n6D80nIxWeKJ+n7VKI7cBXwc/KokaXgkz0It2UEZSlhPUMImnYcOvGIZ7cMr3Q6mf\n6FwD15UQNQKBgQDyAsDz454DvvS/+noJL1qMAPL9tI+pncwQljIXRqVZ0LIO9hoH\nu4kLXjH5aAWGwhxj3o6VYA9cgSIb8jrQFbbXmexnRMbBkGWMOSavCykE2cr0oEfS\nSgbLPPcVtP4HPWZ72tsubH7fg8zbv7v+MOrkW7eX9mxiOrmPb4yFElfSrQKBgA7y\nMLvr91WuSHG/6uChFDEfN9gTLz7A8tAn03NrQwace5xveKHbpLeN3NyOg7hra2Y4\nMfgO/3VR60l2Dg+kBX3HwdgqUeE6ZWrstaRjaQWJwQqtafs196T/zQ0/QiDxoT6P\n25eQhy8F1N8OPHT9y9Lw0/LqyrOycpyyCh+yx1DRAoGAJ/6dlhyQnwSfMAe3mfRC\noiBQG6FkyoeXHHYcoQ/0cSzwp0BwBlar1Z28P7KTGcUNqV+YfK9nF47eoLaTLCmG\nG5du0Ds6m2Eg0sOBBqXHnw6R1PC878tgT/XokNxIsVlF5qRz88q7Rn0J1lzB7+Tl\n2HSAcyIUcmr0gxlhRmC2Jq4=\n-----END PRIVATE KEY-----\n",
+                "client_email": "gopher@developer.gserviceaccount.com",
+                "client_id": "gopher.apps.googleusercontent.com",
+                "token_uri": "https://accounts.google.com/o/gophers/token",
+                "type": "service_account",
+                "audience": "https://testservice.googleapis.com/",
+                "project_id": "test_project"
+            },
+            "type": "impersonated_service_account"
+        }"#;
+
+        let cred_source: FlexibleCredentialSource =
+            serde_json::from_str(impersonate_from_service_key).expect("Valid creds to parse");
+
+        assert!(matches!(
+            cred_source,
+            FlexibleCredentialSource::ImpersonatedServiceAccount(_)
+        ));
+
+        let client = types::client();
+        let creds = cred_source
+            .try_into_service_account(&client)
+            .await
+            .expect("Valid creds to parse");
+
+        assert_eq!(
+            creds
+                .project_id(&client)
+                .await
+                .expect("Project ID to be present"),
+            "test_project".to_string(),
+            "Project ID should be parsed"
+        );
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -97,6 +97,7 @@ mod error;
 mod flexible_credential_source;
 mod gcloud_authorized_user;
 mod jwt;
+mod service_account_impersonation;
 mod types;
 mod util;
 
diff --git a/src/service_account_impersonation.rs b/src/service_account_impersonation.rs
@@ -0,0 +1,127 @@
+use async_trait::async_trait;
+use std::{collections::HashMap, sync::RwLock};
+
+use hyper::header;
+use serde::Serialize;
+
+use crate::{
+    authentication_manager::ServiceAccount, gcloud_authorized_user::DEFAULT_TOKEN_DURATION,
+    types::HyperClient, util::HyperExt, Error, Token,
+};
+
+// This credential uses the `source_credentials` to get a token
+// and then uses that token to get a token impersonating the service
+// account specified by `service_account_impersonation_url`.
+// refresh logic https://github.com/golang/oauth2/blob/a835fc4358f6852f50c4c5c33fddcd1adade5b0a/google/internal/externalaccount/impersonate.go#L57
+//
+// In practice, the api currently referred to by `service_account_impersonation_url` is
+// https://cloud.google.com/iam/docs/reference/credentials/rest/v1/projects.serviceAccounts/generateAccessToken
+pub(crate) struct ImpersonatedServiceAccount {
+    service_account_impersonation_url: String,
+    source_credentials: Box<dyn ServiceAccount>,
+    delegates: Vec<String>,
+    tokens: RwLock<HashMap<Vec<String>, Token>>,
+}
+
+impl ImpersonatedServiceAccount {
+    pub(crate) fn new(
+        source_credentials: Box<dyn ServiceAccount>,
+        service_account_impersonation_url: String,
+        delegates: Vec<String>,
+    ) -> Self {
+        Self {
+            service_account_impersonation_url,
+            source_credentials,
+            delegates,
+            tokens: RwLock::new(HashMap::new()),
+        }
+    }
+}
+
+impl std::fmt::Debug for ImpersonatedServiceAccount {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ImpersonatedServiceAccount")
+            .field(
+                "service_account_impersonation_url",
+                &self.service_account_impersonation_url,
+            )
+            .field("source_credentials", &"Box<dyn ServiceAccount>")
+            .field("delegates", &self.delegates)
+            .finish()
+    }
+}
+
+#[async_trait]
+impl ServiceAccount for ImpersonatedServiceAccount {
+    async fn project_id(&self, hc: &HyperClient) -> Result<String, Error> {
+        self.source_credentials.project_id(hc).await
+    }
+
+    fn get_token(&self, scopes: &[&str]) -> Option<Token> {
+        let key: Vec<_> = scopes.iter().map(|x| x.to_string()).collect();
+        self.tokens.read().unwrap().get(&key).cloned()
+    }
+
+    async fn refresh_token(&self, client: &HyperClient, scopes: &[&str]) -> Result<Token, Error> {
+        let source_token = self
+            .source_credentials
+            .refresh_token(client, scopes)
+            .await?;
+
+        // Then we do a request to get the impersonated token
+        let lifetime_seconds = DEFAULT_TOKEN_DURATION.whole_seconds();
+        #[derive(Serialize, Clone)]
+        // Format from https://github.com/golang/oauth2/blob/a835fc4358f6852f50c4c5c33fddcd1adade5b0a/google/internal/externalaccount/impersonate.go#L21
+        struct AccessTokenRequest {
+            lifetime: String,
+            #[serde(skip_serializing_if = "Vec::is_empty")]
+            scope: Vec<String>,
+            #[serde(skip_serializing_if = "Vec::is_empty")]
+            delegates: Vec<String>,
+        }
+
+        let request = AccessTokenRequest {
+            lifetime: format!("{lifetime_seconds}s"),
+            scope: scopes.iter().map(|s| s.to_string()).collect(),
+            delegates: self.delegates.clone(),
+        };
+        let rqbody =
+            serde_json::to_string(&request).expect("access token request failed to serialize");
+
+        let token_uri = self.service_account_impersonation_url.as_str();
+
+        let mut retries = 0;
+        let response = loop {
+            // We assume bearer tokens only. In the referenced code, other token types are possible
+            // https://github.com/golang/oauth2/blob/a835fc4358f6852f50c4c5c33fddcd1adade5b0a/token.go#L84
+            let request = hyper::Request::post(token_uri)
+                .header(
+                    header::AUTHORIZATION,
+                    format!("Bearer {}", source_token.as_str()),
+                )
+                .header(header::CONTENT_TYPE, "application/json")
+                .body(hyper::Body::from(rqbody.clone()))
+                .unwrap();
+
+            tracing::debug!("requesting impersonation token from service account: {request:?}");
+            let err = match client.request(request).await {
+                // Early return when the request succeeds
+                Ok(response) => break response,
+                Err(err) => err,
+            };
+
+            tracing::warn!(
+                "Failed to refresh impersonation token with service token endpoint {token_uri}: {err}, trying again..."
+            );
+            retries += 1;
+            if retries >= RETRY_COUNT {
+                return Err(Error::OAuthConnectionError(err));
+            }
+        };
+
+        response.deserialize().await
+    }
+}
+
+/// How many times to attempt to fetch a token from the set credentials token endpoint.
+const RETRY_COUNT: u8 = 5;
