aboutsummaryrefslogtreecommitdiff
path: root/src/block/block.rs
blob: 20f57aa532c4e96d6f56a625e11a24fc1aff9dde (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
use std::path::PathBuf;

use bytes::Bytes;
use serde::{Deserialize, Serialize};
use zstd::stream::{decode_all as zstd_decode, Encoder};

use garage_util::data::*;
use garage_util::error::*;

#[derive(Debug, Serialize, Deserialize, Copy, Clone)]
pub enum DataBlockHeader {
	Plain,
	Compressed,
}

/// A possibly compressed block of data
pub enum DataBlock {
	/// Uncompressed data
	Plain(Bytes),
	/// Data compressed with zstd
	Compressed(Bytes),
}

#[derive(Debug)]
pub enum DataBlockPath {
	/// Uncompressed data fail
	Plain(PathBuf),
	/// Compressed data fail
	Compressed(PathBuf),
}

impl DataBlock {
	/// Query whether this block is compressed
	pub fn is_compressed(&self) -> bool {
		matches!(self, DataBlock::Compressed(_))
	}

	/// Get the inner, possibly compressed buffer. You should probably use [`DataBlock::verify_get`]
	/// instead
	pub fn inner_buffer(&self) -> &[u8] {
		use DataBlock::*;
		let (Plain(ref res) | Compressed(ref res)) = self;
		res
	}

	/// Get the buffer, possibly decompressing it, and verify it's integrity.
	/// For Plain block, data is compared to hash, for Compressed block, zstd checksumming system
	/// is used instead.
	pub fn verify_get(self, hash: Hash) -> Result<Bytes, Error> {
		match self {
			DataBlock::Plain(data) => {
				if blake2sum(&data) == hash {
					Ok(data)
				} else {
					Err(Error::CorruptData(hash))
				}
			}
			DataBlock::Compressed(data) => zstd_decode(&data[..])
				.map_err(|_| Error::CorruptData(hash))
				.map(Bytes::from),
		}
	}

	/// Verify data integrity. Allocate less than [`DataBlock::verify_get`] and don't consume self, but
	/// does not return the buffer content.
	pub fn verify(&self, hash: Hash) -> Result<(), Error> {
		match self {
			DataBlock::Plain(data) => {
				if blake2sum(data) == hash {
					Ok(())
				} else {
					Err(Error::CorruptData(hash))
				}
			}
			DataBlock::Compressed(data) => zstd::stream::copy_decode(&data[..], std::io::sink())
				.map_err(|_| Error::CorruptData(hash)),
		}
	}

	pub async fn from_buffer(data: Bytes, level: Option<i32>) -> DataBlock {
		tokio::task::spawn_blocking(move || {
			if let Some(level) = level {
				if let Ok(data) = zstd_encode(&data[..], level) {
					return DataBlock::Compressed(data.into());
				}
			}
			DataBlock::Plain(data)
		})
		.await
		.unwrap()
	}

	pub fn into_parts(self) -> (DataBlockHeader, Bytes) {
		match self {
			DataBlock::Plain(data) => (DataBlockHeader::Plain, data),
			DataBlock::Compressed(data) => (DataBlockHeader::Compressed, data),
		}
	}

	pub fn from_parts(h: DataBlockHeader, bytes: Bytes) -> Self {
		match h {
			DataBlockHeader::Plain => DataBlock::Plain(bytes),
			DataBlockHeader::Compressed => DataBlock::Compressed(bytes),
		}
	}
}

fn zstd_encode<R: std::io::Read>(mut source: R, level: i32) -> std::io::Result<Vec<u8>> {
	let mut result = Vec::<u8>::new();
	let mut encoder = Encoder::new(&mut result, level)?;
	encoder.include_checksum(true)?;
	std::io::copy(&mut source, &mut encoder)?;
	encoder.finish()?;
	Ok(result)
}