flyer 3.0.7

HTTP framework for rust
Documentation
use std::time::Duration;

use aes_gcm::{
    AeadCore,
    Aes256Gcm,
    Key,
    Nonce,
    aead::{Aead, KeyInit, OsRng}
};
use anyhow::Result;
use base64::{Engine, engine::general_purpose};

use crate::{
    hooks::Hook,
    request::Request,
    response::Response,
    routing::next::Next,
    session::Session
};

pub struct CookieSession {
    cookie_name: String,
    encryption_key: String,
    duration: Duration,
}

impl CookieSession {
    pub fn new(cookie_name: impl Into<String>, encryption_key: impl Into<String>, expires: Duration) -> Self {
        return Self {
            cookie_name: cookie_name.into(),
            encryption_key:  Self::string_fixed_length(&encryption_key.into(), 32),
            duration: expires,
        };
    }
}

impl Hook for CookieSession {
    async fn before(&self, mut req: Request, res: Response, next: Next) -> Response {
        let hash = req.cookie(self.cookie_name.clone());

        if hash.is_empty() {
            return next.handle(req, res);
        }

        let result = self.decrypt(&hash);

        if result.is_err() {
            return next.handle(req, res);
        }

        let result = serde_json::from_str::<Session>(&result.unwrap());

        if result.is_err() {
            return next.handle(req, res);
        }

        req.session = result.unwrap();

        return next.handle(req, res);
    }
    
    async fn after(&self, req: Request, mut res: Response, next: Next) -> Response {
        let data = serde_json::to_string(&res.session)
            .unwrap();
        let payload = self
            .encrypt(&data)
            .unwrap();
        let cookie = res
            .set_cookie(self.cookie_name.clone(), payload)
            .set_expires(self.duration)
            .set_same_site(crate::cookies::SameSite::Lax)
            .set_path("/")
            .parse();

        return next.handle(req, res.set_header("Set-Cookie", cookie));
    }
}

impl CookieSession {
    pub fn string_fixed_length(input: &str, length: usize) -> String {
        return match input.len() < length {
            true => format!("{:<width$}", input, width = length),
            false => input.chars().take(length).collect(),
        };
    }

    pub fn encrypt(&self, data: &str) -> Result<String> {
        let key = Key::<Aes256Gcm>::from_slice(self.encryption_key.as_bytes());
        let cipher = Aes256Gcm::new(key);
        let nonce = Aes256Gcm::generate_nonce(OsRng);
        let cipher_text = cipher.encrypt(&nonce, data.as_bytes()).unwrap();
        let mut combined = nonce.to_vec();

        combined.extend_from_slice(&cipher_text);

        let mut buffer = String::new();

        general_purpose::STANDARD.encode_string(combined, &mut buffer);

        return Ok(buffer);
    }

    pub fn decrypt(&self, hash: &str) -> Result<String> {
        let key = Key::<Aes256Gcm>::from_slice(self.encryption_key.as_bytes());
        let cipher = Aes256Gcm::new(key);
        let combined = general_purpose::STANDARD.decode(hash);

        if combined.is_err() {
            return Ok(String::new());
        }

        let session = combined.unwrap();
        let (nonce_bytes, ciphertext) = session.split_at(12);
        let nonce = Nonce::from_slice(nonce_bytes);
        let plaintext = cipher.decrypt(nonce, ciphertext).unwrap();

        return String::from_utf8(plaintext).map_err(|err| err.into());
    }
}