From 59d266a1113f1b1452ffbc5644cf46ed79106e4e Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Fri, 17 Jul 2020 13:31:10 -0700
Subject: [PATCH] Add access_token module

---
 storage-bigtable/Cargo.toml          |  3 +
 storage-bigtable/src/access_token.rs | 91 ++++++++++++++++++++++++++++
 storage-bigtable/src/lib.rs          |  1 +
 3 files changed, 95 insertions(+)
 create mode 100644 storage-bigtable/src/access_token.rs

diff --git a/storage-bigtable/Cargo.toml b/storage-bigtable/Cargo.toml
index c4882527f..12730855d 100644
--- a/storage-bigtable/Cargo.toml
+++ b/storage-bigtable/Cargo.toml
@@ -9,6 +9,9 @@ homepage = "https://solana.com/"
 edition = "2018"
 
 [dependencies]
+log = "0.4.8"
+goauth = "0.7.1"
+smpl_jwt = "0.5.0"
 
 [lib]
 crate-type = ["lib"]
diff --git a/storage-bigtable/src/access_token.rs b/storage-bigtable/src/access_token.rs
new file mode 100644
index 000000000..30eeb7932
--- /dev/null
+++ b/storage-bigtable/src/access_token.rs
@@ -0,0 +1,91 @@
+/// A module for managing a Google API access token
+use goauth::{
+    auth::{JwtClaims, Token},
+    credentials::Credentials,
+    get_token,
+};
+use log::*;
+use smpl_jwt::Jwt;
+use std::time::Instant;
+
+pub use goauth::scopes::Scope;
+
+fn load_credentials() -> Result<Credentials, String> {
+    // Use standard GOOGLE_APPLICATION_CREDENTIALS environment variable
+    let credentials_file = std::env::var("GOOGLE_APPLICATION_CREDENTIALS")
+        .map_err(|_| "GOOGLE_APPLICATION_CREDENTIALS environment variable not found".to_string())?;
+
+    Credentials::from_file(&credentials_file).map_err(|err| {
+        format!(
+            "Failed to read GCP credentials from {}: {}",
+            credentials_file, err
+        )
+    })
+}
+
+pub struct AccessToken {
+    credentials: Credentials,
+    jwt: Jwt<JwtClaims>,
+    token: Option<(Token, Instant)>,
+}
+
+impl AccessToken {
+    pub fn new(scope: &Scope) -> Result<Self, String> {
+        let credentials = load_credentials()?;
+
+        let claims = JwtClaims::new(
+            credentials.iss(),
+            &scope,
+            credentials.token_uri(),
+            None,
+            None,
+        );
+        let jwt = Jwt::new(
+            claims,
+            credentials
+                .rsa_key()
+                .map_err(|err| format!("Invalid rsa key: {}", err))?,
+            None,
+        );
+
+        Ok(Self {
+            credentials,
+            jwt,
+            token: None,
+        })
+    }
+
+    /// The project that this token grants access to
+    pub fn project(&self) -> String {
+        self.credentials.project()
+    }
+
+    /// Call this function regularly, and before calling `access_token()`
+    pub async fn refresh(&mut self) {
+        if let Some((token, last_refresh)) = self.token.as_ref() {
+            if last_refresh.elapsed().as_secs() < token.expires_in() as u64 / 2 {
+                return;
+            }
+        }
+
+        info!("Refreshing token");
+        match get_token(&self.jwt, &self.credentials).await {
+            Ok(new_token) => {
+                info!("Token expires in {} seconds", new_token.expires_in());
+                self.token = Some((new_token, Instant::now()));
+            }
+            Err(err) => {
+                warn!("Failed to get new token: {}", err);
+            }
+        }
+    }
+
+    /// Return an access token suitable for use in an HTTP authorization header
+    pub fn get(&self) -> Result<String, String> {
+        if let Some((token, _)) = self.token.as_ref() {
+            Ok(format!("{} {}", token.token_type(), token.access_token()))
+        } else {
+            Err("Access token not available".into())
+        }
+    }
+}
diff --git a/storage-bigtable/src/lib.rs b/storage-bigtable/src/lib.rs
index e69de29bb..5d8f47dfd 100644
--- a/storage-bigtable/src/lib.rs
+++ b/storage-bigtable/src/lib.rs
@@ -0,0 +1 @@
+mod access_token;