use std::marker::PhantomData;
use constant_time_eq::constant_time_eq;
use digest::{Digest, FixedOutputReset};
use rand::RngCore;
use crate::controller_builder::ControllerBuilder;
use crate::prefixed_api_key::PrefixedApiKey;
#[derive(Clone, Debug)]
pub struct PrefixedApiKeyController<R: RngCore, D: Digest + FixedOutputReset> {
prefix: String,
rng: R,
digest: PhantomData<D>,
short_token_prefix: Option<String>,
short_token_length: usize,
long_token_length: usize,
}
impl<R: RngCore, D: Digest + FixedOutputReset> PrefixedApiKeyController<R, D> {
pub fn new(
prefix: String,
rng: R,
short_token_prefix: Option<String>,
short_token_length: usize,
long_token_length: usize,
) -> PrefixedApiKeyController<R, D> {
PrefixedApiKeyController {
prefix,
rng,
digest: PhantomData,
short_token_prefix,
short_token_length,
long_token_length,
}
}
pub fn configure() -> ControllerBuilder<R, D> {
ControllerBuilder::new()
}
fn get_random_bytes(&mut self, length: usize) -> Vec<u8> {
let mut random_bytes = vec![0u8; length];
self.rng.fill_bytes(&mut random_bytes);
random_bytes
}
fn get_random_token(&mut self, length: usize) -> String {
let bytes = self.get_random_bytes(length);
bs58::encode(bytes).into_string()
}
pub fn generate_key(&mut self) -> PrefixedApiKey {
let mut short_token = self.get_random_token(self.short_token_length);
if self.short_token_prefix.is_some() {
let prefix_string = self.short_token_prefix.as_ref().unwrap().to_owned();
short_token = (prefix_string + &short_token)
.chars()
.take(self.short_token_length)
.collect()
}
let long_token = self.get_random_token(self.long_token_length);
PrefixedApiKey::new(self.prefix.to_owned(), short_token, long_token)
}
pub fn generate_key_and_hash(&mut self) -> (PrefixedApiKey, String) {
let pak = self.generate_key();
let hash = self.long_token_hashed(&pak);
(pak, hash)
}
pub fn long_token_hashed(&self, pak: &PrefixedApiKey) -> String {
let mut digest = D::new();
pak.long_token_hashed(&mut digest)
}
pub fn check_hash(&self, pak: &PrefixedApiKey, hash: &str) -> bool {
let pak_hash = self.long_token_hashed(pak);
constant_time_eq(pak_hash.as_bytes(), hash.as_bytes())
}
}
#[cfg(test)]
mod controller_tests {
use rand::rngs::OsRng;
use sha2::Sha256;
use crate::controller::PrefixedApiKeyController;
use crate::PrefixedApiKey;
#[test]
fn configuration_works() {
let controller = PrefixedApiKeyController::<_, Sha256>::configure()
.default_lengths()
.prefix("mycompany".to_owned())
.rng(OsRng)
.finalize();
assert!(controller.is_ok())
}
#[test]
fn generator() {
let mut generator =
PrefixedApiKeyController::<_, Sha256>::new("mycompany".to_owned(), OsRng, None, 8, 24);
let token_string = generator.generate_key().to_string();
let pak_result = PrefixedApiKey::from_string(&token_string);
assert_eq!(pak_result.is_ok(), true);
let pak_string = pak_result.unwrap().to_string();
assert_eq!(token_string, pak_string);
}
#[test]
fn generator_short_token_prefix() {
let short_length = 8;
let short_prefix = "a".repeat(short_length);
let mut generator = PrefixedApiKeyController::<_, Sha256>::new(
"mycompany".to_owned(),
OsRng,
Some(short_prefix.clone()),
short_length,
24,
);
let pak_short_token = generator.generate_key().short_token().to_owned();
assert_eq!(pak_short_token, short_prefix);
}
#[test]
fn generate_key_and_hash() {
let mut generator =
PrefixedApiKeyController::<_, Sha256>::new("mycompany".to_owned(), OsRng, None, 8, 24);
let (pak, hash) = generator.generate_key_and_hash();
assert!(generator.check_hash(&pak, &hash))
}
#[test]
fn check_long_token_via_generator() {
let pak_string = "mycompany_CEUsS4psCmc_BddpcwWyCT3EkDjHSSTRaSK1dxtuQgbjb";
let hash = "0f01ab6e0833f280b73b2b618c16102d91c0b7c585d42a080d6e6603239a8bee";
let pak: PrefixedApiKey = pak_string.try_into().unwrap();
let generator =
PrefixedApiKeyController::<_, Sha256>::new("mycompany".to_owned(), OsRng, None, 8, 24);
assert_eq!(generator.long_token_hashed(&pak), hash);
}
#[test]
fn generator_digest_resets_after_hashing() {
let pak1_string = "mycompany_CEUsS4psCmc_BddpcwWyCT3EkDjHSSTRaSK1dxtuQgbjb";
let pak1_hash = "0f01ab6e0833f280b73b2b618c16102d91c0b7c585d42a080d6e6603239a8bee";
let pak1: PrefixedApiKey = pak1_string.try_into().unwrap();
let pak2_string = "mycompany_CEUsS4psCmc_BddpcwWyCT3EkDjHSSTRaSK1dxtuQgbjb";
let pak2_hash = "0f01ab6e0833f280b73b2b618c16102d91c0b7c585d42a080d6e6603239a8bee";
let pak2: PrefixedApiKey = pak2_string.try_into().unwrap();
let generator =
PrefixedApiKeyController::<_, Sha256>::new("mycompany".to_owned(), OsRng, None, 8, 24);
assert_eq!(generator.long_token_hashed(&pak1), pak1_hash);
assert_eq!(generator.long_token_hashed(&pak2), pak2_hash);
}
#[test]
fn generator_matches_hash() {
let pak_string = "mycompany_CEUsS4psCmc_BddpcwWyCT3EkDjHSSTRaSK1dxtuQgbjb";
let pak_hash = "0f01ab6e0833f280b73b2b618c16102d91c0b7c585d42a080d6e6603239a8bee";
let pak: PrefixedApiKey = pak_string.try_into().unwrap();
let generator =
PrefixedApiKeyController::<_, Sha256>::new("mycompany".to_owned(), OsRng, None, 8, 24);
assert!(generator.check_hash(&pak, pak_hash));
}
}