Merge pull request #146 from djrodgerspryor/custom_storage_options

Custom token storage

Cargo.toml

@@ -29,6 +29,9 @@ tokio = { version = "1.0", features = ["fs", "macros", "io-std", "io-util", "tim
 url = "2"
 percent-encoding = "2"
 futures = "0.3"
+async-trait = "^0.1"
+anyhow = "1.0.38"
+itertools = "0.10.0"
 
 [dev-dependencies]
 httptest = "0.14"

examples/custom_storage.rs

@@ -0,0 +1,90 @@
+//! Demonstrating how to create a custom token store
+use anyhow::anyhow;
+use async_trait::async_trait;
+use std::sync::RwLock;
+use yup_oauth2::storage::{TokenInfo, TokenStorage};
+
+struct ExampleTokenStore {
+    store: RwLock<Vec<StoredToken>>,
+}
+
+struct StoredToken {
+    scopes: Vec<String>,
+    serialized_token: String,
+}
+
+/// Is this set of scopes covered by the other? Returns true if the other
+/// set is a superset of this one. Use this when implementing TokenStorage.get()
+fn scopes_covered_by(scopes: &[&str], possible_match_or_superset: &[&str]) -> bool {
+    scopes
+        .iter()
+        .all(|s| possible_match_or_superset.iter().any(|t| t == s))
+}
+
+/// Here we implement our own token storage. You could write the serialized token and scope data
+/// to disk, an OS keychain, a database or whatever suits your use-case
+#[async_trait]
+impl TokenStorage for ExampleTokenStore {
+    async fn set(&self, scopes: &[&str], token: TokenInfo) -> anyhow::Result<()> {
+        let data = serde_json::to_string(&token).unwrap();
+
+        println!("Storing token for scopes {:?}", scopes);
+
+        let mut store = self
+            .store
+            .write()
+            .map_err(|_| anyhow!("Unable to lock store for writing"))?;
+
+        store.push(StoredToken {
+            scopes: scopes.iter().map(|str| str.to_string()).collect(),
+            serialized_token: data,
+        });
+
+        Ok(())
+    }
+
+    async fn get(&self, target_scopes: &[&str]) -> Option<TokenInfo> {
+        // Retrieve the token data
+        self.store.read().ok().and_then(|store| {
+            for stored_token in store.iter() {
+                if scopes_covered_by(
+                    target_scopes,
+                    &stored_token
+                        .scopes
+                        .iter()
+                        .map(|s| &s[..])
+                        .collect::<Vec<_>>()[..],
+                ) {
+                    return serde_json::from_str(&stored_token.serialized_token).ok();
+                }
+            }
+
+            None
+        })
+    }
+}
+
+#[tokio::main]
+async fn main() {
+    // Put your client secret in the working directory!
+    let sec = yup_oauth2::read_application_secret("client_secret.json")
+        .await
+        .expect("client secret couldn't be read.");
+    let auth = yup_oauth2::InstalledFlowAuthenticator::builder(
+        sec,
+        yup_oauth2::InstalledFlowReturnMethod::HTTPRedirect,
+    )
+    .with_storage(Box::new(ExampleTokenStore {
+        store: RwLock::new(vec![]),
+    }))
+    .build()
+    .await
+    .expect("InstalledFlowAuthenticator failed to build");
+
+    let scopes = &["https://www.googleapis.com/auth/drive.file"];
+
+    match auth.token(scopes).await {
+        Err(e) => println!("error: {:?}", e),
+        Ok(t) => println!("The token is {:?}", t),
+    }
+}

rustfmt.toml

@@ -0,0 +1 @@
+edition = "2018"

src/authenticator.rs

@@ -5,7 +5,7 @@ use crate::error::Error;
 use crate::installed::{InstalledFlow, InstalledFlowReturnMethod};
 use crate::refresh::RefreshFlow;
 use crate::service_account::{ServiceAccountFlow, ServiceAccountFlowOpts, ServiceAccountKey};
-use crate::storage::{self, Storage};
+use crate::storage::{self, Storage, TokenStorage};
 use crate::types::{AccessToken, ApplicationSecret, TokenInfo};
 use private::AuthFlow;
 
@@ -238,6 +238,7 @@ impl<C, F> AuthenticatorBuilder<C, F> {
                 tokens: Mutex::new(storage::JSONTokens::new()),
             },
             StorageType::Disk(path) => Storage::Disk(storage::DiskStorage::new(path).await?),
+            StorageType::Custom(custom_store) => Storage::Custom(custom_store),
         };
 
         Ok(Authenticator {
@@ -257,6 +258,14 @@ impl<C, F> AuthenticatorBuilder<C, F> {
         }
     }
 
+    /// Use the provided token storage mechanism
+    pub fn with_storage(self, storage: Box<dyn TokenStorage>) -> Self {
+        AuthenticatorBuilder {
+            storage_type: StorageType::Custom(storage),
+            ..self
+        }
+    }
+
     /// Use the provided hyper client.
     pub fn hyper_client<NewC>(
         self,
@@ -515,9 +524,14 @@ where
     }
 }
 
+/// How should the acquired tokens be stored?
 enum StorageType {
+    /// Store tokens in memory (and always log in again to acquire a new token on startup)
     Memory,
+    /// Store tokens to disk in the given file. Warning, this may be insecure unless you configure your operating system to restrict read access to the file.
     Disk(PathBuf),
+    /// Implement your own storage provider
+    Custom(Box<dyn TokenStorage>),
 }
 
 #[cfg(test)]

src/error.rs

@@ -153,6 +153,8 @@ pub enum Error {
     UserError(String),
     /// A lower level IO error.
     LowLevelError(io::Error),
+    /// Other errors produced by a storage provider
+    OtherError(anyhow::Error),
 }
 
 impl From<hyper::Error> for Error {
@@ -179,6 +181,15 @@ impl From<io::Error> for Error {
     }
 }
 
+impl From<anyhow::Error> for Error {
+    fn from(value: anyhow::Error) -> Error {
+        match value.downcast::<io::Error>() {
+            Ok(io_error) => Error::LowLevelError(io_error),
+            Err(err) => Error::OtherError(err),
+        }
+    }
+}
+
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
         match *self {
@@ -194,6 +205,7 @@ impl fmt::Display for Error {
             }
             Error::UserError(ref s) => s.fmt(f),
             Error::LowLevelError(ref e) => e.fmt(f),
+            Error::OtherError(ref e) => e.fmt(f),
         }
     }
 }

src/lib.rs

@@ -77,7 +77,11 @@ mod helper;
 mod installed;
 mod refresh;
 mod service_account;
-mod storage;
+
+/// Interface for storing tokens so that they can be re-used. There are built-in memory and
+/// file-based storage providers. You can implement your own by implementing the TokenStorage trait.
+pub mod storage;
+
 mod types;
 
 #[doc(inline)]

src/storage.rs

@@ -2,13 +2,16 @@
 //
 // See project root for licensing information.
 //
-use crate::types::TokenInfo;
+pub use crate::types::TokenInfo;
 
 use futures::lock::Mutex;
+use itertools::Itertools;
 use std::collections::HashMap;
 use std::io;
 use std::path::{Path, PathBuf};
 
+use async_trait::async_trait;
+
 use serde::{Deserialize, Serialize};
 
 // The storage layer allows retrieving tokens for scopes that have been
@@ -49,6 +52,7 @@ impl ScopeFilter {
     }
 }
 
+/// A set of scopes
 #[derive(Debug)]
 pub(crate) struct ScopeSet<'a, T> {
     hash: ScopeHash,
@@ -73,6 +77,8 @@ impl<'a, T> ScopeSet<'a, T>
 where
     T: AsRef<str>,
 {
+    /// Convert from an array into a ScopeSet. Automatically invoked by the compiler when
+    /// an array reference is passed.
     // implement an inherent from method even though From is implemented. This
     // is because passing an array ref like &[&str; 1] (&["foo"]) will be auto
     // deref'd to a slice on function boundaries, but it will not implement the
@@ -105,9 +111,22 @@ where
     }
 }
 
+/// Implement your own token storage solution by implementing this trait. You need a way to
+/// store and retrieve tokens, each keyed by a set of scopes.
+#[async_trait]
+pub trait TokenStorage: Send + Sync {
+    /// Store a token for the given set of scopes so that it can be retrieved later by get()
+    /// TokenInfo can be serialized with serde.
+    async fn set(&self, scopes: &[&str], token: TokenInfo) -> anyhow::Result<()>;
+
+    /// Retrieve a token stored by set for the given set of scopes
+    async fn get(&self, scopes: &[&str]) -> Option<TokenInfo>;
+}
+
 pub(crate) enum Storage {
     Memory { tokens: Mutex<JSONTokens> },
     Disk(DiskStorage),
+    Custom(Box<dyn TokenStorage>),
 }
 
 impl Storage {
@@ -115,13 +134,29 @@ impl Storage {
         &self,
         scopes: ScopeSet<'_, T>,
         token: TokenInfo,
-    ) -> Result<(), io::Error>
+    ) -> anyhow::Result<()>
     where
         T: AsRef<str>,
     {
         match self {
-            Storage::Memory { tokens } => tokens.lock().await.set(scopes, token),
-            Storage::Disk(disk_storage) => disk_storage.set(scopes, token).await,
+            Storage::Memory { tokens } => Ok(tokens.lock().await.set(scopes, token)?),
+            Storage::Disk(disk_storage) => Ok(disk_storage.set(scopes, token).await?),
+            Storage::Custom(custom_storage) => {
+                let str_scopes: Vec<_> = scopes
+                    .scopes
+                    .iter()
+                    .map(|scope| scope.as_ref())
+                    .sorted()
+                    .unique()
+                    .collect();
+
+                custom_storage
+                    .set(
+                        &str_scopes[..], // TODO: sorted, unique
+                        token,
+                    )
+                    .await
+            }
         }
     }
 
@@ -132,6 +167,17 @@ impl Storage {
         match self {
             Storage::Memory { tokens } => tokens.lock().await.get(scopes),
             Storage::Disk(disk_storage) => disk_storage.get(scopes).await,
+            Storage::Custom(custom_storage) => {
+                let str_scopes: Vec<_> = scopes
+                    .scopes
+                    .iter()
+                    .map(|scope| scope.as_ref())
+                    .sorted()
+                    .unique()
+                    .collect();
+
+                custom_storage.get(&str_scopes[..]).await
+            }
         }
     }
 }

src/types.rs

@@ -56,13 +56,13 @@ impl From<TokenInfo> for AccessToken {
 /// It authenticates certain operations, and must be refreshed once
 /// it reached it's expiry date.
 #[derive(Clone, PartialEq, Debug, Deserialize, Serialize)]
-pub(crate) struct TokenInfo {
+pub struct TokenInfo {
     /// used when authenticating calls to oauth2 enabled services.
-    pub(crate) access_token: String,
+    pub access_token: String,
     /// used to refresh an expired access_token.
-    pub(crate) refresh_token: Option<String>,
+    pub refresh_token: Option<String>,
     /// The time when the token expires.
-    pub(crate) expires_at: Option<DateTime<Utc>>,
+    pub expires_at: Option<DateTime<Utc>>,
 }
 
 impl TokenInfo {