omegaupload/common/src/lib.rs

#![warn(clippy::nursery, clippy::pedantic)]
#![deny(unsafe_code)]

//! Contains common functions and structures used by multiple projects

use std::str::FromStr;

use sha2::{Digest, Sha256};
use thiserror::Error;
pub use url::Url;

use crate::crypto::{Key, Nonce};

pub mod base64 {
    /// URL-safe Base64 encoding.
    pub fn encode(input: impl AsRef<[u8]>) -> String {
        base64::encode_config(input, base64::URL_SAFE)
    }

    /// URL-safe Base64 decoding.
    pub fn decode(input: impl AsRef<[u8]>) -> Result<Vec<u8>, base64::DecodeError> {
        base64::decode_config(input, base64::URL_SAFE)
    }
}

/// Hashes an input to output a usable key.
pub fn hash(data: impl AsRef<[u8]>) -> crypto::Key {
    let mut hasher = Sha256::new();
    hasher.update(data);
    hasher.finalize()
}

pub mod crypto {
    use std::ops::{Deref, DerefMut};

    use chacha20poly1305::aead::generic_array::GenericArray;
    use chacha20poly1305::aead::{Aead, Error, NewAead};
    use chacha20poly1305::XChaCha20Poly1305;
    use chacha20poly1305::XNonce;
    use rand::{thread_rng, Rng};

    pub use chacha20poly1305::Key;

    /// Securely generates a random key and nonce.
    #[must_use]
    pub fn gen_key_nonce() -> (Key, Nonce) {
        let mut rng = thread_rng();
        let mut key: Key = GenericArray::default();
        rng.fill(key.as_mut_slice());
        let mut nonce = Nonce::default();
        rng.fill(nonce.as_mut_slice());
        (key, nonce)
    }

    pub fn seal(plaintext: &[u8], nonce: &Nonce, key: &Key) -> Result<Vec<u8>, Error> {
        let cipher = XChaCha20Poly1305::new(key);
        cipher.encrypt(nonce, plaintext)
    }

    pub fn open(encrypted: &[u8], nonce: &Nonce, key: &Key) -> Result<Vec<u8>, Error> {
        let cipher = XChaCha20Poly1305::new(key);
        cipher.decrypt(nonce, encrypted)
    }

    #[derive(Clone, Copy, PartialEq, Eq, Debug)]
    pub struct Nonce(XNonce);

    impl Default for Nonce {
        fn default() -> Self {
            Self(GenericArray::default())
        }
    }

    impl Deref for Nonce {
        type Target = XNonce;

        fn deref(&self) -> &Self::Target {
            &self.0
        }
    }

    impl DerefMut for Nonce {
        fn deref_mut(&mut self) -> &mut Self::Target {
            &mut self.0
        }
    }

    impl AsRef<[u8]> for Nonce {
        fn as_ref(&self) -> &[u8] {
            self.0.as_ref()
        }
    }

    impl Nonce {
        #[must_use]
        pub fn increment(&self) -> Self {
            todo!()
        }

        #[must_use]
        pub fn from_slice(slice: &[u8]) -> Self {
            Self(*XNonce::from_slice(slice))
        }
    }
}

pub struct ParsedUrl {
    pub sanitized_url: Url,
    pub decryption_key: Key,
    pub nonce: Nonce,
    pub needs_password: bool,
}

#[derive(Default)]
pub struct PartialParsedUrl {
    pub decryption_key: Option<Key>,
    pub nonce: Option<Nonce>,
    pub needs_password: bool,
}

impl From<&str> for PartialParsedUrl {
    fn from(fragment: &str) -> Self {
        let args = fragment.split('!').filter_map(|kv| {
            let (k, v) = {
                let mut iter = kv.split(':');
                (iter.next(), iter.next())
            };

            Some((k?, v))
        });

        let mut decryption_key = None;
        let mut needs_password = false;
        let mut nonce = None;

        for (key, value) in args {
            match (key, value) {
                ("key", Some(value)) => {
                    decryption_key = base64::decode(value)
                        .map(|k| Key::from_slice(&k).clone())
                        .ok();
                }
                ("pw", _) => {
                    needs_password = true;
                }
                ("nonce", Some(value)) => {
                    nonce = base64::decode(value).as_deref().map(Nonce::from_slice).ok();
                }
                _ => (),
            }
        }

        Self {
            decryption_key,
            nonce,
            needs_password,
        }
    }
}

#[derive(Debug, Error)]
pub enum ParseUrlError {
    #[error("The provided url was bad")]
    BadUrl,
    #[error("Missing decryption key")]
    NeedKey,
    #[error("Missing nonce")]
    NeedNonce,
    #[error("Missing decryption key and nonce")]
    NeedKeyAndNonce,
}

impl FromStr for ParsedUrl {
    type Err = ParseUrlError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut url = Url::from_str(s).map_err(|_| ParseUrlError::BadUrl)?;
        let fragment = url.fragment().ok_or(ParseUrlError::NeedKeyAndNonce)?;
        if fragment.is_empty() {
            return Err(ParseUrlError::NeedKeyAndNonce);
        }

        let PartialParsedUrl {
            decryption_key,
            needs_password,
            nonce,
        } = PartialParsedUrl::from(fragment);

        url.set_fragment(None);

        let (decryption_key, nonce) = match (&decryption_key, nonce) {
            (None, None) => Err(ParseUrlError::NeedKeyAndNonce),
            (None, Some(_)) => Err(ParseUrlError::NeedKey),
            (Some(_), None) => Err(ParseUrlError::NeedNonce),
            (Some(k), Some(v)) => Ok((*k, v)),
        }?;

        Ok(Self {
            sanitized_url: url,
            decryption_key,
            needs_password,
            nonce,
        })
    }
}