seshat 2.3.3

// Copyright 2019 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use rand::{thread_rng, Rng};
use std::{
    fs::File,
    io::{BufWriter, Cursor, Error as IoError, ErrorKind, Read, Write},
    path::Path,
};

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

use aes::{
    cipher::{NewCipher, StreamCipher},
    Aes256Ctr,
};
use hkdf::Hkdf;
use hmac::{Hmac, Mac, NewMac};
use pbkdf2::pbkdf2;
use sha2::{Sha256, Sha512};

use tantivy::directory::{
    error::{
        DeleteError, IOError as TvIoError, LockError, OpenDirectoryError, OpenReadError,
        OpenWriteError,
    },
    AntiCallToken, Directory, DirectoryLock, Lock, ReadOnlySource, TerminatingWrite, WatchCallback,
    WatchHandle, WritePtr,
};

use zeroize::Zeroizing;

use crate::index::encrypted_stream::{AesReader, AesWriter};

/// KeyBuffer type that makes sure that the buffer is zeroed out before being
/// dropped.
type KeyBuffer = Zeroizing<Vec<u8>>;

/// Key derivation result type for our initial key derivation. Consists of a
/// tuple containing a encryption key, a MAC key, and a random salt.
type InitialKeyDerivationResult = (KeyBuffer, KeyBuffer, Vec<u8>);

/// Key derivation result for our subsequent key derivations. The salt will be
/// read from our key file and we will re-derive our encryption and MAC keys.
type KeyDerivationResult = (KeyBuffer, KeyBuffer);

// The constants here are chosen to be similar to the constants for the Matrix
// key export format[1].
// [1] https://matrix.org/docs/spec/client_server/r0.5.0#key-exports
const KEYFILE: &str = "seshat-index.key";
// 16 byte random salt.
const SALT_SIZE: usize = 16;
// 16 byte random IV for the AES-CTR mode.
const IV_SIZE: usize = 16;
// 32 byte or 256 bit encryption keys.
const KEY_SIZE: usize = 32;
// 32 byte message authentication code since HMAC-SHA256 is used.
const MAC_LENGTH: usize = 32;
// 1 byte for the store version.
const VERSION: u8 = 1;

#[cfg(test)]
// Tests don't need to protect against brute force attacks.
pub(crate) const PBKDF_COUNT: u32 = 10;

#[cfg(not(test))]
// This is quite a bit lower than the spec since the encrypted key and index
// will not be public. An attacker would need to get hold of the encrypted index,
// its key and only then move on to bruteforcing the key. Even if the attacker
// decrypts the index he wouldn't get to the events itself, only to the index
// of them.
pub(crate) const PBKDF_COUNT: u32 = 10_000;

#[derive(Clone, Debug)]
/// A Directory implementation that wraps a MmapDirectory and adds [AES][aes]
/// based encryption to the file read/write operations.
///
/// When a new EncryptedMmapDirectory is created a random 256 bit AES key is
/// generated, the store key. This key is encrypted using a key that will be
/// derived with the user provided passphrase using [PBKDF2][pbkdf]. We generate
/// a random 128 bit salt and derive a 512 bit key using PBKDF2:
///
/// ```text
///     derived_key = PBKDF2(SHA512, passphrase, salt, count, 512)
/// ```
///
/// After key derivation, the key is split into a 256 bit AES encryption key and
/// a 256 bit MAC key.
///
/// The store key will be encrypted using AES-CTR with the encryption key that
/// was derived before, a random IV will be generated before encrypting the
/// store_key:
///
/// ```text
///     ciphertext = AES256-CTR(iv, store_key)
/// ```
///
/// A MAC of the encrypted ciphertext will be created using
/// the derived MAC key and [HMAC-SHA256][hmac]:
///
/// ```text
///     mac = HMAC-SHA256(mac_key, version || iv || salt || ciphertext)
/// ```
///
/// The store key will be written to a file concatenated with a store version,
/// IV, salt, PBKDF count, and MAC. The PBKDF count will be stored using the big
/// endian byte order:
///
/// ```text
///     key_file = (version || iv || salt || pbkdf_count || mac || key_ciphertext)
/// ```
///
/// Our store key will be used to encrypt the many files that Tantivy generates.
/// For this, the store key will be expanded into an 256 bit encryption key and
/// a 256 bit MAC key using [HKDF][hkdf].
///
/// ```text
///     encryption_key, mac_key = HKDF(SHA512, store_key, "", 512)
/// ```
///
/// Those two keys are used to encrypt and authenticate the Tantivy files. The
/// encryption scheme is similar to the one used for the store key, AES-CTR mode
/// for encryption and HMAC-SHA256 for authentication. For every encrypted file
/// a new random 128 bit IV will be generated.
///
/// The MAC will be calculated only on the ciphertext:
///
/// ```text
///     mac = HMAC-SHA256(mac_key, ciphertext)
/// ```
///
/// The file format differs a bit, the MAC will be at the end of the file and
/// the ciphertext is between the IV and salt:
///
/// ```text
///     file_data = (iv || ciphertext || mac)
/// ```
///
/// [aes]: https://en.wikipedia.org/wiki/Advanced_Encryption_Standard
/// [pbkdf]: https://en.wikipedia.org/wiki/PBKDF2
/// [hkdf]: https://en.wikipedia.org/wiki/HKDF
/// [hmac]: https://en.wikipedia.org/wiki/HMAC
pub struct EncryptedMmapDirectory {
    mmap_dir: tantivy::directory::MmapDirectory,
    encryption_key: KeyBuffer,
    mac_key: KeyBuffer,
}

impl EncryptedMmapDirectory {
    fn new(store_key: KeyBuffer, path: &Path) -> Result<Self, OpenDirectoryError> {
        // Expand the store key into a encryption and MAC key.
        let (encryption_key, mac_key) = EncryptedMmapDirectory::expand_store_key(&store_key)?;

        // Open our underlying bare Tantivy mmap based directory.
        let mmap_dir = tantivy::directory::MmapDirectory::open(&path)?;

        Ok(EncryptedMmapDirectory {
            mmap_dir,
            encryption_key,
            mac_key,
        })
    }
    /// Open a encrypted mmap directory. If the directory is empty a new
    /// directory key will be generated and encrypted with the given passphrase.
    ///
    /// If a new store is created, this method will randomly generated a new
    /// store key and encrypted using the given passphrase.
    ///
    /// # Arguments
    ///
    /// * `path` - The path where the directory should reside in.
    /// * `passphrase` - The passphrase that was used to encrypt our directory
    /// or the one that will be used to encrypt our directory.
    /// * `key_derivation_count` - The number of iterations that our key
    /// derivation function should use. Can't be lower than 1, should be chosen
    /// as high as possible, depending on how much time is acceptable for the
    /// caller to wait. Is only used when a new store is created. The count will
    /// be stored with the store key.
    ///
    /// Returns an error if the path does not exist, if it is not a directory or
    /// if there was an error when trying to decrypt the directory key e.g. the
    /// given passphrase was incorrect.
    pub fn open_or_create<P: AsRef<Path>>(
        path: P,
        passphrase: &str,
        key_derivation_count: u32,
    ) -> Result<Self, OpenDirectoryError> {
        if passphrase.is_empty() {
            return Err(IoError::new(ErrorKind::Other, "empty passphrase").into());
        }

        if key_derivation_count == 0 {
            return Err(IoError::new(ErrorKind::Other, "invalid key derivation count").into());
        }

        let key_path = path.as_ref().join(KEYFILE);
        let key_file = File::open(&key_path);

        // Either load a store key or create a new store key if the key file
        // doesn't exist.
        let store_key = match key_file {
            Ok(k) => {
                let (_, key) = EncryptedMmapDirectory::load_store_key(k, passphrase)?;
                key
            }
            Err(e) => {
                if e.kind() != ErrorKind::NotFound {
                    return Err(e.into());
                }
                EncryptedMmapDirectory::create_new_store(
                    &key_path,
                    passphrase,
                    key_derivation_count,
                )?
            }
        };
        EncryptedMmapDirectory::new(store_key, path.as_ref())
    }

    /// Open a encrypted mmap directory.
    ///
    /// # Arguments
    ///
    /// * `path` - The path where the directory should reside in.
    /// * `passphrase` - The passphrase that was used to encrypt our directory.
    ///
    /// Returns an error if the path does not exist, if it is not a directory or
    /// if there was an error when trying to decrypt the directory key e.g. the
    /// given passphrase was incorrect.

    // This isn't currently used anywhere, but it will make sense if the
    // EncryptedMmapDirectory gets upstreamed.
    #[allow(dead_code)]
    pub fn open<P: AsRef<Path>>(path: P, passphrase: &str) -> Result<Self, OpenDirectoryError> {
        if passphrase.is_empty() {
            return Err(IoError::new(ErrorKind::Other, "empty passphrase").into());
        }

        let key_path = path.as_ref().join(KEYFILE);
        let key_file = File::open(&key_path)?;

        // Expand the store key into a encryption and MAC key.
        let (_, store_key) = EncryptedMmapDirectory::load_store_key(key_file, passphrase)?;
        EncryptedMmapDirectory::new(store_key, path.as_ref())
    }

    /// Change the passphrase that is used to encrypt the store key.
    /// This will decrypt and re-encrypt the store key using the new passphrase.
    ///
    /// # Arguments
    ///
    /// * `path` - The path where the directory resides in.
    /// * `old_passphrase` - The currently used passphrase.
    /// * `new_passphrase` - The passphrase that should be used from now on.
    /// * `new_key_derivation_count` - The key derivation count that should be
    /// used for the re-encrypted store key.
    pub fn change_passphrase<P: AsRef<Path>>(
        path: P,
        old_passphrase: &str,
        new_passphrase: &str,
        new_key_derivation_count: u32,
    ) -> Result<(), OpenDirectoryError> {
        if old_passphrase.is_empty() || new_passphrase.is_empty() {
            return Err(IoError::new(ErrorKind::Other, "empty passphrase").into());
        }
        if new_key_derivation_count == 0 {
            return Err(IoError::new(ErrorKind::Other, "invalid key derivation count").into());
        }

        let key_path = path.as_ref().join(KEYFILE);
        let key_file = File::open(&key_path)?;

        // Load our store key using the old passphrase.
        let (_, store_key) = EncryptedMmapDirectory::load_store_key(key_file, old_passphrase)?;
        // Derive new encryption keys using the new passphrase.
        let (key, hmac_key, salt) =
            EncryptedMmapDirectory::derive_key(new_passphrase, new_key_derivation_count)?;
        // Re-encrypt our store key using the newly derived keys.
        EncryptedMmapDirectory::encrypt_store_key(
            &key,
            &salt,
            new_key_derivation_count,
            &hmac_key,
            &store_key,
            &key_path,
        )?;

        Ok(())
    }

    /// Expand the given store key into an encryption key and HMAC key.
    fn expand_store_key(store_key: &[u8]) -> std::io::Result<KeyDerivationResult> {
        let mut hkdf_result = Zeroizing::new([0u8; KEY_SIZE * 2]);

        let hkdf = Hkdf::<Sha512>::new(None, store_key);
        hkdf.expand(&[], &mut *hkdf_result).map_err(|e| {
            IoError::new(
                ErrorKind::Other,
                format!("unable to expand store key: {:?}", e),
            )
        })?;
        let (key, hmac_key) = hkdf_result.split_at(KEY_SIZE);
        Ok((
            Zeroizing::new(Vec::from(key)),
            Zeroizing::new(Vec::from(hmac_key)),
        ))
    }

    /// Load a store key from the given file and decrypt it using the given
    /// passphrase.
    fn load_store_key(
        mut key_file: File,
        passphrase: &str,
    ) -> Result<(u32, KeyBuffer), OpenDirectoryError> {
        let mut iv = [0u8; IV_SIZE];
        let mut salt = [0u8; SALT_SIZE];
        let mut expected_mac = [0u8; MAC_LENGTH];
        let mut version = [0u8; 1];
        let mut encrypted_key = vec![];

        // Read our iv, salt, mac, and encrypted key from our key file.
        key_file.read_exact(&mut version)?;
        key_file.read_exact(&mut iv)?;
        key_file.read_exact(&mut salt)?;
        let pbkdf_count = key_file.read_u32::<BigEndian>()?;
        key_file.read_exact(&mut expected_mac)?;

        // Our key will be AES encrypted in CTR mode meaning the ciphertext
        // will have the same size as the plaintext. Read at most KEY_SIZE
        // bytes here so we don't end up filling up memory unnecessarily if
        // someone modifies the file.
        key_file
            .take(KEY_SIZE as u64)
            .read_to_end(&mut encrypted_key)?;

        if version[0] != VERSION {
            return Err(IoError::new(ErrorKind::Other, "invalid index store version").into());
        }

        // Re-derive our key using the passphrase and salt.
        let (key, hmac_key) = EncryptedMmapDirectory::rederive_key(passphrase, &salt, pbkdf_count);

        // First check our MAC of the encrypted key.
        let mac = EncryptedMmapDirectory::calculate_hmac(
            version[0],
            &iv,
            &salt,
            &encrypted_key,
            &hmac_key,
        )?;

        if mac.verify(&expected_mac).is_err() {
            return Err(IoError::new(ErrorKind::Other, "invalid MAC of the store key").into());
        }

        let mut decryptor = Aes256Ctr::new_from_slices(&key, &iv).map_err(|e| {
            IoError::new(
                ErrorKind::Other,
                format!("error initializing cipher {:?}", e),
            )
        })?;

        let mut out = Zeroizing::new(encrypted_key);
        decryptor.try_apply_keystream(&mut out).map_err(|_| {
            IoError::new(
                ErrorKind::Other,
                "Decryption error, reached end of the keystream.",
            )
        })?;

        Ok((pbkdf_count, out))
    }

    /// Calculate a HMAC for the given inputs.
    fn calculate_hmac(
        version: u8,
        iv: &[u8],
        salt: &[u8],
        encrypted_data: &[u8],
        hmac_key: &[u8],
    ) -> std::io::Result<Hmac<Sha256>> {
        let mut hmac = Hmac::<Sha256>::new_from_slice(hmac_key)
            .map_err(|e| IoError::new(ErrorKind::Other, format!("error creating hmac: {:?}", e)))?;
        hmac.update(&[version]);
        hmac.update(iv);
        hmac.update(salt);
        hmac.update(encrypted_data);
        Ok(hmac)
    }

    /// Create a new store key, encrypt it with the given passphrase and store
    /// it in the given path.
    fn create_new_store(
        key_path: &Path,
        passphrase: &str,
        pbkdf_count: u32,
    ) -> Result<KeyBuffer, OpenDirectoryError> {
        // Derive a AES key from our passphrase using a randomly generated salt
        // to prevent bruteforce attempts using rainbow tables.
        let (key, hmac_key, salt) = EncryptedMmapDirectory::derive_key(passphrase, pbkdf_count)?;
        // Generate a new random store key. This key will encrypt our Tantivy
        // indexing files. The key itself is stored encrypted using the derived
        // key.
        let store_key = EncryptedMmapDirectory::generate_key()?;

        // Encrypt and save the encrypted store key to a file.
        EncryptedMmapDirectory::encrypt_store_key(
            &key,
            &salt,
            pbkdf_count,
            &hmac_key,
            &store_key,
            key_path,
        )?;

        Ok(store_key)
    }

    /// Encrypt the given store key and save it in the given path.
    fn encrypt_store_key(
        key: &[u8],
        salt: &[u8],
        pbkdf_count: u32,
        hmac_key: &[u8],
        store_key: &[u8],
        key_path: &Path,
    ) -> Result<(), OpenDirectoryError> {
        // Generate a random initialization vector for our AES encryptor.
        let iv = EncryptedMmapDirectory::generate_iv()?;
        let mut encryptor = Aes256Ctr::new_from_slices(key, &iv).map_err(|e| {
            IoError::new(
                ErrorKind::Other,
                format!("error initializing cipher: {:?}", e),
            )
        })?;

        let mut encrypted_key = [0u8; KEY_SIZE];
        encrypted_key.copy_from_slice(store_key);

        let mut key_file = File::create(key_path)?;

        // Write down our public salt and iv first, those will be needed to
        // decrypt the key again.
        key_file.write_all(&[VERSION])?;
        key_file.write_all(&iv)?;
        key_file.write_all(salt)?;
        key_file.write_u32::<BigEndian>(pbkdf_count)?;

        // Encrypt our key.
        encryptor
            .try_apply_keystream(&mut encrypted_key)
            .map_err(|e| {
                IoError::new(
                    ErrorKind::Other,
                    format!("unable to encrypt store key: {:?}", e),
                )
            })?;

        // Calculate a MAC for our encrypted key and store it in the file before
        // the key.
        let mac =
            EncryptedMmapDirectory::calculate_hmac(VERSION, &iv, salt, &encrypted_key, hmac_key)?;
        let mac = mac.finalize();
        let mac = mac.into_bytes();
        key_file.write_all(mac.as_slice())?;

        // Write down the encrypted key.
        key_file.write_all(&encrypted_key)?;

        Ok(())
    }

    /// Generate a random IV.
    fn generate_iv() -> Result<[u8; IV_SIZE], OpenDirectoryError> {
        let mut iv = [0u8; IV_SIZE];
        let mut rng = thread_rng();
        rng.try_fill(&mut iv[..])
            .map_err(|e| IoError::new(ErrorKind::Other, format!("error generating iv: {:?}", e)))?;
        Ok(iv)
    }

    /// Generate a random key.
    fn generate_key() -> Result<KeyBuffer, OpenDirectoryError> {
        let mut key = Zeroizing::new(vec![0u8; KEY_SIZE]);
        let mut rng = thread_rng();
        rng.try_fill(&mut key[..]).map_err(|e| {
            IoError::new(ErrorKind::Other, format!("error generating key: {:?}", e))
        })?;
        Ok(key)
    }

    /// Derive two keys from the given passphrase and the given salt using PBKDF2.
    fn rederive_key(passphrase: &str, salt: &[u8], pbkdf_count: u32) -> KeyDerivationResult {
        let mut pbkdf_result = Zeroizing::new([0u8; KEY_SIZE * 2]);

        pbkdf2::<Hmac<Sha512>>(passphrase.as_bytes(), salt, pbkdf_count, &mut *pbkdf_result);
        let (key, hmac_key) = pbkdf_result.split_at(KEY_SIZE);
        (
            Zeroizing::new(Vec::from(key)),
            Zeroizing::new(Vec::from(hmac_key)),
        )
    }

    /// Generate a random salt and derive two keys from the salt and the given
    /// passphrase.
    fn derive_key(
        passphrase: &str,
        pbkdf_count: u32,
    ) -> Result<InitialKeyDerivationResult, OpenDirectoryError> {
        let mut rng = thread_rng();
        let mut salt = vec![0u8; SALT_SIZE];
        rng.try_fill(&mut salt[..]).map_err(|e| {
            IoError::new(ErrorKind::Other, format!("error generating salt: {:?}", e))
        })?;

        let (key, hmac_key) = EncryptedMmapDirectory::rederive_key(passphrase, &salt, pbkdf_count);
        Ok((key, hmac_key, salt))
    }
}

// The Directory trait[dr] implementation for our EncryptedMmapDirectory.
// [dr] https://docs.rs/tantivy/0.10.2/tantivy/directory/trait.Directory.html
impl Directory for EncryptedMmapDirectory {
    fn open_read(&self, path: &Path) -> Result<ReadOnlySource, OpenReadError> {
        let source = self.mmap_dir.open_read(path)?;

        let mut reader = AesReader::<Aes256Ctr, _>::new::<Hmac<Sha256>>(
            Cursor::new(source.as_slice()),
            &self.encryption_key,
            &self.mac_key,
            IV_SIZE,
            MAC_LENGTH,
        )
        .map_err(TvIoError::from)?;

        let mut decrypted = Vec::new();
        reader
            .read_to_end(&mut decrypted)
            .map_err(TvIoError::from)?;

        Ok(ReadOnlySource::from(decrypted))
    }

    fn delete(&self, path: &Path) -> Result<(), DeleteError> {
        self.mmap_dir.delete(path)
    }

    fn exists(&self, path: &Path) -> bool {
        self.mmap_dir.exists(path)
    }

    fn open_write(&mut self, path: &Path) -> Result<WritePtr, OpenWriteError> {
        let file = match self.mmap_dir.open_write(path)?.into_inner() {
            Ok(f) => f,
            Err(e) => {
                let error = IoError::from(e);
                return Err(TvIoError::from(error).into());
            }
        };

        let writer = AesWriter::<Aes256Ctr, Hmac<Sha256>, _>::new(
            file,
            &self.encryption_key,
            &self.mac_key,
            IV_SIZE,
        )
        .map_err(TvIoError::from)?;
        Ok(BufWriter::new(Box::new(writer)))
    }

    fn atomic_read(&self, path: &Path) -> Result<Vec<u8>, OpenReadError> {
        let data = self.mmap_dir.atomic_read(path)?;

        let mut reader = AesReader::<Aes256Ctr, _>::new::<Hmac<Sha256>>(
            Cursor::new(data),
            &self.encryption_key,
            &self.mac_key,
            IV_SIZE,
            MAC_LENGTH,
        )
        .map_err(TvIoError::from)?;
        let mut decrypted = Vec::new();

        reader
            .read_to_end(&mut decrypted)
            .map_err(TvIoError::from)?;
        Ok(decrypted)
    }

    fn atomic_write(&mut self, path: &Path, data: &[u8]) -> std::io::Result<()> {
        let mut encrypted = Vec::new();
        {
            let mut writer = AesWriter::<Aes256Ctr, Hmac<Sha256>, _>::new(
                &mut encrypted,
                &self.encryption_key,
                &self.mac_key,
                IV_SIZE,
            )?;
            writer.write_all(data)?;
        }

        self.mmap_dir.atomic_write(path, &encrypted)
    }

    fn watch(&self, watch_callback: WatchCallback) -> Result<WatchHandle, tantivy::TantivyError> {
        self.mmap_dir.watch(watch_callback)
    }

    fn acquire_lock(&self, lock: &Lock) -> Result<DirectoryLock, LockError> {
        // The lock files aren't encrypted, this is fine since they won't
        // contain any data. They will be an empty file and a lock will be
        // placed on them using e.g. flock(2) on macOS and Linux.
        self.mmap_dir.acquire_lock(lock)
    }
}

// This Tantivy trait is used to indicate when no more writes are expected to be
// done on a writer.
impl<E: NewCipher + StreamCipher, M: Mac + NewMac, W: Write> TerminatingWrite
    for AesWriter<E, M, W>
{
    fn terminate_ref(&mut self, _: AntiCallToken) -> std::io::Result<()> {
        self.finalize()
    }
}

#[cfg(test)]
use tempfile::tempdir;

#[test]
fn create_new_store_and_reopen() {
    let tmpdir = tempdir().unwrap();
    let dir = EncryptedMmapDirectory::open_or_create(tmpdir.path(), "wordpass", PBKDF_COUNT)
        .expect("Can't create a new store");
    drop(dir);
    let dir = EncryptedMmapDirectory::open(tmpdir.path(), "wordpass")
        .expect("Can't open the existing store");
    drop(dir);
    let dir = EncryptedMmapDirectory::open(tmpdir.path(), "password");
    assert!(
        dir.is_err(),
        "Opened an existing store with the wrong passphrase"
    );
}

#[test]
fn create_store_with_empty_passphrase() {
    let tmpdir = tempdir().unwrap();
    let dir = EncryptedMmapDirectory::open(tmpdir.path(), "");
    assert!(
        dir.is_err(),
        "Opened an existing store with the wrong passphrase"
    );
}

#[test]
fn change_passphrase() {
    let tmpdir = tempdir().unwrap();
    let dir = EncryptedMmapDirectory::open_or_create(tmpdir.path(), "wordpass", PBKDF_COUNT)
        .expect("Can't create a new store");

    drop(dir);
    EncryptedMmapDirectory::change_passphrase(tmpdir.path(), "wordpass", "password", PBKDF_COUNT)
        .expect("Can't change passphrase");
    let dir = EncryptedMmapDirectory::open(tmpdir.path(), "wordpass");
    assert!(
        dir.is_err(),
        "Opened an existing store with the old passphrase"
    );
    let _ = EncryptedMmapDirectory::open(tmpdir.path(), "password")
        .expect("Can't open the store with the new passphrase");
}