sn-fetch PoC: range probe, weighted planner, parallel segments, reassembly + sha256
This commit is contained in:
@@ -0,0 +1,123 @@
|
||||
//! A minimal HTTP server for integration tests. Supports HEAD, GET with
|
||||
//! optional range handling (configurable), `Connection: close` semantics.
|
||||
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
use std::io::{BufRead, BufReader, Write};
|
||||
use std::net::{SocketAddr, TcpListener, TcpStream};
|
||||
use std::sync::Arc;
|
||||
use std::thread::JoinHandle;
|
||||
|
||||
pub struct TestServer {
|
||||
pub addr: SocketAddr,
|
||||
_handle: JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl TestServer {
|
||||
/// Start serving `data` on a random port. When `ranges` is false the
|
||||
/// server behaves like a range-hostile origin: no `Accept-Ranges`
|
||||
/// header, and `Range` requests are answered with the full body.
|
||||
pub fn start(data: Vec<u8>, ranges: bool) -> Self {
|
||||
let data = Arc::new(data);
|
||||
let listener = TcpListener::bind("127.0.0.1:0").expect("bind in test");
|
||||
let addr = listener.local_addr().expect("addr in test");
|
||||
let handle = std::thread::spawn(move || {
|
||||
for stream in listener.incoming().flatten() {
|
||||
let data = Arc::clone(&data);
|
||||
std::thread::spawn(move || {
|
||||
let _ = handle_conn(stream, &data, ranges);
|
||||
});
|
||||
}
|
||||
});
|
||||
Self {
|
||||
addr,
|
||||
_handle: handle,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn url(&self, path: &str) -> String {
|
||||
format!("http://{}{path}", self.addr)
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_conn(stream: TcpStream, data: &[u8], ranges: bool) -> std::io::Result<()> {
|
||||
let mut reader = BufReader::new(stream.try_clone()?);
|
||||
let mut request_line = String::new();
|
||||
reader.read_line(&mut request_line)?;
|
||||
|
||||
let mut range_header: Option<String> = None;
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
reader.read_line(&mut line)?;
|
||||
let trimmed = line.trim_end();
|
||||
if trimmed.is_empty() {
|
||||
break;
|
||||
}
|
||||
let lower = trimmed.to_ascii_lowercase();
|
||||
if let Some(value) = lower.strip_prefix("range:") {
|
||||
range_header = Some(value.trim().to_string());
|
||||
}
|
||||
}
|
||||
|
||||
let mut stream = stream;
|
||||
let len = data.len() as u64;
|
||||
|
||||
let method = request_line
|
||||
.split_whitespace()
|
||||
.next()
|
||||
.unwrap_or("GET")
|
||||
.to_string();
|
||||
|
||||
match method.as_str() {
|
||||
"HEAD" => {
|
||||
let accept = if ranges { "bytes" } else { "none" };
|
||||
write!(
|
||||
stream,
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {len}\r\nAccept-Ranges: {accept}\r\nConnection: close\r\n\r\n"
|
||||
)?;
|
||||
}
|
||||
"GET" => {
|
||||
let range = range_header.as_deref().and_then(parse_range).flatten();
|
||||
match (ranges, range) {
|
||||
(true, Some((start, end))) => {
|
||||
let slice_len = end - start + 1;
|
||||
write!(
|
||||
stream,
|
||||
"HTTP/1.1 206 Partial Content\r\nContent-Length: {slice_len}\r\nContent-Range: bytes {start}-{end}/{len}\r\nConnection: close\r\n\r\n"
|
||||
)?;
|
||||
stream.write_all(&data[start as usize..=(end as usize).min(data.len() - 1)])?;
|
||||
}
|
||||
_ => {
|
||||
write!(
|
||||
stream,
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {len}\r\nConnection: close\r\n\r\n"
|
||||
)?;
|
||||
stream.write_all(data)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
write!(
|
||||
stream,
|
||||
"HTTP/1.1 405 Method Not Allowed\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
stream.flush()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Parse `bytes=a-b` into inclusive `(start, end)`. Anything else
|
||||
/// (open-ended, malformed) is treated as absent.
|
||||
fn parse_range(value: &str) -> Option<Option<(u64, u64)>> {
|
||||
let spec = value.strip_prefix("bytes=")?;
|
||||
let (start_str, end_str) = spec.split_once('-')?;
|
||||
let start: u64 = start_str.trim().parse().ok()?;
|
||||
let end: u64 = end_str.trim().parse().ok()?;
|
||||
if end >= start {
|
||||
Some(Some((start, end)))
|
||||
} else {
|
||||
Some(None)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
//! Integration tests: probe, planning, parallel segmented download,
|
||||
//! reassembly and hash verification against a local HTTP server — plus
|
||||
//! the fallback path for range-hostile servers.
|
||||
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
mod common;
|
||||
|
||||
use std::num::{NonZeroU32, NonZeroU64, NonZeroU8, NonZeroUsize};
|
||||
use std::sync::Arc;
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
use tempfile::TempDir;
|
||||
|
||||
use sn_fetch::error::FetchError;
|
||||
use sn_fetch::exec;
|
||||
use sn_fetch::exit::LocalExit;
|
||||
use sn_fetch::plan::{ExitSpecs, FetchPlan};
|
||||
use sn_fetch::probe;
|
||||
use sn_fetch::types::{
|
||||
ByteRange, ExitId, ExitWeight, FileSize, MinSegmentBytes, RangeSupport, SegmentsPerExit,
|
||||
Sha256Digest,
|
||||
};
|
||||
|
||||
fn exit_specs(n: u8) -> ExitSpecs {
|
||||
(1..=n)
|
||||
.map(|i| {
|
||||
(
|
||||
ExitId::new(NonZeroU8::new(i).expect("test ids are 1..=255")),
|
||||
ExitWeight::new(NonZeroU32::new(1).expect("1 is non-zero")),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn probe_detects_range_support() {
|
||||
let data = vec![42u8; 1024];
|
||||
let server = common::TestServer::start(data, true);
|
||||
let url = sn_fetch::types::TargetUrl::parse(&server.url("/file.bin")).expect("valid url");
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let info = probe::probe(&client, &url).await.expect("probe succeeds");
|
||||
assert_eq!(info.range_support, RangeSupport::Supported);
|
||||
assert_eq!(info.content_length.map(|s| s.value()), Some(1024));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn segmented_download_reassembles_and_hashes() {
|
||||
let data: Vec<u8> = (0..(2 * 1024 * 1024)).map(|i| (i % 251) as u8).collect();
|
||||
let expected_digest = Sha256Digest::from_bytes(Sha256::digest(&data).into());
|
||||
|
||||
let server = common::TestServer::start(data.clone(), true);
|
||||
let url = sn_fetch::types::TargetUrl::parse(&server.url("/file.bin")).expect("valid url");
|
||||
|
||||
let specs = exit_specs(2);
|
||||
let plan = FetchPlan::build(
|
||||
url.clone(),
|
||||
FileSize::try_new(data.len() as u64).expect("non-zero in test"),
|
||||
&specs,
|
||||
SegmentsPerExit::new(NonZeroUsize::new(2).expect("2 is non-zero")),
|
||||
MinSegmentBytes::new(NonZeroU64::new(1024).expect("1024 is non-zero")),
|
||||
)
|
||||
.expect("plan builds");
|
||||
plan.validate().expect("plan validates");
|
||||
|
||||
let mut registry = sn_fetch::exit::ExitRegistry::empty();
|
||||
for id in specs.iter().map(|(id, _)| *id) {
|
||||
registry.register(Arc::new(LocalExit::new(id).expect("local exit")));
|
||||
}
|
||||
|
||||
let parts = TempDir::new().expect("tempdir in test");
|
||||
let stats = exec::run_segments(&plan, ®istry, parts.path())
|
||||
.await
|
||||
.expect("segments succeed");
|
||||
assert_eq!(stats.len(), 4);
|
||||
let total_downloaded: u64 = stats.iter().map(|s| s.bytes).sum();
|
||||
assert_eq!(total_downloaded, data.len() as u64);
|
||||
|
||||
let output = parts.path().join("assembled.bin");
|
||||
let digest = exec::assemble(
|
||||
parts.path(),
|
||||
&output,
|
||||
FileSize::try_new(data.len() as u64).expect("non-zero in test"),
|
||||
Some(&expected_digest),
|
||||
)
|
||||
.await
|
||||
.expect("assembly verifies");
|
||||
|
||||
assert_eq!(digest, expected_digest);
|
||||
let written = std::fs::read(&output).expect("output readable in test");
|
||||
assert_eq!(written, data);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn range_hostile_server_falls_back_to_single_stream() {
|
||||
let data: Vec<u8> = (0..256_000u32).map(|i| (i % 197) as u8).collect();
|
||||
let server = common::TestServer::start(data.clone(), false);
|
||||
let url_str = server.url("/file.bin");
|
||||
let target = sn_fetch::types::TargetUrl::parse(&url_str).expect("valid url");
|
||||
|
||||
let specs = exit_specs(1);
|
||||
let options = sn_fetch::runner::FetchOptions {
|
||||
output: std::env::temp_dir().join("sn-fetch-test-fallback.bin"),
|
||||
exits: specs.clone(),
|
||||
segments_per_exit: SegmentsPerExit::new(NonZeroUsize::new(1).expect("1 is non-zero")),
|
||||
min_segment: MinSegmentBytes::new(NonZeroU64::new(1024).expect("1024 is non-zero")),
|
||||
expected_sha256: None,
|
||||
};
|
||||
let registry = sn_fetch::runner::local_registry(&specs).expect("registry");
|
||||
let parts = TempDir::new().expect("tempdir in test");
|
||||
|
||||
let report = sn_fetch::runner::run(&target, &options, ®istry, parts.path())
|
||||
.await
|
||||
.expect("fallback download succeeds");
|
||||
|
||||
assert_eq!(report.file_size.value(), 256_000);
|
||||
let written = std::fs::read(&options.output).expect("output readable in test");
|
||||
assert_eq!(written, data);
|
||||
let _ = std::fs::remove_file(&options.output);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plan_is_contiguous_and_weighted() {
|
||||
let url = sn_fetch::types::TargetUrl::parse("http://example.com/file.bin").expect("valid url");
|
||||
let size = FileSize::try_new(10_000).expect("non-zero in test");
|
||||
let specs: ExitSpecs = vec![
|
||||
(
|
||||
ExitId::new(NonZeroU8::new(1).expect("1 is non-zero")),
|
||||
ExitWeight::new(NonZeroU32::new(3).expect("3 is non-zero")),
|
||||
),
|
||||
(
|
||||
ExitId::new(NonZeroU8::new(2).expect("2 is non-zero")),
|
||||
ExitWeight::new(NonZeroU32::new(1).expect("1 is non-zero")),
|
||||
),
|
||||
];
|
||||
|
||||
let plan = FetchPlan::build(
|
||||
url,
|
||||
size,
|
||||
&specs,
|
||||
SegmentsPerExit::new(NonZeroUsize::new(1).expect("1 is non-zero")),
|
||||
MinSegmentBytes::new(NonZeroU64::new(100).expect("100 is non-zero")),
|
||||
)
|
||||
.expect("plan builds");
|
||||
|
||||
plan.validate().expect("plan validates");
|
||||
|
||||
let exit1_total: u64 = plan
|
||||
.segments
|
||||
.iter()
|
||||
.filter(|seg| seg.exit.value() == 1)
|
||||
.map(|seg| seg.range.len())
|
||||
.sum();
|
||||
let exit2_total: u64 = plan
|
||||
.segments
|
||||
.iter()
|
||||
.filter(|seg| seg.exit.value() == 2)
|
||||
.map(|seg| seg.range.len())
|
||||
.sum();
|
||||
assert_eq!(exit1_total + exit2_total, 10_000);
|
||||
assert!(
|
||||
exit1_total > exit2_total,
|
||||
"heavier exit must get the larger share"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn byte_range_http_value() {
|
||||
let range = ByteRange::try_new(100, 199).expect("valid range");
|
||||
assert_eq!(range.len(), 100);
|
||||
assert_eq!(range.http_value(), "bytes=100-199");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_url_is_rejected() {
|
||||
assert!(matches!(
|
||||
sn_fetch::types::TargetUrl::parse("ftp://example.com/file"),
|
||||
Err(FetchError::InvalidUrl(_))
|
||||
));
|
||||
}
|
||||
Reference in New Issue
Block a user