Custom token storage

Allow users to build their own token storage system by implementing the `TokenStorage` trait. This allows use of more secure storage mechanisms like OS keychains, encrypted files, or secret-management tools.

Custom storage providers are Box-ed to avoid adding more generics to the API — the indirection cost will only apply if using a custom store.

I've added `anyhow` to allow easy handling of a wide range of errors from custom storage providers.

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;
 
@@ -47,7 +47,7 @@ where
     C: hyper::client::connect::Connect + Clone + Send + Sync + 'static,
 {
     /// Return the current token for the provided scopes.
-    pub async fn token<'a, T>(&'a self, scopes: &'a [T]) -> Result<AccessToken, Error>
+    pub async fn token<'a, T>(&'a mut self, scopes: &'a [T]) -> Result<AccessToken, Error>
     where
         T: AsRef<str>,
     {
@@ -57,7 +57,7 @@ where
     /// Return a token for the provided scopes, but don't reuse cached tokens. Instead,
     /// always fetch a new token from the OAuth server.
     pub async fn force_refreshed_token<'a, T>(
-        &'a self,
+        &'a mut self,
         scopes: &'a [T],
     ) -> Result<AccessToken, Error>
     where
@@ -68,7 +68,7 @@ where
 
     /// Return a cached token or fetch a new one from the server.
     async fn find_token<'a, T>(
-        &'a self,
+        &'a mut self,
         scopes: &'a [T],
         force_refresh: bool,
     ) -> Result<AccessToken, Error>
@@ -219,6 +219,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 {
@@ -236,6 +237,14 @@ impl<C, F> AuthenticatorBuilder<C, F> {
         }
     }
 
+    /// Use the provided token storage mechanism
+    pub fn with_storage(self, storage_type: StorageType) -> Self {
+        AuthenticatorBuilder {
+            storage_type: storage_type,
+            ..self
+        }
+    }
+
     /// Use the provided hyper client.
     pub fn hyper_client<NewC>(
         self,
@@ -494,9 +503,14 @@ where
     }
 }
 
-enum StorageType {
+/// How should the acquired tokens be stored?
+pub 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,8 +52,9 @@ impl ScopeFilter {
     }
 }
 
+/// A set of scopes
 #[derive(Debug)]
-pub(crate) struct ScopeSet<'a, T> {
+pub struct ScopeSet<'a, T> {
     hash: ScopeHash,
     filter: ScopeFilter,
     scopes: &'a [T],
@@ -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
@@ -103,25 +109,72 @@ where
             scopes,
         }
     }
+
+    /// Get the scopes for storage when implementing TokenStorage.set().
+    /// Returned scope strings are unique and sorted.
+    pub fn scopes(&self) -> Vec<String> {
+        self.scopes
+            .iter()
+            .map(|scope| scope.as_ref().to_string())
+            .sorted()
+            .unique()
+            .collect()
+    }
+
+    /// 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()
+    pub fn is_covered_by(&self, other_scopes: &[String]) -> bool {
+        self.scopes
+            .iter()
+            .all(|s| other_scopes.iter().any(|t| t.as_str() == s.as_ref()))
+    }
+}
+
+/// 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()
+    /// ScopeSet implements Hash so that you can easily serialize and store it.
+    /// TokenInfo can be serialized with serde.
+    async fn set(&mut self, scopes: ScopeSet<'_, &str>, token: TokenInfo) -> anyhow::Result<()>;
+
+    /// Retrieve a token stored by set for the given set of scopes
+    async fn get(&self, scopes: ScopeSet<'_, &str>) -> Option<TokenInfo>;
 }
 
 pub(crate) enum Storage {
     Memory { tokens: Mutex<JSONTokens> },
     Disk(DiskStorage),
+    Custom(Box<dyn TokenStorage>),
 }
 
 impl Storage {
     pub(crate) async fn set<T>(
-        &self,
+        &mut 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()).collect();
+
+                (*custom_storage)
+                    .set(
+                        ScopeSet {
+                            hash: scopes.hash,
+                            filter: scopes.filter,
+                            scopes: &str_scopes[..],
+                        },
+                        token,
+                    )
+                    .await
+            }
         }
     }
 
@@ -132,6 +185,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()).collect();
+
+                (*custom_storage)
+                    .get(ScopeSet {
+                        hash: scopes.hash,
+                        filter: scopes.filter,
+                        scopes: &str_scopes[..],
+                    })
+                    .await
+            }
         }
     }
 }

src/types.rs

@@ -56,7 +56,7 @@ 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,
     /// used to refresh an expired access_token.