79 lines
1.9 KiB
Rust
79 lines
1.9 KiB
Rust
use chrono::{DateTime, Duration, Utc};
|
|
use hmac::digest::KeyInit;
|
|
use hmac::Hmac;
|
|
use jwt::{SignWithKey, VerifyWithKey};
|
|
use serde::{de::DeserializeOwned, Deserialize, Serialize};
|
|
use sha2::Sha256;
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum TokenError {
|
|
#[error(transparent)]
|
|
Jwt(#[from] jwt::Error),
|
|
|
|
#[error("Token expired")]
|
|
Expired,
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
|
|
pub enum TokenType {
|
|
AccessToken,
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Deserialize)]
|
|
struct TokenFields<T> {
|
|
#[serde(rename = "d")]
|
|
data: T,
|
|
|
|
#[serde(rename = "exp")]
|
|
expire_at: DateTime<Utc>,
|
|
}
|
|
|
|
pub fn create_token(
|
|
server_key: &str,
|
|
data: impl Serialize,
|
|
expire_duration: Duration,
|
|
) -> Result<String, TokenError> {
|
|
Ok(
|
|
TokenFields {
|
|
data,
|
|
expire_at: Utc::now() + expire_duration,
|
|
}
|
|
.sign_with_key(&generate_hmac_key(server_key))?,
|
|
)
|
|
}
|
|
|
|
fn generate_hmac_key(server_key: &str) -> Hmac<Sha256> {
|
|
Hmac::<Sha256>::new_from_slice(server_key.as_bytes()).expect("invalid server key")
|
|
}
|
|
|
|
pub fn parse_token<T: DeserializeOwned>(server_key: &str, token: &str) -> Result<T, TokenError> {
|
|
let fields =
|
|
VerifyWithKey::<TokenFields<T>>::verify_with_key(token, &generate_hmac_key(server_key))?;
|
|
if fields.expire_at < Utc::now() {
|
|
return Err(TokenError::Expired);
|
|
}
|
|
Ok(fields.data)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn create_token_test() {
|
|
let token_data = "hello appflowy".to_string();
|
|
let token = create_token("server_key", &token_data, Duration::days(2)).unwrap();
|
|
|
|
let parse_token_data = parse_token::<String>("server_key", &token).unwrap();
|
|
assert_eq!(token_data, parse_token_data);
|
|
}
|
|
|
|
#[test]
|
|
#[should_panic]
|
|
fn parser_token_with_different_server_key() {
|
|
let server_key = "123456";
|
|
let token = create_token(server_key, "hello", Duration::days(2)).unwrap();
|
|
let _ = parse_token::<String>("abcdef", &token).unwrap();
|
|
}
|
|
}
|