188 lines
5.3 KiB
Rust
188 lines
5.3 KiB
Rust
pub mod error;
|
|
|
|
use std::{
|
|
borrow::Cow,
|
|
fmt::{Display, Formatter},
|
|
};
|
|
use rand::RngCore;
|
|
use crate::error::Error;
|
|
|
|
#[derive(Debug)]
|
|
pub struct ApiKey {
|
|
pub prefix: String,
|
|
pub short_token: String,
|
|
pub long_token: String,
|
|
|
|
pub short_bytes: Vec<u8>,
|
|
pub long_bytes: Vec<u8>,
|
|
}
|
|
|
|
impl ApiKey {
|
|
pub fn with_prefix(prefix: &str, options: impl Into<Option<Options>>) -> Self {
|
|
generate(prefix, options)
|
|
}
|
|
}
|
|
|
|
impl Display for ApiKey {
|
|
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
|
write!(
|
|
f,
|
|
"{}_{}_{}",
|
|
self.prefix,
|
|
self.short_token,
|
|
self.long_token,
|
|
)
|
|
}
|
|
}
|
|
|
|
pub fn parse(key: impl AsRef<str>) -> Result<ApiKey, Error> {
|
|
let key = key.as_ref();
|
|
let parts: Vec<_> = key.rsplitn(3, "_").collect();
|
|
if parts.len() != 3 {
|
|
return Err(Error::InvalidKeyFormat);
|
|
}
|
|
|
|
let prefix = parts[2];
|
|
let short_token = parts[1];
|
|
let long_token = parts[0];
|
|
|
|
if prefix.is_empty() || short_token.is_empty() || long_token.is_empty() {
|
|
return Err(Error::InvalidKeyFormat);
|
|
}
|
|
|
|
let short_bytes = bs58::decode(short_token)
|
|
.into_vec()
|
|
.map_err(|_| Error::InvalidBase58)?;
|
|
let long_bytes = bs58::decode(long_token)
|
|
.into_vec()
|
|
.map_err(|_| Error::InvalidBase58)?;
|
|
|
|
Ok(ApiKey {
|
|
prefix: prefix.to_string(),
|
|
short_token: short_token.to_string(),
|
|
long_token: long_token.to_string(),
|
|
|
|
short_bytes,
|
|
long_bytes,
|
|
})
|
|
}
|
|
|
|
pub struct Options {
|
|
short_token_length: usize,
|
|
long_token_length: usize,
|
|
}
|
|
|
|
impl Default for Options {
|
|
fn default() -> Self {
|
|
Self {
|
|
short_token_length: 8,
|
|
long_token_length: 24,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn generate(prefix: impl AsRef<str>, options: impl Into<Option<Options>>) -> ApiKey {
|
|
let options = options.into().unwrap_or_default();
|
|
let mut rng = rand::thread_rng();
|
|
|
|
let mut short_bytes = vec![0; options.short_token_length];
|
|
rng.fill_bytes(&mut short_bytes);
|
|
|
|
let mut long_bytes = vec![0; options.long_token_length];
|
|
rng.fill_bytes(&mut long_bytes);
|
|
|
|
let short_raw = bs58::encode(&short_bytes).into_string();
|
|
let short_token = &pad_left(
|
|
&short_raw,
|
|
'0',
|
|
options.short_token_length,
|
|
)[..options.short_token_length];
|
|
|
|
let long_raw = bs58::encode(&long_bytes).into_string();
|
|
let long_token = &pad_left(
|
|
&long_raw,
|
|
'0',
|
|
options.long_token_length,
|
|
)[..options.long_token_length];
|
|
|
|
ApiKey {
|
|
prefix: prefix.as_ref().to_string(),
|
|
short_token: short_token.to_string(),
|
|
long_token: long_token.to_string(),
|
|
|
|
short_bytes: bs58::decode(short_token).into_vec().unwrap(),
|
|
long_bytes: bs58::decode(long_token).into_vec().unwrap(),
|
|
}
|
|
}
|
|
|
|
fn pad_left(s: &str, c: char, n: usize) -> Cow<str> {
|
|
if s.len() >= n {
|
|
return Cow::Borrowed(s);
|
|
}
|
|
|
|
let needed = n - s.len();
|
|
let pad: String = std::iter::repeat(c)
|
|
.take(needed)
|
|
.collect();
|
|
|
|
Cow::Owned(format!("{}{}", pad, s))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod test {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_pad() {
|
|
assert_eq!("00000001", pad_left("1", '0', 8));
|
|
assert_eq!("000001", pad_left("1", '0', 6));
|
|
assert_eq!("12345678", pad_left("12345678", '0', 8));
|
|
assert_eq!("123456789", pad_left("123456789", '0', 8));
|
|
}
|
|
|
|
#[test]
|
|
fn test_bad_parse() {
|
|
assert_eq!(parse("").unwrap_err(), Error::InvalidKeyFormat);
|
|
assert_eq!(parse("__").unwrap_err(), Error::InvalidKeyFormat);
|
|
assert_eq!(parse("___").unwrap_err(), Error::InvalidKeyFormat);
|
|
assert_eq!(parse("uwu_!@#_!@#!@#").unwrap_err(), Error::InvalidBase58);
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse() {
|
|
let key = "mycompany_BRTRKFsL_51FwqftsmMDHHbJAMEXXHCgG";
|
|
let parsed = parse(key).unwrap();
|
|
assert_eq!(parsed.prefix, "mycompany");
|
|
assert_eq!(parsed.short_token, "BRTRKFsL");
|
|
assert_eq!(parsed.long_token, "51FwqftsmMDHHbJAMEXXHCgG");
|
|
assert_eq!(parsed.short_bytes, bs58::decode("BRTRKFsL").into_vec().unwrap());
|
|
assert_eq!(parsed.long_bytes, bs58::decode("51FwqftsmMDHHbJAMEXXHCgG").into_vec().unwrap());
|
|
assert_eq!(parsed.to_string(), key);
|
|
}
|
|
|
|
#[test]
|
|
fn test_multiple_underscore_prefix() {
|
|
let key = "a_b_c_BRTRKFsL_51FwqftsmMDHHbJAMEXXHCgG";
|
|
let parsed = parse(key).unwrap();
|
|
assert_eq!(parsed.prefix, "a_b_c");
|
|
assert_eq!(parsed.short_token, "BRTRKFsL");
|
|
assert_eq!(parsed.long_token, "51FwqftsmMDHHbJAMEXXHCgG");
|
|
assert_eq!(parsed.short_bytes, bs58::decode("BRTRKFsL").into_vec().unwrap());
|
|
assert_eq!(parsed.long_bytes, bs58::decode("51FwqftsmMDHHbJAMEXXHCgG").into_vec().unwrap());
|
|
assert_eq!(parsed.to_string(), key);
|
|
}
|
|
|
|
#[test]
|
|
fn round_trip() {
|
|
const PREFIX: &str = "uwu";
|
|
let key = generate(PREFIX, None);
|
|
let parsed = parse(&key.to_string()).unwrap();
|
|
assert_eq!(parsed.prefix, PREFIX);
|
|
assert_eq!(parsed.short_token, key.short_token);
|
|
assert_eq!(parsed.long_token, key.long_token);
|
|
assert_eq!(parsed.short_bytes, key.short_bytes);
|
|
assert_eq!(parsed.long_bytes, key.long_bytes);
|
|
assert_eq!(parsed.to_string(), key.to_string());
|
|
}
|
|
}
|