use anyhow::{Context, Error, Result, anyhow, bail};
use clap::{Args, ValueEnum};
use smol::{
fs::File,
io::{AsyncBufReadExt, BufReader},
};
use std::{collections::BTreeSet, path::PathBuf};
use ufotofu::{
ExpectedFinalError, ProduceAtLeastError,
codec::{Blame, DecodeError},
prelude::*,
producer::compat::{fn_mut::ClosureProducer, reader::reader_to_bulk_producer},
queues::new_unbounded_elastic,
};
use willow25::{
drop_format::{
DropDecoder, DropSliceMetadata, ImportDropError, SliceStreamProducer, import_drop,
},
entry::NamespaceId,
groupings::Keylike,
storage::{PersistentStore, Store},
};
use crate::util::{
BASE32_ALPHABET, SNEAKERWEB_NAMESPACE_ID_BYTES, domain_style, get_domain, hmm, sneakerweb_dir,
warn_deprecated_domain_encoding,
};
#[derive(Args)]
pub struct ImportArgs {
pub src: PathBuf,
pub mode: Option<ImportMode>,
pub collection: Option<PathBuf>,
}
#[derive(Clone, ValueEnum, Default)]
pub enum ImportMode {
#[default]
All,
Familiar,
}
pub async fn import_sneak(args: &ImportArgs) -> Result<()> {
let file = File::open(&args.src).await?;
let drop_producer = reader_to_bulk_producer(file, new_unbounded_elastic());
let sneakerweb_fs_path = sneakerweb_dir(args.collection.as_ref()).await?;
let namespace_id = NamespaceId::from_bytes(&SNEAKERWEB_NAMESPACE_ID_BYTES);
let mut store = PersistentStore::new(&sneakerweb_fs_path).await?;
let allowed = match args.mode.clone().unwrap_or_default() {
ImportMode::All => None,
ImportMode::Familiar => {
let mut allowed = BTreeSet::new();
store
.subspaces(&namespace_id, &mut (&mut allowed).into_consumer())
.await?;
Some(allowed)
}
};
let mut decoder = DropDecoder::new(drop_producer);
let blocklist_path = sneakerweb_fs_path.join("blocked");
let mut blocked = BTreeSet::new();
if blocklist_path.exists() {
let mut blocklist = BufReader::new(
File::open(blocklist_path)
.await
.context("a blocklist was found but could not be read")?,
);
let mut line = String::new();
while let Ok(n) = blocklist.read_line(&mut line).await
&& n > 0
{
let (domain_id, domain_name) = get_domain(Some(&line), "N/A").context(format!(
"invalid domain '{}' found in blocklist",
line.trim()
))?;
warn_deprecated_domain_encoding(&domain_name, &domain_id).await;
blocked.insert(domain_id);
line.clear();
}
}
let mut skipped = BTreeSet::new();
let mut filtering_decoder =
ClosureProducer::<_, (DropSliceMetadata, SliceStreamProducer<_>), (), Error>::new(
async || {
loop {
let Left((metadata, mut slice_stream)) =
decoder.produce().await.map_err(map_decode_err)?
else {
return Ok(Right(()));
};
let mut skip_reason = "blocked";
let domain_id = metadata.entry().subspace_id();
if !blocked.contains(domain_id) {
match &allowed {
None => return Ok(Left((metadata, slice_stream))),
Some(known) => {
if known.contains(domain_id) {
return Ok(Left((metadata, slice_stream)));
}
skip_reason = "unfamiliar";
}
}
}
if !skipped.contains(domain_id) {
hmm(&format!(
"skipping content from {skip_reason} domain {}",
domain_style(&base32::encode(BASE32_ALPHABET, domain_id.as_bytes()))
))
.await;
}
skipped.insert(domain_id.clone());
slice_stream
.skip(
metadata.expected_bytes().try_into().context(
"an entry in the .snk file was too large for this device",
)?,
)
.await
.map_err(|ProduceAtLeastError { count: _, reason }| match reason {
Ok(()) => anyhow!("unexpected end of .snk file"),
Err(err) => map_decode_err(err),
})?;
slice_stream
.produce_final()
.await
.map_err(|err| match err {
ExpectedFinalError::Item(_) => {
anyhow!("an entry in the .snk file had incorrect metadata")
}
ExpectedFinalError::Error(err) => map_decode_err(err),
})?;
}
},
);
import_drop(&mut store, &mut filtering_decoder, true)
.await
.map_err(|err| match err {
ImportDropError::SliceStreamBytesMismatch => {
anyhow!("an entry in the .snk file had incorrect metadata")
}
ImportDropError::ImportProducerError(err) => err,
ImportDropError::SliceProducerError(err) => map_decode_err(err),
ImportDropError::StoreError(err) => err.into(),
ImportDropError::VerificationError => {
anyhow!("an entry in the .snk file could not be verified")
}
ImportDropError::EntryDeleted => {
anyhow!("an entry was deleted concurrently with being imported")
}
ImportDropError::ArchitectureTooSmall => {
anyhow!("an entry in the .snk file was too large for this device")
}
ImportDropError::PartialPayloadSliceEncountered => {
unreachable!("we are not treating incompatible slices as fatal errors")
}
})?;
match store.flush().await {
Ok(_) => Ok(()),
Err(_) => bail!("failed to flush changes to the sneakerweb store"),
}
}
fn map_decode_err<E, P: Into<Error>>(err: DecodeError<E, P, Blame>) -> Error {
match err {
DecodeError::UnexpectedEndOfInput(_) => {
anyhow!("unexpected end of .snk file")
}
DecodeError::ProducerError(err) => anyhow!(err),
DecodeError::Other(blame) => match blame {
Blame::TheirFault => {
anyhow!("malformed encoding of entry in .snk file")
}
Blame::OurFault => anyhow!("failed to decode entry from .snk file"),
},
}
}