use std::collections::{HashMap, HashSet};
use std::cell::RefCell;
use time;
use Cookie;
pub struct CookieJar<'a> {
flavor: Flavor<'a>,
}
enum Flavor<'a> {
Child(Child<'a>),
Root(Root),
}
struct Child<'a> {
parent: &'a CookieJar<'a>,
read: Read,
write: Write,
}
type Read = fn(&Root, Cookie) -> Option<Cookie>;
type Write = fn(&Root, Cookie) -> Cookie;
struct Root {
map: RefCell<HashMap<String, Cookie>>,
new_cookies: RefCell<HashSet<String>>,
removed_cookies: RefCell<HashSet<String>>,
key: Vec<u8>,
}
pub struct Iter<'a> {
jar: &'a CookieJar<'a>,
keys: Vec<String>,
}
impl<'a> CookieJar<'a> {
pub fn new(key: &[u8]) -> CookieJar<'static> {
let normalized_key = if key.len() >= secure::MIN_KEY_LEN {
key.to_vec()
} else {
secure::prepare_key(key)
};
CookieJar {
flavor: Flavor::Root(Root {
map: RefCell::new(HashMap::new()),
new_cookies: RefCell::new(HashSet::new()),
removed_cookies: RefCell::new(HashSet::new()),
key: normalized_key,
})
}
}
fn root<'b>(&'b self) -> &'b Root {
let mut cur = self;
loop {
match cur.flavor {
Flavor::Child(ref child) => cur = child.parent,
Flavor::Root(ref me) => return me,
}
}
}
pub fn add_original(&mut self, cookie: Cookie) {
match self.flavor {
Flavor::Child(..) => panic!("can't add an original cookie to a child jar!"),
Flavor::Root(ref mut root) => {
let name = cookie.name.clone();
root.map.borrow_mut().insert(name, cookie);
}
}
}
pub fn add(&self, mut cookie: Cookie) {
let mut cur = self;
let root = self.root();
loop {
match cur.flavor {
Flavor::Child(ref child) => {
cookie = (child.write)(root, cookie);
cur = child.parent;
}
Flavor::Root(..) => break,
}
}
let name = cookie.name.clone();
root.map.borrow_mut().insert(name.clone(), cookie);
root.removed_cookies.borrow_mut().remove(&name);
root.new_cookies.borrow_mut().insert(name);
}
pub fn remove(&self, cookie: &str) {
let root = self.root();
let cookie = cookie.to_string();
root.map.borrow_mut().remove(&cookie);
root.new_cookies.borrow_mut().remove(&cookie);
root.removed_cookies.borrow_mut().insert(cookie);
}
pub fn find(&self, name: &str) -> Option<Cookie> {
let name = name.to_string();
let root = self.root();
if root.removed_cookies.borrow().contains(&name) {
return None
}
root.map.borrow().get(&name).and_then(|c| self.try_read(root, c.clone()))
}
pub fn signed<'b>(&'b self) -> CookieJar<'b> {
return CookieJar {
flavor: Flavor::Child(Child {
parent: self,
read: secure::design,
write: secure::sign,
})
};
}
pub fn encrypted<'b>(&'b self) -> CookieJar<'b> {
return CookieJar {
flavor: Flavor::Child(Child {
parent: self,
read: secure::design_and_decrypt,
write: secure::encrypt_and_sign,
})
};
}
pub fn permanent<'b>(&'b self) -> CookieJar<'b> {
return CookieJar {
flavor: Flavor::Child(Child {
parent: self,
read: read,
write: write,
})
};
fn read(_root: &Root, cookie: Cookie) -> Option<Cookie> {
Some(cookie)
}
fn write(_root: &Root, mut cookie: Cookie) -> Cookie {
cookie.max_age = Some(3600 * 24 * 365 * 20);
let mut now = time::now();
now.tm_year += 20;
cookie.expires = Some(now);
cookie
}
}
pub fn delta(&self) -> Vec<Cookie> {
let mut ret = Vec::new();
let root = self.root();
for cookie in root.removed_cookies.borrow().iter() {
let mut c = Cookie::new(cookie.clone(), String::new());
c.max_age = Some(0);
let mut now = time::now();
now.tm_year -= 1;
c.expires = Some(now);
ret.push(c);
}
let map = root.map.borrow();
for cookie in root.new_cookies.borrow().iter() {
ret.push(map.get(cookie).unwrap().clone());
}
return ret;
}
fn try_read(&self, root: &Root, mut cookie: Cookie) -> Option<Cookie> {
let mut jar = self;
loop {
match jar.flavor {
Flavor::Child(Child { read, parent, .. }) => {
cookie = match read(root, cookie) {
Some(c) => c, None => return None,
};
jar = parent;
}
Flavor::Root(..) => return Some(cookie),
}
}
}
pub fn iter(&self) -> Iter {
let map = self.root().map.borrow();
Iter { jar: self, keys: map.keys().cloned().collect() }
}
}
impl<'a> Iterator for Iter<'a> {
type Item = Cookie;
fn next(&mut self) -> Option<Cookie> {
loop {
let key = match self.keys.pop() {
Some(v) => v,
None => return None,
};
let root = self.jar.root();
let map = root.map.borrow();
let cookie = match map.get(&key) {
Some(cookie) => cookie.clone(),
None => continue,
};
match self.jar.try_read(root, cookie) {
Some(cookie) => return Some(cookie),
None => {}
}
}
}
}
mod secure {
use Cookie;
use jar::{Root};
use openssl::crypto::{hmac, hash, memcmp, symm};
use serialize::hex::{ToHex, FromHex};
use std::io::prelude::*;
pub const MIN_KEY_LEN: usize = 32;
pub fn sign(root: &Root, mut cookie: Cookie) -> Cookie {
let signature = dosign(root, cookie.value.as_slice());
cookie.value.push_str("--");
cookie.value.push_str(signature.as_slice().to_hex().as_slice());
cookie
}
fn split_value(val: &str) -> Option<(&str, Vec<u8>)> {
let parts = val.split("--");
let ext = match parts.last() {
Some(ext) => ext,
_ => return None,
};
let val_len = val.len();
if ext.len() == val_len { return None }
let text = &val[..val_len - ext.len() - 2];
let ext = match ext.from_hex() {
Ok(sig) => sig, Err(..) => return None,
};
Some((text, ext))
}
pub fn design(root: &Root, mut cookie: Cookie) -> Option<Cookie> {
let len = {
let (text, signature) = match split_value(cookie.value.as_slice()) {
Some(pair) => pair, None => return None
};
let expected = dosign(root, text);
if !memcmp::eq(expected.as_slice(), signature.as_slice()) {
return None
}
text.len()
};
cookie.value.truncate(len);
Some(cookie)
}
fn dosign(root: &Root, val: &str) -> Vec<u8> {
let mut hmac = hmac::HMAC::new(hash::Type::SHA1, root.key.as_slice());
hmac.write_all(val.as_bytes()).unwrap();
hmac.finish()
}
pub fn encrypt_and_sign(root: &Root, mut cookie: Cookie) -> Cookie {
let encrypted_data = encrypt_data(root, cookie.value.as_slice());
cookie.value = encrypted_data;
sign(root, cookie)
}
fn encrypt_data(root: &Root, val: &str) -> String {
let iv = random_iv();
let iv_str = iv.as_slice().to_hex();
let mut encrypted_data = symm::encrypt(
symm::Type::AES_256_CBC,
&root.key[..MIN_KEY_LEN],
iv,
val.as_bytes()
).as_slice().to_hex();
encrypted_data.push_str("--");
encrypted_data.push_str(iv_str.as_slice());
encrypted_data
}
pub fn design_and_decrypt(root: &Root, cookie: Cookie) -> Option<Cookie> {
let mut cookie = match design(root, cookie) {
Some(cookie) => cookie,
None => return None
};
let decrypted_data = decrypt_data(root, cookie.value.as_slice()).and_then(|data| String::from_utf8(data).ok());
match decrypted_data {
Some(val) => { cookie.value = val; Some(cookie) },
None => return None
}
}
fn decrypt_data(root: &Root, val: &str) -> Option<Vec<u8>> {
let (val, iv) = match split_value(val) {
Some(pair) => pair, None => return None
};
let actual = match val.as_slice().from_hex() {
Ok(actual) => actual, Err(_) => return None
};
Some(symm::decrypt(
symm::Type::AES_256_CBC,
&root.key[..MIN_KEY_LEN],
iv,
actual.as_slice()
))
}
fn random_iv() -> Vec<u8> {
::openssl::crypto::rand::rand_bytes(16)
}
pub fn prepare_key(key: &[u8]) -> Vec<u8> {
hash::hash(hash::Type::SHA256, key)
}
}
#[cfg(test)]
mod test {
use {Cookie, CookieJar};
const KEY: &'static [u8] = b"f8f9eaf1ecdedff5e5b749c58115441e";
#[test]
fn short_key() {
CookieJar::new(b"foo");
}
#[test]
fn simple() {
let c = CookieJar::new(KEY);
c.add(Cookie::new("test".to_string(), "".to_string()));
c.add(Cookie::new("test2".to_string(), "".to_string()));
c.remove("test");
assert!(c.find("test").is_none());
assert!(c.find("test2").is_some());
}
macro_rules! secure_behaviour {
($c:ident, $secure:ident) => ({
$c.$secure().add(Cookie::new("test".to_string(), "test".to_string()));
assert!($c.find("test").unwrap().value.as_slice() != "test");
assert!($c.$secure().find("test").unwrap().value.as_slice() == "test");
let mut cookie = $c.find("test").unwrap();
cookie.value.push('l');
$c.add(cookie);
assert!($c.$secure().find("test").is_none());
let mut cookie = $c.find("test").unwrap();
cookie.value = "foobar".to_string();
$c.add(cookie);
assert!($c.$secure().find("test").is_none());
})
}
#[test]
fn signed() {
let c = CookieJar::new(KEY);
secure_behaviour!(c, signed)
}
#[test]
fn encrypted() {
let c = CookieJar::new(KEY);
secure_behaviour!(c, encrypted)
}
#[test]
fn permanent() {
let c = CookieJar::new(KEY);
c.permanent().add(Cookie::new("test".to_string(), "test".to_string()));
let cookie = c.find("test").unwrap();
assert_eq!(cookie.value.as_slice(), "test");
assert_eq!(c.permanent().find("test").unwrap().value.as_slice(), "test");
assert!(cookie.expires.is_some());
assert!(cookie.max_age.is_some());
}
#[test]
fn chained() {
let c = CookieJar::new(KEY);
c.permanent().signed()
.add(Cookie::new("test".to_string(), "test".to_string()));
let cookie = c.signed().find("test").unwrap();
assert_eq!(cookie.value.as_slice(), "test");
assert!(cookie.expires.is_some());
assert!(cookie.max_age.is_some());
}
#[test]
fn iter() {
let mut c = CookieJar::new(KEY);
c.add_original(Cookie::new("original".to_string(),
"original".to_string()));
c.add(Cookie::new("test".to_string(), "test".to_string()));
c.add(Cookie::new("test2".to_string(), "test2".to_string()));
c.add(Cookie::new("test3".to_string(), "test3".to_string()));
c.add(Cookie::new("test4".to_string(), "test4".to_string()));
c.signed()
.add(Cookie::new("signed".to_string(), "signed".to_string()));
c.encrypted()
.add(Cookie::new("encrypted".to_string(), "encrypted".to_string()));
c.remove("test");
let cookies = c.iter().collect::<Vec<_>>();
assert_eq!(cookies.len(), 6);
let encrypted_cookies = c.encrypted().iter().collect::<Vec<_>>();
assert_eq!(encrypted_cookies.len(), 1);
assert_eq!(encrypted_cookies[0].name.as_slice(), "encrypted");
let signed_cookies = c.signed().iter().collect::<Vec<_>>();
assert_eq!(signed_cookies.len(), 2);
assert!(signed_cookies[0].name.as_slice() == "signed" ||
signed_cookies[1].name.as_slice() == "signed");
}
}