infinitree 0.11.0

use super::{ObjectError, ObjectId, Result, WriteObject};
use crate::{
    backends::Backend,
    compress,
    crypto::{ChunkKey, CryptoOps, Digest, IndexKey, StorageKey},
    ChunkPointer,
};
use ring::rand::{SecureRandom, SystemRandom};
use std::{
    io::{Seek, SeekFrom},
    sync::Arc,
};

#[derive(Clone)]
enum Mode {
    SealRoot(u64),
    Data,
}

impl Mode {
    fn skip(&self) -> u64 {
        match self {
            Self::SealRoot(skip) => *skip,
            Self::Data => 0,
        }
    }
}

pub trait Writer: Send {
    fn write(&mut self, data: &[u8]) -> Result<ChunkPointer>;
    fn write_chunk(&mut self, hash: &Digest, data: &[u8]) -> Result<ChunkPointer>;
    fn flush(&mut self) -> Result<()>;
}

pub struct AEADWriter {
    backend: Arc<dyn Backend>,
    random: SystemRandom,
    crypto: CryptoOps,
    object: WriteObject,
    mode: Mode,
    rewrite: Vec<ObjectId>,
}

impl AEADWriter {
    pub(crate) fn new(backend: Arc<dyn Backend>, crypto: ChunkKey) -> Self {
        let mut object = WriteObject::default();
        let random = SystemRandom::new();
        reset_id(&mut object, &random);

        AEADWriter {
            backend,
            object,
            random,
            crypto: crypto.into_inner(),
            mode: Mode::Data,
            rewrite: vec![],
        }
    }

    pub(crate) fn for_storage(backend: Arc<dyn Backend>, crypto: StorageKey) -> Self {
        let mut object = WriteObject::default();
        let random = SystemRandom::new();
        reset_id(&mut object, &random);

        AEADWriter {
            backend,
            object,
            random,
            crypto: crypto.into_inner(),
            mode: Mode::Data,
            rewrite: vec![],
        }
    }

    pub(crate) fn for_root(
        backend: Arc<dyn Backend>,
        crypto: IndexKey,
        header_size: u64,
        mut rewrite: Vec<ObjectId>,
    ) -> Self {
        let random = SystemRandom::new();
        let mut object = WriteObject::default();
        object.seek(SeekFrom::Start(header_size)).unwrap();

        rewrite_or_reset_id(&mut object, &random, &mut rewrite);

        AEADWriter {
            backend,
            object,
            rewrite,
            random,
            crypto: crypto.into_inner(),
            mode: Mode::SealRoot(header_size),
        }
    }

    pub(crate) fn flush_root_head(&mut self, id: ObjectId, head: &[u8]) -> Result<()> {
        self.object.set_id(id);
        self.write_head(head);
        self.finalize()?;
        self.backend.write_object(&self.object)?;

        self.rewrite.clear();
        reset_id(&mut self.object, &self.random);
        self.object.seek(SeekFrom::Start(self.mode.skip()))?;

        Ok(())
    }

    fn finalize(&mut self) -> Result<()> {
        self.random
            .fill(self.object.tail_mut())
            .map_err(|_| ObjectError::Fatal)?;
        Ok(())
    }

    fn write_head(&mut self, content: &[u8]) {
        self.object.head_mut(content.len()).copy_from_slice(content);
    }
}

impl Clone for AEADWriter {
    fn clone(&self) -> Self {
        let mut object = self.object.clone();
        reset_id(&mut object, &self.random);

        AEADWriter {
            object,
            random: self.random.clone(),
            backend: self.backend.clone(),
            crypto: self.crypto.clone(),
            mode: self.mode.clone(),
            rewrite: vec![],
        }
    }
}

impl Drop for AEADWriter {
    fn drop(&mut self) {
        if self.object.position() > self.mode.skip() as usize {
            self.flush().unwrap();
        }
    }
}

impl Writer for AEADWriter {
    fn write_chunk(&mut self, hash: &Digest, data: &[u8]) -> Result<ChunkPointer> {
        let size = {
            let buffer = self.object.tail_mut();

            match compress::compress_into(data, buffer) {
                Ok(size) => size,
                Err(_e) => {
                    self.flush()?;

                    let buffer = self.object.tail_mut();
                    compress::compress_into(data, buffer).map_err(|_| {
                        ObjectError::ChunkTooLarge {
                            size: data.len(),
                            max_size: ((self.object.capacity() - 16 - 4) as f64 / 1.1) as usize,
                        }
                    })?
                }
            }
        };

        let oid = *self.object.id();
        let pointer = {
            let offs = self.object.position() as u32;
            let buffer = self.object.tail_mut();

            self.crypto
                .encrypt_chunk(oid, offs, hash, &mut buffer[..size])
        };

        *self.object.position_mut() += pointer.size();

        Ok(pointer)
    }

    fn flush(&mut self) -> Result<()> {
        if let Mode::SealRoot(header_size) = self.mode {
            self.random
                .fill(self.object.head_mut(header_size.try_into().unwrap()))
                .map_err(|_| ObjectError::Fatal)?;
        }
        self.finalize()?;
        self.backend.write_object(&self.object)?;

        rewrite_or_reset_id(&mut self.object, &self.random, &mut self.rewrite);

        self.object.seek(SeekFrom::Start(self.mode.skip()))?;

        Ok(())
    }

    fn write(&mut self, data: &[u8]) -> Result<ChunkPointer> {
        self.write_chunk(&self.crypto.hash(data), data)
    }
}

#[inline(always)]
fn rewrite_or_reset_id(
    object: &mut WriteObject,
    random: &impl SecureRandom,
    rewrite: &mut Vec<ObjectId>,
) {
    object.set_id(rewrite.pop().unwrap_or_else(|| ObjectId::new(random)));
}

#[inline(always)]
fn reset_id(object: &mut WriteObject, random: &impl SecureRandom) {
    object.set_id(ObjectId::new(random));
}