use super::{
data::{encrypt_blob, to_chunk, Blob, Spot},
Client,
};
use crate::{
client::{client_api::data::DataMapLevel, utils::encryption, Error, Result},
messaging::data::{DataCmd, DataQuery, QueryResponse},
types::{BytesAddress, Chunk, ChunkAddress, Encryption, PublicKey, Scope},
};
use bincode::deserialize;
use bytes::Bytes;
use futures::future::join_all;
use itertools::Itertools;
use self_encryption::{self, ChunkInfo, DataMap, EncryptedChunk};
use tokio::task;
use tracing::trace;
use xor_name::XorName;
struct HeadChunk {
chunk: Chunk,
address: BytesAddress,
}
impl Client {
#[instrument(skip(self), level = "debug")]
pub async fn read_bytes(&self, address: BytesAddress) -> Result<Bytes> {
let chunk = self.get_chunk(address.name()).await?;
if let Ok(data_map) = self
.unpack_head_chunk(HeadChunk {
chunk: chunk.clone(),
address,
})
.await
{
self.read_all(data_map).await
} else {
self.get_bytes(chunk, address.scope())
}
}
#[instrument(skip_all, level = "trace")]
pub async fn read_from(
&self,
address: BytesAddress,
position: usize,
length: usize,
) -> Result<Bytes>
where
Self: Sized,
{
trace!(
"Reading {:?} bytes at: {:?}, starting from position: {:?}",
&length,
&address,
&position,
);
let chunk = self.get_chunk(address.name()).await?;
if let Ok(data_map) = self
.unpack_head_chunk(HeadChunk {
chunk: chunk.clone(),
address,
})
.await
{
return self.seek(data_map, position, length).await;
}
let bytes = self.get_bytes(chunk, address.scope())?;
let chunk_size = bytes.len();
if chunk_size < position + length {
return Err(Error::InvalidPositionOrLength(format!(
"slicing from chunk at: {:?} of size: {}, with position: {}, and length: {}",
&address, chunk_size, position, length
)));
}
Ok(bytes.slice(position..length))
}
#[instrument(skip(self), level = "trace")]
pub(crate) async fn get_chunk(&self, name: &XorName) -> Result<Chunk> {
let res = self
.send_query(DataQuery::GetChunk(ChunkAddress(*name)))
.await?;
let operation_id = res.operation_id;
let chunk: Chunk = match res.response {
QueryResponse::GetChunk(result) => {
result.map_err(|err| Error::from((err, operation_id)))
}
_ => return Err(Error::ReceivedUnexpectedEvent),
}?;
Ok(chunk)
}
#[instrument(skip_all, level = "trace")]
pub fn chunk_bytes(&self, bytes: Bytes, scope: Scope) -> Result<(BytesAddress, Vec<Chunk>)> {
if let Ok(blob) = Blob::new(bytes.clone()) {
Self::encrypt_blob(blob, scope, self.public_key())
} else {
let spot = Spot::new(bytes)?;
let (address, chunk) = Self::package_spot(spot, scope, self.public_key())?;
Ok((address, vec![chunk]))
}
}
#[instrument(skip(blob), level = "trace")]
fn encrypt_blob(
blob: Blob,
scope: Scope,
public_key: PublicKey,
) -> Result<(BytesAddress, Vec<Chunk>)> {
let owner = encryption(scope, public_key);
encrypt_blob(blob.bytes(), owner.as_ref())
}
fn package_spot(
spot: Spot,
scope: Scope,
public_key: PublicKey,
) -> Result<(BytesAddress, Chunk)> {
let encryption = encryption(scope, public_key);
let chunk = to_chunk(spot.bytes(), encryption.as_ref())?;
if chunk.value().len() >= self_encryption::MIN_ENCRYPTABLE_BYTES {
return Err(Error::Generic("You might need to pad the `Spot` contents and then store it as a `Blob`, as the encryption has made it slightly too big".to_string()));
}
let name = *chunk.name();
let address = if encryption.is_some() {
BytesAddress::Private(name)
} else {
BytesAddress::Public(name)
};
Ok((address, chunk))
}
#[instrument(skip(self, bytes), level = "debug")]
pub async fn upload(&self, bytes: Bytes, scope: Scope) -> Result<BytesAddress> {
if let Ok(blob) = Blob::new(bytes.clone()) {
self.upload_blob(blob, scope).await
} else {
let spot = Spot::new(bytes)?;
self.upload_spot(spot, scope).await
}
}
#[instrument(skip(bytes), level = "debug")]
pub fn calculate_address(bytes: Bytes, scope: Scope) -> Result<BytesAddress> {
let public_key = PublicKey::Bls(bls::SecretKey::random().public_key());
if let Ok(blob) = Blob::new(bytes.clone()) {
let (head_address, _all_chunks) = Self::encrypt_blob(blob, scope, public_key)?;
Ok(head_address)
} else {
let spot = Spot::new(bytes)?;
let (address, _chunk) = Self::package_spot(spot, scope, public_key)?;
Ok(address)
}
}
#[instrument(skip_all, level = "trace")]
async fn upload_blob(&self, blob: Blob, scope: Scope) -> Result<BytesAddress> {
let (head_address, all_chunks) = Self::encrypt_blob(blob, scope, self.public_key())?;
let tasks = all_chunks.into_iter().map(|chunk| {
let writer = self.clone();
task::spawn(async move { writer.send_cmd(DataCmd::StoreChunk(chunk)).await })
});
let _ = join_all(tasks)
.await
.into_iter()
.flatten() .collect_vec();
Ok(head_address)
}
#[instrument(skip_all, level = "trace")]
async fn upload_spot(&self, spot: Spot, scope: Scope) -> Result<BytesAddress> {
let (address, chunk) = Self::package_spot(spot, scope, self.public_key())?;
self.send_cmd(DataCmd::StoreChunk(chunk)).await?;
Ok(address)
}
async fn read_all(&self, data_map: DataMap) -> Result<Bytes> {
let encrypted_chunks = Self::try_get_chunks(self.clone(), data_map.infos()).await?;
self_encryption::decrypt_full_set(&data_map, &encrypted_chunks)
.map_err(Error::SelfEncryption)
}
#[instrument(skip_all, level = "trace")]
async fn seek(&self, data_map: DataMap, pos: usize, len: usize) -> Result<Bytes> {
let info = self_encryption::seek_info(data_map.file_size(), pos, len);
let range = &info.index_range;
let all_infos = data_map.infos();
let encrypted_chunks = Self::try_get_chunks(
self.clone(),
(range.start..range.end + 1)
.clone()
.map(|i| all_infos[i].clone())
.collect_vec(),
)
.await?;
self_encryption::decrypt_range(&data_map, &encrypted_chunks, info.relative_pos, len)
.map_err(Error::SelfEncryption)
}
#[instrument(skip_all, level = "trace")]
async fn try_get_chunks(reader: Client, keys: Vec<ChunkInfo>) -> Result<Vec<EncryptedChunk>> {
let expected_count = keys.len();
let tasks = keys.into_iter().map(|key| {
let reader = reader.clone();
task::spawn(async move {
match reader.get_chunk(&key.dst_hash).await {
Ok(chunk) => Some(EncryptedChunk {
index: key.index,
content: chunk.value().clone(),
}),
Err(e) => {
warn!(
"Reading chunk {} from network, resulted in error {:?}.",
&key.dst_hash, e
);
None
}
}
})
});
let encrypted_chunks = join_all(tasks)
.await
.into_iter()
.flatten()
.flatten()
.collect_vec();
if expected_count > encrypted_chunks.len() {
Err(Error::NotEnoughChunks(
expected_count,
encrypted_chunks.len(),
))
} else {
Ok(encrypted_chunks)
}
}
#[instrument(skip_all, level = "trace")]
async fn unpack_head_chunk(&self, chunk: HeadChunk) -> Result<DataMap> {
let HeadChunk { mut chunk, address } = chunk;
loop {
let bytes = self.get_bytes(chunk, address.scope())?;
match deserialize(&bytes)? {
DataMapLevel::First(data_map) => {
return Ok(data_map);
}
DataMapLevel::Additional(data_map) => {
let serialized_chunk = self.read_all(data_map).await?;
chunk = deserialize(&serialized_chunk)?;
}
}
}
}
#[instrument(skip_all, level = "trace")]
fn get_bytes(&self, chunk: Chunk, scope: Scope) -> Result<Bytes> {
if matches!(scope, Scope::Public) {
Ok(chunk.value().clone())
} else {
let owner = encryption(scope, self.public_key())
.ok_or_else(|| Error::Generic("Could not get an encryption object.".to_string()))?;
Ok(owner.decrypt(chunk.value().clone())?)
}
}
}
#[cfg(test)]
mod tests {
use super::Spot;
use crate::client::utils::test_utils::create_test_client_with;
use crate::client::{
client_api::blob_apis::Blob,
utils::test_utils::{create_test_client, init_test_logger},
Client,
};
use crate::routing::log_markers::LogMarker;
use crate::types::{utils::random_bytes, BytesAddress, Keypair, Scope};
use bytes::Bytes;
use eyre::Result;
use futures::future::join_all;
use rand::rngs::OsRng;
use tokio::time::Instant;
use tracing::Instrument;
const MIN_BLOB_SIZE: usize = self_encryption::MIN_ENCRYPTABLE_BYTES;
const DELAY_DIVIDER: usize = 500_000;
#[test]
fn deterministic_chunking() -> Result<()> {
init_test_logger();
let keypair = Keypair::new_ed25519(&mut OsRng);
let blob = random_bytes(MIN_BLOB_SIZE);
use crate::client::client_api::data::encrypt_blob;
use crate::client::utils::encryption;
let owner = encryption(Scope::Private, keypair.public_key());
let (first_address, mut first_chunks) = encrypt_blob(blob.clone(), owner.as_ref())?;
first_chunks.sort();
for _ in 0..100 {
let owner = encryption(Scope::Private, keypair.public_key());
let (head_address, mut all_chunks) = encrypt_blob(blob.clone(), owner.as_ref())?;
assert_eq!(first_address, head_address);
all_chunks.sort();
assert_eq!(first_chunks, all_chunks);
}
Ok(())
}
#[tokio::test(flavor = "multi_thread")]
async fn store_and_read_3kb() -> Result<()> {
init_test_logger();
let _start_span = tracing::info_span!("store_and_read_3kb").entered();
let client = create_test_client().await?;
let blob = Blob::new(random_bytes(MIN_BLOB_SIZE))?;
let private_address = client.upload_blob(blob.clone(), Scope::Private).await?;
let delay = tokio::time::Duration::from_secs(usize::max(
1,
blob.bytes().len() / DELAY_DIVIDER,
) as u64);
tokio::time::sleep(delay).await;
let read_data = client.read_bytes(private_address).await?;
compare(blob.bytes(), read_data)?;
let address = client
.upload_blob(blob.clone(), Scope::Private)
.instrument(tracing::info_span!(
"checking no conflict on same private upload"
))
.await?;
assert_eq!(address, private_address);
let public_address = client
.upload_blob(blob.clone(), Scope::Public)
.instrument(tracing::info_span!("checking no conflict on public upload"))
.await?;
assert_ne!(public_address, private_address);
let read_data = client
.read_bytes(public_address)
.instrument(tracing::info_span!("reading_public"))
.await?;
compare(blob.bytes(), read_data)?;
Ok(())
}
#[tokio::test(flavor = "multi_thread")]
async fn seek_in_data() -> Result<()> {
init_test_logger();
let _outer_span = tracing::info_span!("seek_in_data").entered();
let client = create_test_client().await?;
for i in 1..5 {
let size = i * MIN_BLOB_SIZE;
let _outer_span = tracing::info_span!("size:", size).entered();
for divisor in 2..5 {
let _outer_span = tracing::info_span!("divisor", divisor).entered();
let len = size / divisor;
let blob = Blob::new(random_bytes(size))?;
let address = store_for_seek(blob.clone(), &client).await?;
let read_data_1 = {
let pos = 0;
get_for_seek(blob.clone(), address, pos, len, &client).await?
};
let read_data_2 = {
let pos = len;
get_for_seek(blob.clone(), address, pos, len, &client).await?
};
let read_data: Bytes = [read_data_1, read_data_2]
.iter()
.flat_map(|bytes| bytes.clone())
.collect();
compare(blob.bytes().slice(0..(2 * len)), read_data)?;
}
}
Ok(())
}
#[tokio::test(flavor = "multi_thread")]
#[ignore = "Testnet network_assert_ tests should be excluded from normal tests runs, they need to be run in sequence to ensure validity of checks"]
async fn blob_network_assert_expected_log_counts() -> Result<()> {
init_test_logger();
let _outer_span = tracing::info_span!("blob_network_assert").entered();
let network_assert_delay: u64 = std::env::var("NETWORK_ASSERT_DELAY")
.unwrap_or_else(|_| "0".to_string())
.parse()?;
let bytes = random_bytes(MIN_BLOB_SIZE / 3);
let client = create_test_client().await?;
let mut the_logs = crate::testnet_grep::NetworkLogState::new()?;
let address = client.upload(bytes.clone(), Scope::Public).await?;
let delay = tokio::time::Duration::from_secs(network_assert_delay);
debug!("Running network asserts with delay of {:?}", delay);
tokio::time::sleep(delay).await;
the_logs
.assert_count(LogMarker::ChunkStoreReceivedAtElder, 3)
.await?;
the_logs
.assert_count(LogMarker::ServiceMsgToBeHandled, 3)
.await?;
the_logs.assert_count(LogMarker::StoringChunk, 12).await?;
the_logs.assert_count(LogMarker::StoredNewChunk, 12).await?;
let _ = client.read_bytes(address).await?;
tokio::time::sleep(delay).await;
the_logs
.assert_count(LogMarker::ChunkQueryReceviedAtElder, 3)
.await?;
the_logs
.assert_count(LogMarker::ChunkQueryReceviedAtAdult, 12)
.await?;
the_logs
.assert_count(LogMarker::ChunkQueryResponseReceviedFromAdult, 12)
.await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread")]
async fn store_and_read_1kb() -> Result<()> {
store_and_read_spot(MIN_BLOB_SIZE / 3, Scope::Public).await?;
store_and_read_spot(MIN_BLOB_SIZE / 3, Scope::Private).await
}
#[tokio::test(flavor = "multi_thread")]
async fn store_and_read_1mb() -> Result<()> {
init_test_logger();
let _outer_span = tracing::info_span!("store_and_read_1mb").entered();
let client = create_test_client().await?;
store_and_read_blob(&client, 1024 * 1024, Scope::Public).await?;
store_and_read_blob(&client, 1024 * 1024, Scope::Private).await
}
#[tokio::test(flavor = "multi_thread")]
async fn ae_checks_blob_test() -> Result<()> {
init_test_logger();
let _outer_span = tracing::info_span!("ae_checks_blob_test").entered();
let client = create_test_client_with(None, None, false).await?;
store_and_read_blob(&client, 10 * 1024 * 1024, Scope::Private).await
}
#[tokio::test(flavor = "multi_thread")]
async fn store_and_read_10mb() -> Result<()> {
init_test_logger();
let _outer_span = tracing::info_span!("store_and_read_10mb").entered();
let client = create_test_client().await?;
store_and_read_blob(&client, 10 * 1024 * 1024, Scope::Private).await
}
#[tokio::test(flavor = "multi_thread")]
#[ignore = "too heavy for CI"]
async fn store_and_read_20mb() -> Result<()> {
init_test_logger();
let _outer_span = tracing::info_span!("store_and_read_20mb").entered();
let client = create_test_client().await?;
store_and_read_blob(&client, 20 * 1024 * 1024, Scope::Private).await
}
#[tokio::test(flavor = "multi_thread")]
#[ignore = "too heavy for CI"]
async fn store_and_read_40mb() -> Result<()> {
init_test_logger();
let _outer_span = tracing::info_span!("store_and_read_40mb").entered();
let client = create_test_client().await?;
store_and_read_blob(&client, 40 * 1024 * 1024, Scope::Private).await
}
#[tokio::test(flavor = "multi_thread")]
#[ignore = "too heavy for CI"]
async fn parallel_timings() -> Result<()> {
init_test_logger();
let _outer_span = tracing::info_span!("parallel_timings").entered();
let client = create_test_client().await?;
let handles = (0..1000_usize)
.map(|i| (i, client.clone()))
.map(|(i, client)| {
tokio::spawn(async move {
let blob = Blob::new(random_bytes(MIN_BLOB_SIZE))?;
let _ = client.upload_blob(blob, Scope::Public).await?;
println!("Iter: {}", i);
let res: Result<()> = Ok(());
res
})
});
let results = join_all(handles).await;
for res1 in results {
if let Ok(res2) = res1 {
if res2.is_err() {
println!("Error: {:?}", res2);
}
} else {
println!("Error: {:?}", res1);
}
}
Ok(())
}
#[tokio::test(flavor = "multi_thread")]
#[ignore = "too heavy for CI"]
async fn one_by_one_timings() -> Result<()> {
init_test_logger();
let _outer_span = tracing::info_span!("test__one_by_one_timings").entered();
let client = create_test_client().await?;
for i in 0..1000_usize {
let blob = Blob::new(random_bytes(MIN_BLOB_SIZE))?;
let now = Instant::now();
let _ = client.upload_blob(blob, Scope::Public).await?;
let elapsed = now.elapsed();
println!("Iter: {}, in {} millis", i, elapsed.as_millis());
}
Ok(())
}
async fn store_and_read_blob(client: &Client, size: usize, scope: Scope) -> Result<()> {
let _outer_span = if scope == Scope::Public {
tracing::info_span!("store_and_read_public_blob", size).entered()
} else {
tracing::info_span!("store_and_read_private_blob", size).entered()
};
let blob_bytes = random_bytes(size);
let blob = Blob::new(blob_bytes.clone())?;
let expected_address = Client::calculate_address(blob_bytes, scope)?;
let address = client.upload_blob(blob.clone(), scope).await?;
assert_eq!(address, expected_address);
let delay = tokio::time::Duration::from_secs(usize::max(1, size / DELAY_DIVIDER) as u64);
tokio::time::sleep(delay).await;
let read_data = client.read_bytes(address).await?;
compare(blob.bytes(), read_data)?;
Ok(())
}
async fn store_and_read_spot(size: usize, scope: Scope) -> Result<()> {
init_test_logger();
let _outer_span = if scope == Scope::Public {
tracing::info_span!("store_and_read_public_spot", size).entered()
} else {
tracing::info_span!("store_and_read_private_spot", size).entered()
};
let spot_bytes = random_bytes(size);
let spot = Spot::new(spot_bytes.clone())?;
let client = create_test_client().await?;
let expected_address = Client::calculate_address(spot_bytes, scope)?;
let address = client.upload_spot(spot.clone(), scope).await?;
assert_eq!(address, expected_address);
let delay = tokio::time::Duration::from_secs(usize::max(1, size / DELAY_DIVIDER) as u64);
tokio::time::sleep(delay).await;
let read_data = client.read_bytes(address).await?;
compare(spot.bytes(), read_data)?;
Ok(())
}
async fn store_for_seek(blob: Blob, client: &Client) -> Result<BytesAddress> {
let address = client.upload_blob(blob.clone(), Scope::Public).await?;
let delay = tokio::time::Duration::from_secs(usize::max(
1,
blob.bytes().len() / DELAY_DIVIDER,
) as u64);
tokio::time::sleep(delay).await;
Ok(address)
}
async fn get_for_seek(
blob: Blob,
address: BytesAddress,
pos: usize,
len: usize,
client: &Client,
) -> Result<Bytes> {
let read_data = client.read_from(address, pos, len).await?;
compare(blob.bytes().slice(pos..(pos + len)), read_data.clone())?;
Ok(read_data)
}
fn compare(original: Bytes, result: Bytes) -> Result<()> {
assert_eq!(original.len(), result.len());
for (counter, (a, b)) in original.into_iter().zip(result).enumerate() {
if a != b {
return Err(eyre::eyre!(format!("Not equal! Counter: {}", counter)));
}
}
Ok(())
}
}