bsv/libbsv/src/simple_db/object.rs

// This file is part of bsv.
//
// bsv is free software: you can redistribute it and/or modify it under the
// terms of the GNU Affero General Public License as published by the Free
// Software Foundation, either version 3 of the License, or (at your option)
// any later version.
//
// cdb is distributed in the hope that it will be useful, but WITHOUT ANY
// WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public License for
// more details.
//
// You should have received a copy of the Affero GNU General Public License
// along with cdb.  If not, see <https://www.gnu.org/licenses/>.


extern crate digest;
extern crate sha2;
extern crate flate2;

use std::io::{Read, Seek, SeekFrom, Write, copy, sink};

use digest::Digest;

use flate2::{
    Compression,
    write::GzEncoder,
    read::GzDecoder,
};

use crate::core::error::*;
use crate::core::ObjectId;
use crate::core::ObjectType;


pub struct ObjectWriter<W: Write, D: Digest> {
    writer: GzEncoder<W>,
    digest: D,
    size: u64,
    written_size: u64,
}


impl<W: Write, D: Digest> ObjectWriter<W, D> {
    pub fn new(writer: W, otype: &ObjectType, size: u64)
        -> Result<ObjectWriter<W, D>> {

        let mut digest = D::new();

        digest.update(otype);
        digest.update(&size.to_be_bytes());

        let mut zwriter = GzEncoder::new(writer, Compression::default());

        zwriter.write_all(otype)?;
        zwriter.write_all(&size.to_be_bytes())?;

        Ok(ObjectWriter {
            writer: zwriter,
            digest,
            size,
            written_size: 0,
        })
    }

    pub fn finish(mut self) -> Result<(ObjectId, W)> {
        self.writer.try_finish()?;

        if self.written_size == self.size {
            Ok((
                ObjectId::new(&self.digest.finalize()),
                self.writer.finish()?,
            ))
        }
        else {
            Err(ErrorKind::MismatchingObjectSize(self.written_size, self.size).into())
        }
    }
}


impl<W: Write, D: Digest> Write for ObjectWriter<W, D> {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        let size = self.writer.write(buf)?;
        self.digest.update(&buf[..size]);
        self.written_size += size as u64;
        Ok(size)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.writer.flush()
    }
}


pub struct ObjectReader<R: Read> {
    otype: ObjectType,
    size: u64,
    reader: GzDecoder<R>,
}


impl<R: Read> ObjectReader<R> {
    pub fn new(reader: R)
        -> Result<ObjectReader<R>> {

        let mut zreader = GzDecoder::new(reader);

        let mut buffer = [0u8; 12];
        zreader.read_exact(&mut buffer)?;

        let otype = {
            let mut otype = [0; 4];
            otype.copy_from_slice(&buffer[0..4]);
            otype
        };
        let size = {
            let mut size_bytes = [0; 8];
            size_bytes.copy_from_slice(&buffer[4..12]);
            u64::from_be_bytes(size_bytes)
        };

        Ok(ObjectReader {
            otype,
            size,
            reader: zreader,
        })
    }

    pub fn object_type(&self) -> ObjectType {
        self.otype
    }

    pub fn size(&self) -> u64 {
        self.size
    }

    pub fn close(self) -> Result<R> {
        Ok(self.reader.into_inner())
    }
}


impl<R: Read> Read for ObjectReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.reader.read(buf)
    }
}


pub fn write_object<R, W>(reader: &mut R, writer: W,
        otype: &ObjectType, size: u64)
    -> Result<(ObjectId, W)> where
        R: Read,
        W: Write {
    let mut owriter: ObjectWriter<W, sha2::Sha224>
        = ObjectWriter::new(writer, otype, size)?;

    copy(reader, &mut owriter)?;

    owriter.finish()
}


pub trait WriteAsObject {
    fn write_as_object<W: Write>(&mut self, writer: W, otype: &ObjectType) -> Result<ObjectId>;

    fn object_id(&mut self, otype: &ObjectType) -> Result<ObjectId> {
        self.write_as_object(sink(), otype)
    }
}


impl<T: Read + Seek> WriteAsObject for T {
    fn write_as_object<W: Write>(&mut self, writer: W, otype: &ObjectType) -> Result<ObjectId> {
        let start = self.seek(SeekFrom::Current(0))?;
        let end = self.seek(SeekFrom::End(0))?;
        self.seek(SeekFrom::Start(start))?;

        let (oid, _) = write_object(self, writer, otype, end - start)?;
        Ok(oid)
    }
}


#[cfg(test)]
mod tests {
    use std::str::{FromStr};

    use crate::core::OTYPE_BLOB;

    use super::*;

    const PAYLOAD: &[u8; 12] = b"Hello World!";
    const PAYLOAD_OID: &str = "c3b4032160b015b2261530532a6c49f2bdadbe0687fb1f5a6a32e083";

    #[test]
    fn object_read_write() -> Result<()> {
        use std::io::{Cursor, Seek, SeekFrom};

        let mut source = Cursor::new(PAYLOAD);
        let mut fake_file = Cursor::new(vec![]);
        let mut output = vec![];

        let (oid, _) = write_object(&mut source, &mut fake_file,
            OTYPE_BLOB, PAYLOAD.len() as u64)?;

        assert_eq!(oid, ObjectId::from_str(PAYLOAD_OID)?);

        fake_file.seek(SeekFrom::Start(0))?;
        let mut reader = ObjectReader::new(fake_file)?;
        let read_size = reader.read_to_end(&mut output)?;

        assert_eq!(reader.object_type(), *OTYPE_BLOB);
        assert_eq!(reader.size(), PAYLOAD.len() as u64);
        assert_eq!(read_size, PAYLOAD.len());
        assert_eq!(output, PAYLOAD);

        Ok(())
    }

    #[test]
    fn object_write_invalid_size() {
        let mut source = PAYLOAD.as_ref();
        let fake_file = vec![];

        let result = write_object(&mut source, fake_file, OTYPE_BLOB, 13);

        match result {
            Err(Error(ErrorKind::MismatchingObjectSize(actual, expected), _))
                if actual == 12 && expected == 13 => (),
            Err(error) => panic!("Unexpected error: {:?}", error),
            Ok(_) => panic!("Unexpected success"),
        }
    }

    #[test]
    fn write_as_object() -> Result<()> {
        use std::io::{Cursor, sink};

        let mut source = Cursor::new(PAYLOAD);
        let oid = source.write_as_object(sink(), OTYPE_BLOB)?;

        assert_eq!(oid, ObjectId::from_str(PAYLOAD_OID)?);

        Ok(())
    }
}