use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use dig_dht::ProviderRecord;
use dig_nat::{AvailabilityItem, AvailabilityResponse, RangeFrame, RangeRequest};
use dig_peer::DigPeer;
use tokio::io::{AsyncRead, AsyncReadExt};
use crate::error::DownloadError;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct RangeMeta {
pub total_length: Option<u64>,
pub chunk_lens: Option<Vec<u64>>,
pub chunk_count: Option<u64>,
pub chunk_index: Option<u64>,
pub root: Option<String>,
pub inclusion_proof: Option<String>,
}
impl RangeMeta {
pub fn from_frame(frame: &RangeFrame) -> Self {
RangeMeta {
total_length: frame.total_length,
chunk_lens: frame.chunk_lens.clone(),
chunk_count: frame.chunk_count,
chunk_index: frame.chunk_index,
root: frame.root.clone(),
inclusion_proof: frame.inclusion_proof.clone(),
}
}
pub fn declaring_layout(
mut self,
total_length: u64,
chunk_lens: Vec<u64>,
chunk_count: u64,
) -> Self {
self.total_length = Some(total_length);
self.chunk_count = Some(chunk_count);
self.chunk_lens = Some(chunk_lens);
self
}
pub fn declaring_anchor(mut self, root: String, inclusion_proof: Option<String>) -> Self {
self.root = Some(root);
self.inclusion_proof = inclusion_proof;
self
}
pub fn declaring_chunk_index(mut self, chunk_index: u64) -> Self {
self.chunk_index = Some(chunk_index);
self
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FetchedRange {
pub request_offset: u64,
pub bytes: Vec<u8>,
pub meta: RangeMeta,
}
#[async_trait]
pub trait RangeTransport: Send + Sync {
async fn query_availability(
&self,
provider: &ProviderRecord,
items: Vec<AvailabilityItem>,
) -> Result<AvailabilityResponse, DownloadError>;
async fn fetch_range(
&self,
provider: &ProviderRecord,
req: &RangeRequest,
) -> Result<FetchedRange, DownloadError>;
}
#[derive(Debug, Clone, Default)]
pub struct SourceHealth {
pub failures: u32,
pub served: u64,
pub backoff_until: Option<Instant>,
}
#[derive(Debug, Default)]
pub struct SourceTracker {
health: HashMap<String, SourceHealth>,
base_backoff: Duration,
max_backoff: Duration,
}
impl SourceTracker {
pub fn new(base_backoff: Duration, max_backoff: Duration) -> Self {
SourceTracker {
health: HashMap::new(),
base_backoff,
max_backoff,
}
}
pub fn is_available(&self, peer_id: &str, now: Instant) -> bool {
match self.health.get(peer_id) {
Some(h) => match h.backoff_until {
Some(until) => now >= until,
None => true,
},
None => true,
}
}
pub fn record_success(&mut self, peer_id: &str) {
let h = self.health.entry(peer_id.to_string()).or_default();
h.failures = 0;
h.served += 1;
h.backoff_until = None;
}
pub fn record_failure(&mut self, peer_id: &str, now: Instant) {
let base = self.base_backoff;
let max = self.max_backoff;
let h = self.health.entry(peer_id.to_string()).or_default();
h.failures = h.failures.saturating_add(1);
let shift = h.failures.saturating_sub(1).min(16);
let backoff = base.checked_mul(1u32 << shift).unwrap_or(max).min(max);
h.backoff_until = Some(now + backoff);
}
pub fn served(&self, peer_id: &str) -> u64 {
self.health.get(peer_id).map(|h| h.served).unwrap_or(0)
}
pub fn failures(&self, peer_id: &str) -> u32 {
self.health.get(peer_id).map(|h| h.failures).unwrap_or(0)
}
}
#[derive(Debug, Default)]
struct StreamIdentity {
root: Option<String>,
total_length: Option<u64>,
chunk_count: Option<u64>,
highest_chunk_index: Option<u64>,
prologue_frontier: u64,
}
#[derive(Debug, PartialEq, Eq)]
enum LaterFrame {
Continuation,
NewProloguePage {
offset: u64,
entries: usize,
},
}
impl StreamIdentity {
fn from_first_frame(frame: &RangeFrame) -> Self {
let page_entries = frame.chunk_lens.as_ref().map_or(0, Vec::len) as u64;
StreamIdentity {
root: frame.root.clone(),
total_length: frame.total_length,
chunk_count: frame.chunk_count,
highest_chunk_index: frame.chunk_index,
prologue_frontier: frame
.chunk_lens_offset
.unwrap_or(0)
.saturating_add(page_entries),
}
}
fn check_later_frame(&mut self, frame: &RangeFrame) -> Result<LaterFrame, String> {
check_unrevised("root", self.root.as_deref(), frame.root.as_deref())?;
check_unrevised("total_length", self.total_length, frame.total_length)?;
check_unrevised("chunk_count", self.chunk_count, frame.chunk_count)?;
if frame.inclusion_proof.is_some() {
return Err("a later frame restates inclusion_proof, which is once-per-stream".into());
}
if let Some(index) = frame.chunk_index {
if let Some(highest) = self.highest_chunk_index {
if index < highest {
return Err(format!(
"chunk_index {index} rewinds below {highest}; frames arrive in ascending offset"
));
}
}
self.highest_chunk_index = Some(index);
}
let Some(page) = &frame.chunk_lens else {
return Ok(LaterFrame::Continuation);
};
let offset = frame.chunk_lens_offset.unwrap_or(0);
if offset < self.prologue_frontier {
return Err(format!(
"a later frame's chunk_lens page at offset {offset} re-covers entries below {}, which an \
earlier page already filled; a paged prologue must ADVANCE, never restate",
self.prologue_frontier
));
}
self.prologue_frontier = offset.saturating_add(page.len() as u64);
Ok(LaterFrame::NewProloguePage {
offset,
entries: page.len(),
})
}
fn entries_delivered(&self) -> u64 {
self.prologue_frontier
}
}
fn check_unrevised<T: PartialEq + std::fmt::Debug>(
field: &str,
committed: Option<T>,
declared: Option<T>,
) -> Result<(), String> {
match (committed, declared) {
(_, None) => Ok(()), (Some(committed), Some(declared)) if committed == declared => Ok(()),
(Some(committed), Some(declared)) => Err(format!(
"{field} changed mid-stream: first frame declared {committed:?}, a later frame {declared:?}"
)),
(None, Some(declared)) => Err(format!(
"{field} {declared:?} appears only on a later frame; the first frame left it unstated"
)),
}
}
pub async fn assemble_range_stream<R: AsyncRead + Unpin>(
reader: &mut R,
max_len: u64,
) -> Result<(Vec<u8>, RangeMeta), DownloadError> {
let mut buf: Vec<u8> = Vec::new();
let mut meta = RangeMeta::default();
let mut identity: Option<StreamIdentity> = None;
let mut byte_frontier: u64 = 0;
loop {
let frame = RangeFrame::decode(reader)
.await
.map_err(|e| DownloadError::Transport {
provider: String::new(),
reason: format!("range frame decode: {e}"),
})?;
let Some(frame) = frame else {
break; };
match identity.as_mut() {
None => {
meta = RangeMeta::from_frame(&frame);
identity = Some(StreamIdentity::from_first_frame(&frame));
}
Some(identity) => {
let verdict = identity.check_later_frame(&frame).map_err(|reason| {
DownloadError::Transport {
provider: String::new(),
reason,
}
})?;
if let LaterFrame::NewProloguePage { .. } = verdict {
return Err(DownloadError::PagedPrologueUnsupported {
provider: String::new(),
chunk_count: identity.chunk_count.unwrap_or_default(),
delivered: identity.entries_delivered(),
});
}
}
}
if max_len == 0 {
break;
}
if frame.offset >= max_len {
return Err(DownloadError::Transport {
provider: String::new(),
reason: format!(
"range frame at offset {} starts beyond expected length {max_len}",
frame.offset
),
});
}
let start = frame.offset as usize;
let take = frame.bytes.len().min((max_len - frame.offset) as usize);
let end = start + take;
if buf.len() < end {
buf.try_reserve(end - buf.len())
.map_err(|e| DownloadError::Transport {
provider: String::new(),
reason: format!("cannot allocate a {end}-byte range assembly buffer: {e}"),
})?;
buf.resize(end, 0); }
buf[start..end].copy_from_slice(&frame.bytes[..take]);
if frame.complete || buf.len() as u64 >= max_len {
break;
}
if end as u64 <= byte_frontier {
return Err(DownloadError::Transport {
provider: String::new(),
reason: format!(
"a frame at offset {} extends the range to {end}, past nothing (already at \
{byte_frontier}), and does not complete it; the stream cannot progress",
frame.offset
),
});
}
byte_frontier = end as u64;
}
Ok((buf, meta))
}
const MAX_TRAILER_DRAIN: u64 = 64 * 1024;
pub async fn drain_trailer_bounded<R: AsyncRead + Unpin>(reader: &mut R, cap: u64) -> u64 {
let mut scratch = [0u8; 4096];
let mut drained: u64 = 0;
while drained < cap {
let want = ((cap - drained) as usize).min(scratch.len());
match reader.read(&mut scratch[..want]).await {
Ok(0) => break, Ok(n) => drained += n as u64,
Err(_) => break, }
}
drained
}
type PooledConn = Arc<tokio::sync::Mutex<DigPeer>>;
pub struct NatRangeTransport {
node: std::sync::Arc<dig_nat::NodeCert>,
config: dig_nat::NatConfig,
network_id: String,
runtime: Arc<dig_nat::NatRuntime>,
pool: tokio::sync::Mutex<HashMap<String, PooledConn>>,
}
impl NatRangeTransport {
pub fn new(
node: std::sync::Arc<dig_nat::NodeCert>,
config: dig_nat::NatConfig,
network_id: impl Into<String>,
) -> Self {
Self::new_with_runtime(
node,
config,
network_id,
Arc::new(dig_nat::NatRuntime::default()),
)
}
pub fn new_with_runtime(
node: std::sync::Arc<dig_nat::NodeCert>,
config: dig_nat::NatConfig,
network_id: impl Into<String>,
runtime: Arc<dig_nat::NatRuntime>,
) -> Self {
NatRangeTransport {
node,
config,
network_id: network_id.into(),
runtime,
pool: tokio::sync::Mutex::new(HashMap::new()),
}
}
pub fn provider_dial_targets(
&self,
provider: &ProviderRecord,
) -> Result<Vec<(String, dig_nat::PeerTarget)>, DownloadError> {
let peer_id = provider.provider_peer_id().ok_or_else(|| {
DownloadError::transport(&provider.provider_peer_id, "malformed provider peer_id")
})?;
let mut targets = Vec::new();
for candidate in crate::addr::dial_candidates(provider) {
match crate::addr::candidate_socket(candidate) {
Ok(socket) => targets.push((
socket.to_string(),
dig_nat::PeerTarget::with_addr(peer_id, socket, self.network_id.clone()),
)),
Err(e) => tracing::warn!(
peer = %crate::error::hex64_or_sentinel(&provider.provider_peer_id, "peer-id"),
candidate = %crate::addr::display(candidate),
error = %e,
"skipping unusable provider candidate address"
),
}
}
targets.push((
"relay-only".to_string(),
dig_nat::PeerTarget::relay_only(peer_id, self.network_id.clone()),
));
Ok(targets)
}
pub fn provider_to_target(
&self,
provider: &ProviderRecord,
) -> Result<dig_nat::PeerTarget, DownloadError> {
let (_, target) = self
.provider_dial_targets(provider)?
.into_iter()
.next()
.expect("dial targets always include the relay-only fallback");
Ok(target)
}
async fn connect(&self, provider: &ProviderRecord) -> Result<DigPeer, DownloadError> {
let mut last_error = None;
for (addr, target) in self.provider_dial_targets(provider)? {
match DigPeer::connect_with_runtime(&target, &self.node, &self.config, &self.runtime)
.await
{
Ok(peer) => return Ok(peer),
Err(e) => {
tracing::debug!(
peer = %crate::error::hex64_or_sentinel(&provider.provider_peer_id, "peer-id"),
candidate = %addr,
error = %e,
"provider dial candidate failed; trying the next address"
);
last_error = Some(format!("dial {addr}: {e}"));
}
}
}
Err(DownloadError::transport(
&provider.provider_peer_id,
last_error.unwrap_or_else(|| "no dialable candidate address".to_string()),
))
}
async fn pooled_conn(&self, provider: &ProviderRecord) -> Result<PooledConn, DownloadError> {
let key = provider.provider_peer_id.clone();
if let Some(conn) = self.pool.lock().await.get(&key).cloned() {
return Ok(conn);
}
let fresh = Arc::new(tokio::sync::Mutex::new(self.connect(provider).await?));
let mut pool = self.pool.lock().await;
Ok(pool.entry(key).or_insert(fresh).clone())
}
async fn evict(&self, provider: &ProviderRecord) {
self.pool.lock().await.remove(&provider.provider_peer_id);
}
}
#[async_trait]
impl RangeTransport for NatRangeTransport {
async fn query_availability(
&self,
provider: &ProviderRecord,
items: Vec<AvailabilityItem>,
) -> Result<AvailabilityResponse, DownloadError> {
let conn = self.pooled_conn(provider).await?;
let res = {
let mut guard = conn.lock().await;
guard.get_availability(items).await
};
match res {
Ok(resp) => Ok(resp),
Err(e) => {
self.evict(provider).await;
Err(DownloadError::transport(&provider.provider_peer_id, e))
}
}
}
async fn fetch_range(
&self,
provider: &ProviderRecord,
req: &RangeRequest,
) -> Result<FetchedRange, DownloadError> {
let conn = self.pooled_conn(provider).await?;
let stream = {
let mut guard = conn.lock().await;
guard.fetch_range(req).await
};
let mut stream = match stream {
Ok(s) => s,
Err(e) => {
self.evict(provider).await;
return Err(DownloadError::transport(&provider.provider_peer_id, e));
}
};
let (bytes, meta) = assemble_range_stream(&mut stream, req.length)
.await
.map_err(|e| e.attributed_to(&provider.provider_peer_id))?;
let _ = drain_trailer_bounded(&mut stream, MAX_TRAILER_DRAIN).await;
Ok(FetchedRange {
request_offset: req.offset,
bytes,
meta,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use dig_dht::{CandidateAddr, ProviderRecord};
use dig_nat::PeerId;
fn test_root() -> String {
"aa".repeat(32)
}
fn encode(frame: &RangeFrame) -> Vec<u8> {
frame
.encode()
.expect("fixture frame must be within the dig-nat framing ceilings")
}
fn provider(peer: u8, host: &str, port: u16) -> ProviderRecord {
ProviderRecord::new(
&dig_dht::Key::from_bytes([0xAB; 32]),
&PeerId::from_bytes([peer; 32]),
vec![CandidateAddr::direct(host, port)],
u64::MAX,
)
}
#[test]
fn provider_to_target_uses_direct_address() {
let t = NatRangeTransport::new(
fake_node_cert(),
dig_nat::NatConfig::default(),
"DIG_MAINNET",
);
let p = provider(1, "203.0.113.7", 9444);
let target = t.provider_to_target(&p).unwrap();
assert_eq!(
target.direct_addr().unwrap().to_string(),
"203.0.113.7:9444"
);
assert_eq!(target.network_id, "DIG_MAINNET");
}
#[test]
fn new_with_runtime_builds_a_full_ladder_transport() {
let runtime = std::sync::Arc::new(dig_nat::NatRuntime::default());
let t = NatRangeTransport::new_with_runtime(
fake_node_cert(),
dig_nat::NatConfig::default(),
"DIG_MAINNET",
runtime,
);
let p = provider(1, "203.0.113.7", 9444);
let target = t.provider_to_target(&p).unwrap();
assert_eq!(
target.direct_addr().unwrap().to_string(),
"203.0.113.7:9444"
);
assert_eq!(target.network_id, "DIG_MAINNET");
}
#[test]
fn provider_to_target_accepts_v4_mapped_v6_host() {
let t = NatRangeTransport::new(
fake_node_cert(),
dig_nat::NatConfig::default(),
"DIG_MAINNET",
);
let p = provider(1, "::ffff:172.31.79.22", 9444);
let target = t
.provider_to_target(&p)
.expect("v4-mapped v6 host must resolve");
assert_eq!(
target.direct_addr().unwrap(),
std::net::SocketAddr::new("::ffff:172.31.79.22".parse().unwrap(), 9444)
);
}
#[test]
fn provider_to_target_accepts_plain_v6_host() {
let t = NatRangeTransport::new(
fake_node_cert(),
dig_nat::NatConfig::default(),
"DIG_MAINNET",
);
let p = provider(1, "2001:db8::1", 9444);
let target = t.provider_to_target(&p).expect("v6 host must resolve");
assert_eq!(
target.direct_addr().unwrap(),
std::net::SocketAddr::new("2001:db8::1".parse().unwrap(), 9444)
);
}
#[test]
fn unusable_first_candidate_falls_through_to_the_ipv4_one() {
let t = NatRangeTransport::new(
fake_node_cert(),
dig_nat::NatConfig::default(),
"DIG_MAINNET",
);
let p = ProviderRecord::new(
&dig_dht::Key::from_bytes([0xAB; 32]),
&PeerId::from_bytes([3; 32]),
vec![
CandidateAddr::direct("not-an-ip-literal", 9444),
CandidateAddr::direct("10.0.0.1", 9444),
],
u64::MAX,
);
let target = t
.provider_to_target(&p)
.expect("the v4 candidate is dialable");
assert_eq!(
target.direct_addr().unwrap(),
"10.0.0.1:9444".parse::<std::net::SocketAddr>().unwrap()
);
}
#[test]
fn dial_targets_order_v6_then_v4_then_relay() {
let t = NatRangeTransport::new(
fake_node_cert(),
dig_nat::NatConfig::default(),
"DIG_MAINNET",
);
let p = ProviderRecord::new(
&dig_dht::Key::from_bytes([0xAB; 32]),
&PeerId::from_bytes([4; 32]),
vec![
CandidateAddr::direct("172.31.79.22", 9444),
CandidateAddr::direct("::ffff:172.31.79.22", 9444),
],
u64::MAX,
);
let addrs: Vec<String> = t
.provider_dial_targets(&p)
.unwrap()
.into_iter()
.map(|(addr, _)| addr)
.collect();
assert_eq!(
addrs,
vec![
"[::ffff:172.31.79.22]:9444",
"172.31.79.22:9444",
"relay-only"
]
);
}
#[tokio::test]
async fn connect_tries_every_candidate_before_failing() {
let t = NatRangeTransport::new(
fake_node_cert(),
dig_nat::NatConfig::default(),
"DIG_MAINNET",
);
let p = ProviderRecord::new(
&dig_dht::Key::from_bytes([0xAB; 32]),
&PeerId::from_bytes([5; 32]),
vec![
CandidateAddr::direct("::1", 1),
CandidateAddr::direct("127.0.0.1", 1),
],
u64::MAX,
);
let err = t.connect(&p).await.expect_err("no listener is up");
let reason = err.to_string();
assert!(
reason.contains("relay-only"),
"the last attempt must be named: {reason}"
);
}
#[test]
fn provider_to_target_relay_only_without_address() {
let t = NatRangeTransport::new(
fake_node_cert(),
dig_nat::NatConfig::default(),
"DIG_MAINNET",
);
let p = ProviderRecord::new(
&dig_dht::Key::from_bytes([0xAB; 32]),
&PeerId::from_bytes([2; 32]),
vec![CandidateAddr::relay_marker()],
u64::MAX,
);
let target = t.provider_to_target(&p).unwrap();
assert!(target.direct_addr().is_none());
}
#[tokio::test]
async fn assemble_reassembles_ordered_frames() {
let f0 = RangeFrame::data(0, b"ABC".to_vec())
.with_identity(test_root(), 6, 2)
.with_chunk_lens_page(0, vec![3, 3])
.with_chunk_index(0)
.with_inclusion_proof("proof");
let f1 = RangeFrame::data(3, b"DEF".to_vec())
.with_complete(true)
.with_identity(test_root(), 6, 2)
.with_chunk_index(1);
let mut wire = encode(&f0);
wire.extend_from_slice(&encode(&f1));
let mut cur = std::io::Cursor::new(wire);
let (bytes, meta) = assemble_range_stream(&mut cur, 6).await.unwrap();
assert_eq!(bytes, b"ABCDEF");
assert_eq!(meta.total_length, Some(6));
assert_eq!(meta.chunk_lens, Some(vec![3, 3]));
assert_eq!(meta.chunk_index, Some(0));
assert_eq!(meta.root, Some("aa".repeat(32)));
assert_eq!(meta.inclusion_proof, Some("proof".into()));
}
#[tokio::test]
async fn assemble_rejects_frame_starting_beyond_window() {
let f = RangeFrame::data(8, vec![0u8; 4]).with_complete(true);
let mut cur = std::io::Cursor::new(encode(&f));
let err = assemble_range_stream(&mut cur, 5).await;
assert!(matches!(err, Err(DownloadError::Transport { .. })));
}
#[tokio::test]
async fn assemble_clips_chunk_granular_frame_to_one_byte_probe() {
let chunk = vec![0x5Au8; 4096];
let f = RangeFrame::data(0, chunk)
.with_complete(true)
.with_identity(test_root(), 1_048_576, 256)
.with_chunk_lens_page(0, vec![4096; 256])
.with_chunk_index(0)
.with_inclusion_proof("proof");
let mut cur = std::io::Cursor::new(encode(&f));
let (bytes, meta) = assemble_range_stream(&mut cur, 1).await.unwrap();
assert_eq!(
bytes,
vec![0x5Au8],
"clipped to exactly the requested window"
);
assert_eq!(meta.total_length, Some(1_048_576));
assert_eq!(meta.chunk_lens, Some(vec![4096; 256]));
assert_eq!(meta.chunk_index, Some(0));
assert_eq!(meta.root, Some("aa".repeat(32)));
assert_eq!(meta.inclusion_proof, Some("proof".into()));
}
#[tokio::test]
async fn assemble_clips_only_the_overshooting_last_frame() {
let f0 = RangeFrame::data(0, b"ABC".to_vec())
.with_identity(test_root(), 9, 2)
.with_chunk_lens_page(0, vec![3, 6])
.with_chunk_index(0);
let f1 = RangeFrame::data(3, b"DEFGHI".to_vec())
.with_complete(true)
.with_identity(test_root(), 9, 2)
.with_chunk_index(1);
let mut wire = encode(&f0);
wire.extend_from_slice(&encode(&f1));
let mut cur = std::io::Cursor::new(wire);
let (bytes, meta) = assemble_range_stream(&mut cur, 5).await.unwrap();
assert_eq!(bytes, b"ABCDE");
assert_eq!(meta.total_length, Some(9));
}
#[tokio::test]
async fn assemble_stops_once_the_window_is_full() {
let f0 = RangeFrame::data(0, b"WXYZ".to_vec())
.with_identity(test_root(), 8, 2)
.with_chunk_lens_page(0, vec![4, 4])
.with_chunk_index(0);
let f1 = RangeFrame::data(4, b"nope".to_vec())
.with_complete(true)
.with_identity(test_root(), 8, 2)
.with_chunk_index(1);
let mut wire = encode(&f0);
wire.extend_from_slice(&encode(&f1));
let mut cur = std::io::Cursor::new(wire);
let (bytes, _) = assemble_range_stream(&mut cur, 4).await.unwrap();
assert_eq!(bytes, b"WXYZ");
}
#[tokio::test]
async fn drain_trailer_is_bounded_by_cap() {
let flood = vec![0u8; 1_000_000];
let mut cur = std::io::Cursor::new(flood);
let drained = drain_trailer_bounded(&mut cur, 64 * 1024).await;
assert_eq!(drained, 64 * 1024, "drain must stop exactly at the cap");
assert!((cur.position() as usize) < 1_000_000);
}
#[tokio::test]
async fn drain_trailer_stops_at_eof_below_cap() {
let mut cur = std::io::Cursor::new(vec![0u8; 100]);
assert_eq!(drain_trailer_bounded(&mut cur, 64 * 1024).await, 100);
let mut empty = std::io::Cursor::new(Vec::<u8>::new());
assert_eq!(drain_trailer_bounded(&mut empty, 64 * 1024).await, 0);
}
#[tokio::test]
async fn assemble_stops_on_clean_eof() {
let f = RangeFrame::data(0, b"hi".to_vec())
.with_identity(test_root(), 2, 1)
.with_chunk_lens_page(0, vec![2])
.with_chunk_index(0);
let mut cur = std::io::Cursor::new(encode(&f));
let (bytes, meta) = assemble_range_stream(&mut cur, 2).await.unwrap();
assert_eq!(bytes, b"hi");
assert_eq!(meta.total_length, Some(2));
}
#[tokio::test]
async fn a_payload_at_the_ceiling_round_trips_and_one_byte_over_is_refused() {
let ceiling = dig_nat::MAX_RANGE_FRAME_PAYLOAD;
let at_ceiling = vec![0x7Eu8; ceiling];
let f = RangeFrame::data(0, at_ceiling.clone())
.with_complete(true)
.with_identity(test_root(), ceiling as u64, 1)
.with_chunk_lens_page(0, vec![ceiling as u64])
.with_chunk_index(0);
let wire = f
.encode()
.expect("a payload AT MAX_RANGE_FRAME_PAYLOAD is conforming and must encode");
let mut cur = std::io::Cursor::new(wire);
let (bytes, meta) = assemble_range_stream(&mut cur, ceiling as u64)
.await
.expect("a ceiling-sized frame decodes and assembles");
assert_eq!(
bytes, at_ceiling,
"every byte of a ceiling-sized window survives the round trip"
);
assert_eq!(meta.total_length, Some(ceiling as u64));
assert_eq!(meta.chunk_index, Some(0));
let over = RangeFrame::data(0, vec![0x7Eu8; ceiling + 1]).with_complete(true);
let err = over
.encode()
.expect_err("one byte past the ceiling has no conforming frame and must be refused");
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
}
#[tokio::test]
async fn a_paged_prologue_is_refused_rather_than_adopted_as_a_complete_layout() {
let chunk_count = dig_nat::MAX_CHUNK_LENS_PER_FRAME + 952;
let chunk_lens: Vec<u64> = (0..chunk_count).map(|i| 64 + (i as u64 % 7)).collect();
let total_length: u64 = chunk_lens.iter().sum();
let page0 = chunk_lens[..dig_nat::MAX_CHUNK_LENS_PER_FRAME].to_vec();
let f = RangeFrame::data(0, b"AB".to_vec())
.with_complete(true)
.with_identity(test_root(), total_length, chunk_count as u64)
.with_chunk_lens_page(0, page0.clone())
.with_chunk_index(0);
let wire = f.encode().expect(
"a first page of MAX_CHUNK_LENS_PER_FRAME entries is within the framing ceiling",
);
let mut cur = std::io::Cursor::new(wire);
let (_bytes, meta) = assemble_range_stream(&mut cur, 2).await.unwrap();
assert_eq!(
meta.chunk_lens.as_deref(),
Some(&page0[..]),
"the assembler surfaces the page it was given, unpadded and unguessed"
);
let err = crate::verify::ResourceCommitment::from_first_frame(
total_length,
meta.chunk_lens.expect("the page is present"),
meta.root,
meta.inclusion_proof,
)
.expect_err("an incomplete chunk_lens must never become a commitment");
assert!(
format!("{err}").contains("chunk_lens sum"),
"and it is refused for the reason that makes it unusable — the array does not describe the declared resource: {err}"
);
}
#[test]
fn source_tracker_backoff_and_recovery() {
let mut t = SourceTracker::new(Duration::from_millis(100), Duration::from_secs(10));
let now = Instant::now();
assert!(t.is_available("p", now));
t.record_failure("p", now);
assert!(!t.is_available("p", now)); assert_eq!(t.failures("p"), 1);
assert!(t.is_available("p", now + Duration::from_millis(101)));
t.record_success("p");
assert!(t.is_available("p", now));
assert_eq!(t.failures("p"), 0);
assert_eq!(t.served("p"), 1);
}
#[test]
fn source_tracker_backoff_is_exponential_and_capped() {
let mut t = SourceTracker::new(Duration::from_millis(100), Duration::from_millis(250));
let now = Instant::now();
t.record_failure("p", now); assert!(t.is_available("p", now + Duration::from_millis(150)));
t.record_failure("p", now); assert!(!t.is_available("p", now + Duration::from_millis(150)));
t.record_failure("p", now); assert!(t.is_available("p", now + Duration::from_millis(260)));
}
fn fake_node_cert() -> std::sync::Arc<dig_nat::NodeCert> {
use sha2::{Digest, Sha256};
let seed: [u8; 32] = Sha256::digest(b"dig-download/tests/fake-node-cert").into();
let bls_sk = dig_tls::bls::SecretKey::from_seed(&seed);
std::sync::Arc::new(dig_nat::NodeCert::generate_signed(&bls_sk).unwrap())
}
#[tokio::test]
async fn an_unsatisfiable_assembly_buffer_is_a_recoverable_error_not_an_abort() {
let f = RangeFrame::data(u64::MAX - 4, vec![0xAB; 2]);
let mut cur = std::io::Cursor::new(encode(&f));
let err = assemble_range_stream(&mut cur, u64::MAX)
.await
.expect_err("an unsatisfiable window allocation is refused, not fatal");
assert!(
err.is_recoverable(),
"and it is RECOVERABLE, so the scheduler re-fetches the range elsewhere: {err}"
);
}
fn identity_pair() -> (RangeFrame, RangeFrame) {
let first = RangeFrame::data(0, b"ABC".to_vec())
.with_identity(test_root(), 6, 2)
.with_chunk_lens_page(0, vec![3, 3])
.with_chunk_index(0)
.with_inclusion_proof("proof");
let second = RangeFrame::data(3, b"DEF".to_vec())
.with_complete(true)
.with_identity(test_root(), 6, 2)
.with_chunk_index(1);
(first, second)
}
async fn reject_reason(first: &RangeFrame, second: &RangeFrame) -> String {
let mut wire = encode(first);
wire.extend_from_slice(&encode(second));
let mut cur = std::io::Cursor::new(wire);
match assemble_range_stream(&mut cur, 6).await {
Err(DownloadError::Transport { reason, .. }) => reason,
other => panic!("a revised identity must be rejected; got {other:?}"),
}
}
#[tokio::test]
async fn a_later_frame_revising_the_root_is_rejected() {
let (first, mut second) = identity_pair();
second.root = Some("bb".repeat(32));
let reason = reject_reason(&first, &second).await;
assert!(
reason.contains("root changed mid-stream"),
"must name the revised field, not fail generically; got {reason}"
);
}
#[tokio::test]
async fn a_later_frame_revising_the_total_length_is_rejected() {
let (first, mut second) = identity_pair();
second.total_length = Some(7);
let reason = reject_reason(&first, &second).await;
assert!(
reason.contains("total_length changed mid-stream"),
"got {reason}"
);
}
#[tokio::test]
async fn a_later_frame_revising_the_chunk_count_is_rejected() {
let (first, mut second) = identity_pair();
second.chunk_count = Some(3);
let reason = reject_reason(&first, &second).await;
assert!(
reason.contains("chunk_count changed mid-stream"),
"got {reason}"
);
}
#[tokio::test]
async fn a_later_frame_revising_the_chunk_count_downward_is_also_rejected() {
let (first, mut second) = identity_pair();
second.chunk_count = Some(1);
let reason = reject_reason(&first, &second).await;
assert!(
reason.contains("chunk_count changed mid-stream"),
"a revision is a revision in EITHER direction; got {reason}"
);
}
#[tokio::test]
async fn an_identity_field_appearing_only_on_a_later_frame_is_rejected() {
let (mut first, second) = identity_pair();
first.chunk_count = None;
let reason = reject_reason(&first, &second).await;
assert!(
reason.contains("appears only on a later frame"),
"got {reason}"
);
}
#[tokio::test]
async fn a_later_frame_restating_an_identical_chunk_lens_page_is_rejected() {
let (first, second) = identity_pair();
let second = second.with_chunk_lens_page(0, vec![3, 3]);
let reason = reject_reason(&first, &second).await;
assert!(
reason.contains("must ADVANCE, never restate"),
"an identical restatement is still a restatement; got {reason}"
);
}
#[tokio::test]
async fn a_later_frame_carrying_an_unstamped_chunk_lens_page_is_rejected() {
let (first, mut second) = identity_pair();
second.chunk_lens = Some(vec![3, 3]);
second.chunk_lens_offset = None;
let reason = reject_reason(&first, &second).await;
assert!(
reason.contains("offset 0 re-covers entries below 2"),
"an unstamped page begins at 0, which is already filled; got {reason}"
);
}
#[tokio::test]
async fn a_later_frame_whose_chunk_lens_page_partially_overlaps_is_rejected() {
let (first, second) = identity_pair();
let second = second.with_chunk_lens_page(1, vec![9]);
let reason = reject_reason(&first, &second).await;
assert!(
reason.contains("offset 1 re-covers entries below 2"),
"a partially overlapping page is a restatement too; got {reason}"
);
}
#[tokio::test]
async fn a_conforming_next_chunk_lens_page_is_recognised_as_paging_not_a_violation() {
let first = RangeFrame::data(0, b"ABC".to_vec())
.with_identity(test_root(), 12, 4)
.with_chunk_lens_page(0, vec![3, 3])
.with_chunk_index(0)
.with_inclusion_proof("proof");
let second = RangeFrame::data(3, b"DEF".to_vec())
.with_complete(true)
.with_identity(test_root(), 12, 4)
.with_chunk_lens_page(2, vec![3, 3])
.with_chunk_index(1);
let mut wire = encode(&first);
wire.extend_from_slice(&encode(&second));
let mut cur = std::io::Cursor::new(wire);
match assemble_range_stream(&mut cur, 6).await {
Err(DownloadError::PagedPrologueUnsupported {
chunk_count,
delivered,
..
}) => {
assert_eq!(chunk_count, 4, "the whole array's declared entry count");
assert_eq!(delivered, 4, "both pages were placed before the reader gave up");
}
other => panic!(
"a conforming next page must be reported as a READER limitation, never as a restatement \
or a transport violation; got {other:?}"
),
}
}
#[tokio::test]
async fn the_paged_prologue_error_survives_provider_attribution_and_is_recoverable() {
let first = RangeFrame::data(0, b"ABC".to_vec())
.with_identity(test_root(), 12, 4)
.with_chunk_lens_page(0, vec![3, 3])
.with_chunk_index(0);
let second = RangeFrame::data(3, b"DEF".to_vec())
.with_complete(true)
.with_identity(test_root(), 12, 4)
.with_chunk_lens_page(2, vec![3, 3]);
let mut wire = encode(&first);
wire.extend_from_slice(&encode(&second));
let mut cur = std::io::Cursor::new(wire);
let raw = assemble_range_stream(&mut cur, 6)
.await
.expect_err("this reader cannot assemble a paged prologue");
let peer = "ab".repeat(32);
let attributed = raw.attributed_to(&peer);
let DownloadError::PagedPrologueUnsupported { provider, .. } = &attributed else {
panic!("attribution must not change the variant; got {attributed:?}");
};
assert_eq!(provider, &peer, "the peer it was read from is stamped in");
assert!(
attributed.is_recoverable(),
"one holder's unassemblable stream must not be terminal for the download"
);
}
#[tokio::test]
async fn a_holder_streaming_empty_non_final_frames_is_refused() {
let first = RangeFrame::data(0, b"AB".to_vec())
.with_identity(test_root(), 64, 2)
.with_chunk_lens_page(0, vec![32, 32])
.with_chunk_index(0);
let empty = RangeFrame::data(0, Vec::new());
let mut wire = encode(&first);
for _ in 0..64 {
wire.extend_from_slice(&encode(&empty));
}
let mut cur = std::io::Cursor::new(wire);
let outcome = tokio::time::timeout(
std::time::Duration::from_secs(5),
assemble_range_stream(&mut cur, 64),
)
.await
.expect("the reader must REFUSE a non-progressing stream, not consume it");
let Err(DownloadError::Transport { reason, .. }) = outcome else {
panic!("a stream that cannot progress must be an error; got {outcome:?}");
};
assert!(
reason.contains("cannot progress"),
"and must say why; got {reason}"
);
}
#[tokio::test]
async fn a_holder_resending_an_already_written_prefix_is_refused() {
let first = RangeFrame::data(0, b"AB".to_vec())
.with_identity(test_root(), 64, 2)
.with_chunk_lens_page(0, vec![32, 32])
.with_chunk_index(0);
let resend = RangeFrame::data(0, b"AB".to_vec());
let mut wire = encode(&first);
for _ in 0..64 {
wire.extend_from_slice(&encode(&resend));
}
let mut cur = std::io::Cursor::new(wire);
let outcome = tokio::time::timeout(
std::time::Duration::from_secs(5),
assemble_range_stream(&mut cur, 64),
)
.await
.expect("a re-sent prefix advances nothing and must be refused, not read forever");
assert!(
matches!(outcome, Err(DownloadError::Transport { .. })),
"got {outcome:?}"
);
}
#[tokio::test]
async fn a_frame_advancing_the_window_by_one_byte_is_accepted() {
let first = RangeFrame::data(0, b"A".to_vec())
.with_identity(test_root(), 3, 1)
.with_chunk_lens_page(0, vec![3])
.with_chunk_index(0);
let second = RangeFrame::data(1, b"B".to_vec());
let third = RangeFrame::data(2, b"C".to_vec()).with_complete(true);
let mut wire = encode(&first);
wire.extend_from_slice(&encode(&second));
wire.extend_from_slice(&encode(&third));
let mut cur = std::io::Cursor::new(wire);
let (bytes, _) = assemble_range_stream(&mut cur, 3)
.await
.expect("one byte at a time is slow, not hostile");
assert_eq!(bytes, b"ABC");
}
#[tokio::test]
async fn wrapping_a_typed_error_loses_the_variant_that_attributing_keeps() {
let peer = "ab".repeat(32);
let typed = || DownloadError::PagedPrologueUnsupported {
provider: String::new(),
chunk_count: 4,
delivered: 4,
};
match typed().attributed_to(&peer) {
DownloadError::PagedPrologueUnsupported { provider, .. } => assert_eq!(
provider, peer,
"attribution must fill the provider in, not merely keep the variant"
),
other => panic!("attribution must leave the variant alone; got {other:?}"),
}
assert!(
matches!(
DownloadError::transport(&peer, typed()),
DownloadError::Transport { .. }
),
"wrapping flattens it to Transport, so `is_recoverable` can no longer tell it apart and the \
stable error catalogue promises a variant no caller can ever match"
);
}
async fn assemble_with_terse_continuation(
strip: impl FnOnce(&mut RangeFrame),
) -> Result<(Vec<u8>, RangeMeta), DownloadError> {
let (first, mut second) = identity_pair();
strip(&mut second);
let mut wire = encode(&first);
wire.extend_from_slice(&encode(&second));
let mut cur = std::io::Cursor::new(wire);
assemble_range_stream(&mut cur, 6).await
}
#[tokio::test]
async fn a_later_frame_omitting_the_root_is_accepted() {
let (bytes, meta) = assemble_with_terse_continuation(|f| f.root = None)
.await
.expect("a terse continuation asserts nothing and must be accepted");
assert_eq!(bytes, b"ABCDEF", "and its bytes still land in the window");
assert_eq!(
meta.root,
Some(test_root()),
"the stream's identity stays the FIRST frame's declaration"
);
}
#[tokio::test]
async fn a_later_frame_omitting_the_total_length_is_accepted() {
let (bytes, meta) = assemble_with_terse_continuation(|f| f.total_length = None)
.await
.expect("a terse continuation asserts nothing and must be accepted");
assert_eq!(bytes, b"ABCDEF");
assert_eq!(meta.total_length, Some(6));
}
#[tokio::test]
async fn a_later_frame_omitting_the_chunk_count_is_accepted() {
let (bytes, meta) = assemble_with_terse_continuation(|f| f.chunk_count = None)
.await
.expect("a terse continuation asserts nothing and must be accepted");
assert_eq!(bytes, b"ABCDEF");
assert_eq!(meta.chunk_count, Some(2));
}
#[tokio::test]
async fn a_later_frame_omitting_every_identity_field_is_accepted() {
let (bytes, _) = assemble_with_terse_continuation(|f| {
f.root = None;
f.total_length = None;
f.chunk_count = None;
f.chunk_index = None;
})
.await
.expect("a bare continuation frame is conforming");
assert_eq!(bytes, b"ABCDEF");
}
#[tokio::test]
async fn a_later_frame_restating_the_inclusion_proof_is_rejected() {
let (first, second) = identity_pair();
let second = second.with_inclusion_proof("proof");
let reason = reject_reason(&first, &second).await;
assert!(reason.contains("restates inclusion_proof"), "got {reason}");
}
#[tokio::test]
async fn a_later_frame_rewinding_the_chunk_index_is_rejected() {
let (mut first, mut second) = identity_pair();
first.chunk_index = Some(1);
second.chunk_index = Some(0);
let reason = reject_reason(&first, &second).await;
assert!(reason.contains("rewinds below"), "got {reason}");
}
#[tokio::test]
async fn a_later_frame_repeating_its_chunk_index_is_accepted() {
let (first, mut second) = identity_pair();
second.chunk_index = Some(0);
let mut wire = encode(&first);
wire.extend_from_slice(&encode(&second));
let mut cur = std::io::Cursor::new(wire);
let (bytes, _) = assemble_range_stream(&mut cur, 6)
.await
.expect("an equal chunk_index is not a rewind");
assert_eq!(bytes, b"ABCDEF");
}
}