Skip to content

Commit

Permalink
feat(storage): Implement DiskTokenStorage
Browse files Browse the repository at this point in the history
DiskTokenStorage is a TokenStorage that stores its tokens in a JSON file
on disk. That file can be read in later, and the tokens in it reused.
(The idea for a cache file is from here:
https://developers.google.com/drive/v3/web/quickstart/go)
  • Loading branch information
dermesser committed Apr 15, 2016
1 parent 2aa95c0 commit 2cb5250
Show file tree
Hide file tree
Showing 3 changed files with 205 additions and 68 deletions.
66 changes: 2 additions & 64 deletions src/helper.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,75 +11,12 @@ use std::convert::From;
use common::{Token, FlowType, ApplicationSecret};
use device::{PollInformation, RequestError, DeviceFlow, PollError};
use refresh::{RefreshResult, RefreshFlow};
use storage::{TokenStorage};
use chrono::{DateTime, UTC, Local};
use std::time::Duration;
use hyper;


/// Implements a specialized storage to set and retrieve `Token` instances.
/// The `scope_hash` represents the signature of the scopes for which the given token
/// should be stored or retrieved.
/// For completeness, the underlying, sorted scopes are provided as well. They might be
/// useful for presentation to the user.
pub trait TokenStorage {
type Error: 'static + Error;

/// If `token` is None, it is invalid or revoked and should be removed from storage.
/// Otherwise, it should be saved.
fn set(&mut self, scope_hash: u64, scopes: &Vec<&str>, token: Option<Token>) -> Result<(), Self::Error>;
/// A `None` result indicates that there is no token for the given scope_hash.
fn get(&self, scope_hash: u64, scopes: &Vec<&str>) -> Result<Option<Token>, Self::Error>;
}

/// A storage that remembers nothing.
#[derive(Default)]
pub struct NullStorage;
#[derive(Debug)]
pub struct NullError;

impl Error for NullError {
fn description(&self) -> &str {
"NULL"
}
}

impl fmt::Display for NullError {
fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
"NULL-ERROR".fmt(f)
}
}

impl TokenStorage for NullStorage {
type Error = NullError;
fn set(&mut self, _: u64, _: &Vec<&str>, _: Option<Token>) -> Result<(), NullError> { Ok(()) }
fn get(&self, _: u64, _: &Vec<&str>) -> Result<Option<Token>, NullError> { Ok(None) }
}

/// A storage that remembers values for one session only.
#[derive(Default)]
pub struct MemoryStorage {
pub tokens: HashMap<u64, Token>
}

impl TokenStorage for MemoryStorage {
type Error = NullError;

fn set(&mut self, scope_hash: u64, _: &Vec<&str>, token: Option<Token>) -> Result<(), NullError> {
match token {
Some(t) => self.tokens.insert(scope_hash, t),
None => self.tokens.remove(&scope_hash),
};
Ok(())
}

fn get(&self, scope_hash: u64, _: &Vec<&str>) -> Result<Option<Token>, NullError> {
match self.tokens.get(&scope_hash) {
Some(t) => Ok(Some(t.clone())),
None => Ok(None),
}
}
}

/// A generalized authenticator which will keep tokens valid and store them.
///
/// It is the go-to helper to deal with any kind of supported authentication flow,
Expand Down Expand Up @@ -477,6 +414,7 @@ mod tests {
use super::super::device::tests::MockGoogleAuth;
use super::super::common::tests::SECRET;
use super::super::common::{ConsoleApplicationSecret};
use storage::MemoryStorage;
use std::default::Default;
use hyper;

Expand Down
9 changes: 5 additions & 4 deletions src/lib.rs.in
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,14 @@ extern crate mime;
extern crate url;
extern crate itertools;

mod device;
mod refresh;
mod common;
mod device;
mod helper;
mod refresh;
mod storage;

pub use device::{DeviceFlow, PollInformation, PollError};
pub use refresh::{RefreshFlow, RefreshResult};
pub use common::{Token, FlowType, ApplicationSecret, ConsoleApplicationSecret, Scheme, TokenType};
pub use helper::{TokenStorage, NullStorage, MemoryStorage, Authenticator,
AuthenticatorDelegate, Retry, DefaultAuthenticatorDelegate, GetToken};
pub use helper::{Authenticator, AuthenticatorDelegate, Retry, DefaultAuthenticatorDelegate, GetToken};
pub use storage::{TokenStorage, DiskTokenStorage, NullStorage, MemoryStorage};
198 changes: 198 additions & 0 deletions src/storage.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,198 @@
/*
* partially (c) 2016 Google Inc. (Lewin Bormann, [email protected])
*
* See project root for licensing information.
*/

extern crate serde_json;

use std::collections::HashMap;
use std::hash::Hasher;
use std::error::Error;
use std::fmt;
use std::fs;
use std::io;
use std::io::{Read, Write};

use common::Token;

/// Implements a specialized storage to set and retrieve `Token` instances.
/// The `scope_hash` represents the signature of the scopes for which the given token
/// should be stored or retrieved.
/// For completeness, the underlying, sorted scopes are provided as well. They might be
/// useful for presentation to the user.
pub trait TokenStorage {
type Error: 'static + Error;

/// If `token` is None, it is invalid or revoked and should be removed from storage.
/// Otherwise, it should be saved.
fn set(&mut self,
scope_hash: u64,
scopes: &Vec<&str>,
token: Option<Token>)
-> Result<(), Self::Error>;
/// A `None` result indicates that there is no token for the given scope_hash.
fn get(&self, scope_hash: u64, scopes: &Vec<&str>) -> Result<Option<Token>, Self::Error>;
}

/// A storage that remembers nothing.
#[derive(Default)]
pub struct NullStorage;
#[derive(Debug)]
pub struct NullError;

impl Error for NullError {
fn description(&self) -> &str {
"NULL"
}
}

impl fmt::Display for NullError {
fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
"NULL-ERROR".fmt(f)
}
}

impl TokenStorage for NullStorage {
type Error = NullError;
fn set(&mut self, _: u64, _: &Vec<&str>, _: Option<Token>) -> Result<(), NullError> {
Ok(())
}
fn get(&self, _: u64, _: &Vec<&str>) -> Result<Option<Token>, NullError> {
Ok(None)
}
}

/// A storage that remembers values for one session only.
#[derive(Default)]
pub struct MemoryStorage {
pub tokens: HashMap<u64, Token>,
}

impl TokenStorage for MemoryStorage {
type Error = NullError;

fn set(&mut self,
scope_hash: u64,
_: &Vec<&str>,
token: Option<Token>)
-> Result<(), NullError> {
match token {
Some(t) => self.tokens.insert(scope_hash, t),
None => self.tokens.remove(&scope_hash),
};
Ok(())
}

fn get(&self, scope_hash: u64, _: &Vec<&str>) -> Result<Option<Token>, NullError> {
match self.tokens.get(&scope_hash) {
Some(t) => Ok(Some(t.clone())),
None => Ok(None),
}
}
}

/// A single stored token.
#[derive(Serialize, Deserialize)]
struct JSONToken {
pub hash: u64,
pub token: Token,
}

/// List of tokens in a JSON object
#[derive(Serialize, Deserialize)]
struct JSONTokens {
pub tokens: Vec<JSONToken>,
}

/// Serializes tokens to a JSON file on disk.
#[derive(Default)]
pub struct DiskTokenStorage {
location: String,
tokens: HashMap<u64, Token>,
}

impl DiskTokenStorage {
pub fn new(location: &String) -> DiskTokenStorage {
use std::fs;
let mut dts = DiskTokenStorage {
location: location.clone(),
tokens: HashMap::new(),
};

// best-effort
let _ = dts.load_from_file();
dts
}

fn load_from_file(&mut self) -> Result<(), io::Error> {
let mut f = try!(fs::OpenOptions::new().read(true).open(&self.location));
let mut contents = String::new();

match f.read_to_string(&mut contents) {
Result::Err(e) => return Result::Err(e),
Result::Ok(_sz) => (),
}

let tokens: JSONTokens;

match serde_json::from_str(&contents) {
Result::Err(e) => return Result::Err(io::Error::new(io::ErrorKind::InvalidData, e)),
Result::Ok(t) => tokens = t,
}

for t in tokens.tokens {
self.tokens.insert(t.hash, t.token);
}
return Result::Ok(());
}

pub fn dump_to_file(&mut self) -> Result<(), io::Error> {
let mut jsontokens = JSONTokens { tokens: Vec::new() };

for (hash, token) in self.tokens.iter() {
jsontokens.tokens.push(JSONToken {
hash: *hash,
token: token.clone(),
});
}

let serialized;;

match serde_json::to_string(&jsontokens) {
Result::Err(e) => return Result::Err(io::Error::new(io::ErrorKind::InvalidData, e)),
Result::Ok(s) => serialized = s,
}

let mut f = try!(fs::OpenOptions::new()
.create(true)
.write(true)
.truncate(true)
.open(&self.location));
f.write(serialized.as_ref()).map(|_| ())
}
}

impl TokenStorage for DiskTokenStorage {
type Error = io::Error;
fn set(&mut self,
scope_hash: u64,
_: &Vec<&str>,
token: Option<Token>)
-> Result<(), Self::Error> {
match token {
None => {
self.tokens.remove(&scope_hash);
()
}
Some(t) => {
self.tokens.insert(scope_hash, t.clone());
()
}
}
self.dump_to_file()
}
fn get(&self, scope_hash: u64, _: &Vec<&str>) -> Result<Option<Token>, Self::Error> {
Result::Ok(self.tokens.get(&scope_hash).map(|tok| tok.clone()))
}
}

0 comments on commit 2cb5250

Please sign in to comment.