// 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 . 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 { writer: GzEncoder, digest: D, size: u64, written_size: u64, } impl ObjectWriter { pub fn new(writer: W, otype: &ObjectType, size: u64) -> Result> { 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 Write for ObjectWriter { fn write(&mut self, buf: &[u8]) -> std::io::Result { 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 { otype: ObjectType, size: u64, reader: GzDecoder, } impl ObjectReader { pub fn new(reader: R) -> Result> { 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 { Ok(self.reader.into_inner()) } } impl Read for ObjectReader { fn read(&mut self, buf: &mut [u8]) -> std::io::Result { self.reader.read(buf) } } pub fn write_object(reader: &mut R, writer: W, otype: &ObjectType, size: u64) -> Result<(ObjectId, W)> where R: Read, W: Write { let mut owriter: ObjectWriter = ObjectWriter::new(writer, otype, size)?; copy(reader, &mut owriter)?; owriter.finish() } pub trait WriteAsObject { fn write_as_object(&mut self, writer: W, otype: &ObjectType) -> Result; fn object_id(&mut self, otype: &ObjectType) -> Result { self.write_as_object(sink(), otype) } } impl WriteAsObject for T { fn write_as_object(&mut self, writer: W, otype: &ObjectType) -> Result { 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(()) } }