snapper-box 0.0.4

//! Middle level abstraction logic.
//!
//! This module models a file as an LSM, with encrypted updates

use std::{
    collections::HashMap,
    fs::{File, OpenOptions},
    hash::Hash,
    io::{Read, Seek, SeekFrom, Write},
    path::Path,
    sync::Arc,
};

use serde::{de::DeserializeOwned, Deserialize, Serialize};
use serde_bytes::ByteBuf;
use snafu::ResultExt;
use tracing::{debug, instrument, warn};

use crate::{
    crypto::{CipherText, ClearText, Key},
    error::{BackendError, EntryIO},
    file::segment::Segment,
};
mod segment;

/// Internal, data containing struct used for serializing log entries
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
struct Entry {
    /// Items being stored
    items: HashMap<ByteBuf, ByteBuf>,
}

/// LSM Abstraction over something [`File`] like (i.e. implementing [`Read`], [`Write`], and [`Seek`]).
///
/// This type maintains a write cache in memory, and you _must_ call the `flush` method for the changes
/// to be persisted. This type _will not_ automatically flush.l
pub struct LsmFile<T, K>
where
    T: Read + Write + Seek,
    K: Key,
{
    /// Underlying [`File`](std::fs::File) like object
    file: T,
    /// Segment offset table
    offsets: HashMap<Vec<u8>, u64>,
    /// Cache containing yet-to-be-written items
    cache: Option<Entry>,
    /// compression level for the file
    compression: Option<i32>,
    /// Key for the file
    key: Arc<K>,
    /// Maximum number of cache entries
    max_cache_entries: usize,
}

impl<T, K> LsmFile<T, K>
where
    T: Read + Write + Seek,
    K: Key,
{
    /// Creates a new `LsmFile` from the provided, already existing, [`File`] like object.
    ///
    /// Accepts values for the compression level for the file, the key to be used for encryption/decryption,
    /// as well as the maximum number of cache entries before a mandatory flush occurs, which defaults to 100.
    ///
    /// # Errors
    ///
    ///   * `Error::EntryIO` if an error occurs while reading the entries
    ///   * Will bubble up any segment errors`
    #[instrument(skip(file, key))]
    pub fn new(
        mut file: T,
        compression: Option<i32>,
        key: Arc<K>,
        max_cache_entries: Option<usize>,
    ) -> Result<Self, BackendError> {
        // Seek to start of file
        file.seek(SeekFrom::Start(0)).context(EntryIO)?;
        // Start trying to read segments
        let mut offset: u64 = 0;
        let mut offsets: HashMap<Vec<u8>, u64> = HashMap::new();
        while let Ok(segment) = Segment::read_owned(&mut file) {
            match CipherText::try_from(segment) {
                Ok(x) => {
                    let cleartext = x.decrypt(&*key)?;
                    let entry: Entry = cleartext.deserialize()?;
                    for (k, _) in entry.items {
                        offsets.insert(k.to_vec(), offset);
                    }
                }
                Err(e) => {
                    warn!(?e, "Failed to decode a segment");
                }
            }
            offset = file.stream_position().context(EntryIO)?;
        }
        Ok(Self {
            file,
            offsets,
            cache: None,
            compression,
            key,
            max_cache_entries: max_cache_entries.unwrap_or(100),
        })
    }

    /// Retries a value from the store given a key, if the value exists in the store.
    ///
    /// The most recently inserted version of the value will be returned
    ///
    /// # Errors
    ///
    /// Will return an error if any IO error occurs, or if the value fails to deserialize.
    pub fn get<C, V>(&mut self, key: &C) -> Result<Option<V>, BackendError>
    where
        C: Serialize,
        V: DeserializeOwned,
    {
        // Attempt to serialize the key
        let key = match serde_cbor::to_vec(key) {
            Ok(k) => ByteBuf::from(k),
            // Intentionally hide the underlying serde error, to avoid leaking sensitive data into
            // the logs
            Err(_) => return Err(BackendError::ItemSerialization),
        };
        // Look up the key in the cache, return from it if possible
        if let Some(cache) = self.cache.as_ref() {
            if let Some(bytes) = cache.items.get(&key) {
                match serde_cbor::from_slice(bytes) {
                    Ok(v) => return Ok(Some(v)),
                    // Intentionally hide the underlying serde error, to avoid leaking sensitive
                    // data into the logs
                    Err(_) => return Err(BackendError::ItemDeserialization),
                }
            }
        }
        // Look up the key in the offsets table
        if let Some(offset) = self.offsets.get(&*key) {
            // Find the segment
            self.file.seek(SeekFrom::Start(*offset)).context(EntryIO)?;
            let segment = Segment::read_owned(&mut self.file)?;
            // Decrypt it
            let ciphertext = CipherText::try_from(segment)?;
            let cleartext = ciphertext.decrypt(&*self.key)?;
            let entry: Entry = cleartext.deserialize()?;
            if let Some(bytes) = entry.items.get(&key) {
                match serde_cbor::from_slice(bytes) {
                    Ok(v) => Ok(Some(v)),
                    // Intentionally hide the underlying serde error, to avoid leaking sensitive
                    // data into the logs
                    Err(_) => Err(BackendError::ItemDeserialization),
                }
            } else {
                warn!("Offsets table contained a offset fora nonexistent value");
                Ok(None)
            }
        } else {
            Ok(None)
        }
    }

    /// Inserts an object into the store
    ///
    /// This method will write to the cache, however, if this insertion makes the cache size greater than or
    /// equal to the maximum number of entries, the current cache will be flushed.
    ///
    /// # Errors
    ///   * Will return an error if either the key or value fail to serialize
    ///   * Will bubble up any IO errors that occur, if a flush occurs
    #[instrument(skip(self, key, value))]
    pub fn insert<C, V>(&mut self, key: &C, value: &V) -> Result<(), BackendError>
    where
        C: Serialize,
        V: Serialize,
    {
        // Get the cache
        let mut cache = self.cache.take().unwrap_or_else(|| Entry {
            items: HashMap::new(),
        });
        // Serialize the values
        // Intentionally hide the serde errors, to avoid leaking sensitive information
        let key = serde_cbor::to_vec(key).map_err(|_| BackendError::ItemSerialization)?;
        let value = serde_cbor::to_vec(value).map_err(|_| BackendError::ItemSerialization)?;
        // Insert it into the cache
        cache.items.insert(ByteBuf::from(key), ByteBuf::from(value));
        let length = cache.items.len();
        self.cache = Some(cache);
        // Flush if needed
        if length >= self.max_cache_entries {
            debug!("Flushing cache");
            self.flush()?;
        }
        Ok(())
    }

    /// Flushes the cache to disk
    ///
    /// # Errors
    ///
    /// Will bubble up any IO errors
    pub fn flush(&mut self) -> Result<(), BackendError> {
        // Get the end of the file
        let offset = self.file.seek(SeekFrom::End(0)).context(EntryIO)?;
        // Take the cache
        if let Some(cache) = self.cache.take() {
            // Update the offset table
            for k in cache.items.keys() {
                self.offsets.insert((&*k).to_vec(), offset);
            }
            // Serialize it and encrypt it
            let plaintext = ClearText::new(&cache)?;
            let ciphertext = plaintext.encrypt(&*self.key, self.compression)?;
            let segment: Segment<'_> = ciphertext.into();
            segment.write(&mut self.file)?;
            self.file.flush().context(EntryIO)?;

            Ok(())
        } else {
            Ok(())
        }
    }

    /// Converts the current contents of the store to a [`HashMap`].
    ///
    /// Iterates through this `LsmFIle`'s keys, and insert them all into the returned [`HashMap`]. The most
    /// up to date value will be used for each key.
    ///
    /// # Errors
    ///
    /// Will bubble up any error that occurs during a `get` operation
    pub fn to_hashmap<C, V>(&mut self) -> Result<HashMap<C, V>, BackendError>
    where
        C: DeserializeOwned + Serialize + Hash + Eq,
        V: DeserializeOwned,
    {
        let mut ret = HashMap::new();
        let keys: Vec<Vec<u8>> = self.offsets.keys().cloned().collect();
        // Insert from storage
        for key in keys {
            // Deserialize the key
            let key: C =
                serde_cbor::from_slice(&key).map_err(|_| BackendError::ItemDeserialization)?;
            // get the value
            let value = self.get(&key)?.ok_or(BackendError::InvalidLsmState)?;
            ret.insert(key, value);
        }
        // Insert from cache
        if let Some(cache) = self.cache.as_ref() {
            let keys: Vec<ByteBuf> = cache.items.keys().cloned().collect();
            for key in keys {
                // Deserialize the key
                let key_deser: C =
                    serde_cbor::from_slice(&key).map_err(|_| BackendError::ItemDeserialization)?;
                // get the value
                let value = cache.items.get(&key).ok_or(BackendError::InvalidLsmState)?;
                let value: V =
                    serde_cbor::from_slice(value).map_err(|_| BackendError::ItemDeserialization)?;
                ret.insert(key_deser, value);
            }
        }
        Ok(ret)
    }
    /// Converts the current contents of the store to a [`Vec`] of pairs.
    ///
    /// Iterates through this `LsmFIle`'s keys, and insert them all into the returned [`Vec`]]. The most
    /// up to date value will be used for each key.
    ///
    /// # Errors
    ///
    /// Will bubble up any error that occurs during a `get` operation
    pub fn to_pairs<C, V>(&mut self) -> Result<Vec<(C, V)>, BackendError>
    where
        C: DeserializeOwned + Serialize + Eq + Clone,
        V: DeserializeOwned + Clone,
    {
        let mut ret: Vec<(C, V)> = Vec::new();
        let keys: Vec<Vec<u8>> = self.offsets.keys().cloned().collect();
        // Insert from storage
        for key in keys {
            // Deserialize the key
            let key: C =
                serde_cbor::from_slice(&key).map_err(|_| BackendError::ItemDeserialization)?;
            // get the value
            let value = self.get(&key)?.ok_or(BackendError::InvalidLsmState)?;
            if let Some(index) =
                ret.iter()
                    .cloned()
                    .enumerate()
                    .find_map(|(x, (k, _))| if k == key { Some(x) } else { None })
            {
                ret[index] = (key, value);
            } else {
                ret.push((key, value));
            }
        }
        // Insert from cache
        if let Some(cache) = self.cache.as_ref() {
            let keys: Vec<ByteBuf> = cache.items.keys().cloned().collect();
            for key in keys {
                // Deserialize the key
                let key_deser: C =
                    serde_cbor::from_slice(&key).map_err(|_| BackendError::ItemDeserialization)?;
                // get the value
                let value = cache.items.get(&key).ok_or(BackendError::InvalidLsmState)?;
                let value: V =
                    serde_cbor::from_slice(value).map_err(|_| BackendError::ItemDeserialization)?;
                if let Some(index) = ret.iter().cloned().enumerate().find_map(|(x, (k, _))| {
                    if k == key_deser {
                        Some(x)
                    } else {
                        None
                    }
                }) {
                    ret[index] = (key_deser, value);
                } else {
                    ret.push((key_deser, value));
                }
            }
        }
        Ok(ret)
    }
    /// Consumes self and returns the inner [`File`] like object
    pub fn into_inner(self) -> T {
        self.file
    }
}

impl<K> LsmFile<File, K>
where
    K: Key,
{
    /// Opens an existing `LsmFile` from the provided path.
    ///
    /// Will open the file in read/write mode. Will fail if the file does not exists.
    ///
    /// # Errors
    ///   * If the file does not exist
    ///   * If any other IO occurs
    ///   * If any of the decryption or deserialization of control structures fails
    #[instrument(skip(key))]
    pub fn open(
        path: impl AsRef<Path> + std::fmt::Debug,
        compression: Option<i32>,
        key: Arc<K>,
        max_cache_entries: Option<usize>,
    ) -> Result<Self, BackendError> {
        // open the file read/write
        let file = OpenOptions::new()
            .read(true)
            .write(true)
            .create(false)
            .open(path.as_ref())
            .context(EntryIO)?;
        Self::new(file, compression, key, max_cache_entries)
    }

    /// Opens an existing `LsmFile` from the provided path.
    ///
    /// Will create the file in read/write mode, failing if the file exists
    ///
    /// # Errors
    ///   * If the file already exists
    ///   * If any other IO occurs
    #[instrument(skip(key))]
    pub fn create(
        path: impl AsRef<Path> + std::fmt::Debug,
        compression: Option<i32>,
        key: Arc<K>,
        max_cache_entries: Option<usize>,
    ) -> Result<Self, BackendError> {
        // open the file read/write
        let file = OpenOptions::new()
            .read(true)
            .write(true)
            .create_new(true)
            .open(path.as_ref())
            .context(EntryIO)?;
        Self::new(file, compression, key, max_cache_entries)
    }
}

/// Unit tests for the module
#[cfg(test)]
mod tests {
    use super::*;
    use crate::crypto::RootKey;
    use proptest::prelude::*;
    use std::io::Cursor;
    use tempfile::{tempdir, tempfile};
    /// Make sure that `LsmFile` can create and open files
    #[test]
    fn file_create_open() -> Result<(), BackendError> {
        let dir = tempdir().expect("Unable to make tempdir");
        let file_path = dir.path().join("test_lsm");
        println!("{:?}", file_path);
        // Create the comparison hashmap
        let hashmap: HashMap<u64, u64> = [(1, 2), (3, 4), (5, 6), (7, 8)].into_iter().collect();
        // Get a key
        let root_key = Arc::new(RootKey::random());
        // Create the LsmFile
        let mut lsm_file = LsmFile::create(&file_path, None, root_key.clone(), Some(3))?;
        // Insert all the k/v pairs into the LsmFile
        for (k, v) in &hashmap {
            println!("k: {} v: {}", k, v);
            lsm_file.insert(&k, &v).expect("Unable to insert k/v pair");
        }
        // Flush the file and drop it
        lsm_file.flush()?;
        std::mem::drop(lsm_file);
        // Reopen it
        let mut lsm_file = LsmFile::open(&file_path, None, root_key, Some(3))?;
        // Get the output hashmap
        let output: HashMap<u64, u64> = lsm_file.to_hashmap()?;
        assert_eq!(output, hashmap);
        Ok(())
    }

    proptest! {
        /// Use a random list of pairs to insert, and test the behavior compared to a hash map
        #[test]
        fn round_trip_hashmap_smoke(pairs: Vec<(u64,u64)>) {
            // Get the comparison HashMap
            let hashmap = pairs.iter().copied().collect::<HashMap<u64,u64>>();
            // Open up the `LsmFile`
            let backing = Cursor::new(Vec::<u8>::new());
            let root_key = Arc::new(RootKey::random());
            let mut lsm_file: LsmFile<_, RootKey> =
                LsmFile::new(backing, None, root_key, Some(10))
                .expect("Failed to open file");
            // Insert all the k/v pairs into the LsmFile
            for (k,v) in pairs {
                println!("k: {} v: {}", k, v);
                lsm_file.insert(&k, &v).expect("Unable to insert k/v pair");
            }
            // Get the output hashmap
            let output: HashMap<u64,u64> = lsm_file.to_hashmap().expect("Failed to covert to hashmap");
            for (k,v) in &output {
                if Some(v) != hashmap.get(k) {
                    panic!("Output hashmap contains k/v pair not in input: k: {} v: {}", k, v);
                }
            }
            for (k,v) in &hashmap {
                if Some(v) != output.get(k) {
                    panic!("Input hashmap contains k/v pair not in output: k: {} v: {}", k, v);
                }
            }
        }
        /// Use a random list of pairs to insert, and test the behavior compared to a hash map, deconstructing
        /// and reconstructing the `LsmFile` before reading the output hashmap.
        #[test]
        fn round_trip_hashmap_flush(pairs: Vec<(u64,u64)>) {
            // Get the comparison HashMap
            let hashmap = pairs.iter().copied().collect::<HashMap<u64,u64>>();
            // Open up the `LsmFile`
            let backing = Cursor::new(Vec::<u8>::new());
            let root_key = Arc::new(RootKey::random());
            let mut lsm_file: LsmFile<_, RootKey> =
                LsmFile::new(backing, None, root_key.clone(), Some(10))
                .expect("Failed to open file");
            // Insert all the k/v pairs into the LsmFile
            for (k,v) in pairs {
                println!("k: {} v: {}", k, v);
                lsm_file.insert(&k, &v).expect("Unable to insert k/v pair");
            }
            // Flush and reconstruct the lsm_file
            lsm_file.flush().expect("Failed to flush lsm file");
            let backing = lsm_file.into_inner();
            let mut lsm_file: LsmFile<_, RootKey> =
                LsmFile::new(backing, None, root_key, Some(10))
                .expect("Failed to open file");
            // Get the output hashmap
            let output: HashMap<u64,u64> = lsm_file.to_hashmap().expect("Failed to covert to hashmap");
            for (k,v) in &output {
                if Some(v) != hashmap.get(k) {
                    panic!("Output hashmap contains k/v pair not in input: k: {} v: {}", k, v);
                }
            }
            for (k,v) in &hashmap {
                if Some(v) != output.get(k) {
                    panic!("Input hashmap contains k/v pair not in output: k: {} v: {}", k, v);
                }
            }
        }
        /// Use a random list of pairs to insert, and test the behavior compared to a hash map, deconstructing
        /// and reconstructing the `LsmFile` before reading the output hashmap.
        ///
        /// This variant of the test uses an actual temporary file
        #[test]
        fn round_trip_hashmap_file(pairs: Vec<(u64,u64)>) {
            // Get the comparison HashMap
            let hashmap = pairs.iter().copied().collect::<HashMap<u64,u64>>();
            // Open up the `LsmFile`
            let backing = tempfile().expect("Failed to open tempfile");
            let root_key = Arc::new(RootKey::random());
            let mut lsm_file: LsmFile<_, RootKey> =
                LsmFile::new(backing, None, root_key.clone(), Some(10))
                .expect("Failed to open file");
            // Insert all the k/v pairs into the LsmFile
            for (k,v) in pairs {
                println!("k: {} v: {}", k, v);
                lsm_file.insert(&k, &v).expect("Unable to insert k/v pair");
            }
            // Flush and reconstruct the lsm_file
            lsm_file.flush().expect("Failed to flush lsm file");
            let backing = lsm_file.into_inner();
            let mut lsm_file: LsmFile<_, RootKey> =
                LsmFile::new(backing, None, root_key, Some(10))
                .expect("Failed to open file");
            // Get the output hashmap
            let output: HashMap<u64,u64> = lsm_file.to_hashmap().expect("Failed to covert to hashmap");
            for (k,v) in &output {
                if Some(v) != hashmap.get(k) {
                    panic!("Output hashmap contains k/v pair not in input: k: {} v: {}", k, v);
                }
            }
            for (k,v) in &hashmap {
                if Some(v) != output.get(k) {
                    panic!("Input hashmap contains k/v pair not in output: k: {} v: {}", k, v);
                }
            }
        }
    }
}