dkl/src/dls/store.rs

use openssl::symm::Mode;
use std::borrow::Cow;
use std::path::{Path, PathBuf};
use tokio::{fs, io::AsyncWriteExt};

pub type Salt = [u8; 16];
pub type Key = [u8; 32];
pub type Hash = Vec<u8>;

pub fn hash_password(salt: &[u8], passphrase: &str) -> Result<Key> {
    let hash = argon2::hash_raw(
        passphrase.as_bytes(),
        salt,
        &argon2::Config {
            variant: argon2::Variant::Argon2id,
            hash_length: 32,
            time_cost: 1,
            mem_cost: 64 << 10,
            thread_mode: argon2::ThreadMode::Parallel,
            lanes: 4,
            ..Default::default()
        },
    )
    .map_err(Error::Hash)?;

    Ok(hash.try_into().unwrap())
}

pub struct Store {
    path: PathBuf,
    key: Option<Key>,
}

impl Store {
    pub fn new(path: impl Into<PathBuf>) -> Self {
        Self {
            path: path.into(),
            key: None,
        }
    }

    pub async fn unlock(&mut self, passphrase: &str) -> Result<()> {
        let keys = self.read_keys().await?;

        let salt = keys.salt;

        let user_key = hash_password(&salt, passphrase)?;
        let user_hash = openssl::sha::sha512(&user_key);

        for nk in keys.keys {
            if nk.hash != user_hash {
                continue;
            }

            // build iv+data from salt
            let mut enc_data = Vec::with_capacity(salt.len() + nk.enc_key.len());
            enc_data.extend_from_slice(&salt);
            enc_data.extend_from_slice(&nk.enc_key);

            let key = crypt(&enc_data, &user_key, Mode::Decrypt)?;
            let key = key.try_into().map_err(|_| Error::InvalidKey)?;

            self.key = Some(key);
            return Ok(());
        }

        Err(Error::KeyNotFound)
    }

    pub async fn read(&self, path: impl AsRef<Path>) -> Result<Vec<u8>> {
        let enc_data = self.read_file(path.as_ref().with_extension("data")).await?;
        self.decrypt(&enc_data)
    }
    pub async fn read_to_string(&self, path: impl AsRef<Path>) -> Result<String> {
        let path = path.as_ref();
        String::from_utf8(self.read(path).await?).map_err(|_| Error::InvalidUtf8(path.into()))
    }

    pub async fn write(&self, path: impl AsRef<Path>, data: &[u8]) -> Result<()> {
        let enc_data = self.encrypt(data)?;
        safe_write(self.path.join(&path).with_extension("data"), &enc_data).await
    }

    async fn read_keys(&self) -> Result<Keys<'_>> {
        let keys = self.read_file(".keys").await?;
        let Some(keys) = keys.strip_prefix(b"{json}") else {
            return Err(Error::InvalidKeys);
        };

        serde_json::from_slice(keys).map_err(Error::KeysParse)
    }

    async fn read_file(&self, subpath: impl AsRef<Path>) -> Result<Vec<u8>> {
        let path = self.path.join(subpath);
        fs::read(&path).await.map_err(|e| Error::Read(path, e))
    }

    fn encrypt(&self, src: &[u8]) -> Result<Vec<u8>> {
        self.crypt(src, Mode::Encrypt)
    }
    fn decrypt(&self, src: &[u8]) -> Result<Vec<u8>> {
        self.crypt(src, Mode::Decrypt)
    }
    fn crypt(&self, src: &[u8], mode: Mode) -> Result<Vec<u8>> {
        let key = self.key.as_ref().ok_or(Error::NotUnlocked)?;
        crypt(src, key, mode)
    }
}

async fn safe_write(path: impl AsRef<Path>, contents: &[u8]) -> Result<()> {
    let path = path.as_ref();
    let tmp = path.with_added_extension("new");

    let tmp_err = |e| Error::Write(tmp.clone(), e);

    let mut file = fs::OpenOptions::new()
        .create(true)
        .truncate(true)
        .write(true)
        .open(&tmp)
        .await
        .map_err(tmp_err)?;

    file.write_all(contents).await.map_err(tmp_err)?;
    file.sync_all().await.map_err(tmp_err)?;
    file.shutdown().await.map_err(tmp_err)?;

    fs::rename(tmp, &path)
        .await
        .map_err(|e| Error::Write(path.into(), e))
}

fn crypt(mut src: &[u8], key: &Key, mode: Mode) -> Result<Vec<u8>> {
    use openssl::symm::{Cipher, Crypter};

    let mut iv = [0u8; 16];
    let mut dst: Vec<u8>;
    let crypt_dst: &mut [u8];

    match mode {
        Mode::Encrypt => {
            getrandom::fill(&mut iv).unwrap();
            dst = vec![0u8; iv.len() + src.len()];
            dst[..iv.len()].copy_from_slice(&iv);
            crypt_dst = &mut dst[iv.len()..];
        }
        Mode::Decrypt => {
            iv = src[..iv.len()]
                .try_into()
                .map_err(|_| Error::DecryptInputToSmall)?;
            src = &src[iv.len()..];
            dst = vec![0u8; src.len()];
            crypt_dst = &mut dst;
        }
    }

    let cipher = Cipher::aes_256_cfb128();
    Crypter::new(cipher, mode, key, Some(&iv))
        .expect("Failed to init AES")
        .update(src, crypt_dst)
        .expect("AES CFB encrypt/decrypt failed");

    Ok(dst)
}

pub type Result<T> = std::result::Result<T, Error>;

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("hash error: {0}")]
    Hash(argon2::Error),
    #[error("read {0} failed: {1}")]
    Read(PathBuf, std::io::Error),
    #[error("write {0} failed: {1}")]
    Write(PathBuf, std::io::Error),
    #[error("read {0} failed: invalid UTF-8")]
    InvalidUtf8(PathBuf),
    #[error("invalid keys data")]
    InvalidKeys,
    #[error("invalid key")]
    InvalidKey,
    #[error("keys parse error: {0}")]
    KeysParse(serde_json::Error),
    #[error("key not found")]
    KeyNotFound,
    #[error("store not unlocked")]
    NotUnlocked,
    #[error("decrypt input too small")]
    DecryptInputToSmall,
}

#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "PascalCase")]
struct Keys<'t> {
    salt: Salt,
    keys: Vec<NamedKey<'t>>,
}

#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "PascalCase")]
struct NamedKey<'t> {
    name: Cow<'t, str>,
    hash: Hash,
    enc_key: Key,
}