refactor(peer): split download pipeline into modules
The download pipeline had grown into one large file that mixed sentinel transaction handling, peer planning, transport, retry, and top-level orchestration. Split it into a download/ module tree with one file per concern so future lifecycle changes can be reviewed at the right boundary. The public crate surface remains download::download_game_files. Helper types and functions are kept pub(super) or private so the refactor does not widen the API or encourage new callers to depend on internals. The version.ini transaction helpers stay local to version_ini.rs; the proposed fs_util extraction is intentionally left for the later atomic-index work, where a second caller exists. There is no intended runtime behavior change. Test Plan: - just fmt - just test - just clippy - just build Refs: none
This commit is contained in:
@@ -0,0 +1,259 @@
|
||||
use std::{
|
||||
net::SocketAddr,
|
||||
path::{Path, PathBuf},
|
||||
sync::Arc,
|
||||
};
|
||||
|
||||
use lanspread_db::db::GameFileDescription;
|
||||
use tokio::{
|
||||
fs::OpenOptions,
|
||||
io::{AsyncSeekExt, AsyncWriteExt},
|
||||
};
|
||||
use tokio_util::{
|
||||
codec::{FramedWrite, LengthDelimitedCodec},
|
||||
sync::CancellationToken,
|
||||
};
|
||||
|
||||
use super::{
|
||||
planning::{ChunkDownloadResult, DownloadChunk, PeerDownloadPlan},
|
||||
version_ini::VersionIniBuffer,
|
||||
};
|
||||
use crate::{network::connect_to_peer, path_validation::validate_game_file_path};
|
||||
|
||||
fn ensure_download_not_cancelled(
|
||||
cancel_token: &CancellationToken,
|
||||
game_id: &str,
|
||||
) -> eyre::Result<()> {
|
||||
if cancel_token.is_cancelled() {
|
||||
eyre::bail!("download cancelled for game {game_id}");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Downloads a single chunk from a peer.
|
||||
async fn download_chunk(
|
||||
conn: &mut s2n_quic::Connection,
|
||||
base_dir: &Path,
|
||||
game_id: &str,
|
||||
chunk: &DownloadChunk,
|
||||
version_buffer: Option<Arc<VersionIniBuffer>>,
|
||||
) -> eyre::Result<()> {
|
||||
use futures::SinkExt;
|
||||
use lanspread_proto::{Message, Request};
|
||||
|
||||
let stream = conn.open_bidirectional_stream().await?;
|
||||
let (mut rx, tx) = stream.split();
|
||||
let mut framed_tx = FramedWrite::new(tx, LengthDelimitedCodec::new());
|
||||
|
||||
let request = Request::GetGameFileChunk {
|
||||
game_id: game_id.to_string(),
|
||||
relative_path: chunk.relative_path.clone(),
|
||||
offset: chunk.offset,
|
||||
length: chunk.length,
|
||||
};
|
||||
framed_tx.send(request.encode()).await?;
|
||||
|
||||
framed_tx.close().await?;
|
||||
|
||||
if let Some(buffer) = version_buffer
|
||||
&& buffer.matches(&chunk.relative_path)
|
||||
{
|
||||
return download_version_ini_chunk(&mut rx, chunk, &buffer).await;
|
||||
}
|
||||
|
||||
// Validate the path to prevent directory traversal
|
||||
let validated_path = validate_game_file_path(base_dir, &chunk.relative_path)?;
|
||||
let mut file = OpenOptions::new()
|
||||
.create(true)
|
||||
.write(true)
|
||||
.truncate(false)
|
||||
.open(&validated_path)
|
||||
.await?;
|
||||
if chunk.length == 0 && chunk.offset == 0 {
|
||||
// fallback-to-whole-file path replaces any existing partial data
|
||||
file.set_len(0).await?;
|
||||
}
|
||||
file.seek(std::io::SeekFrom::Start(chunk.offset)).await?;
|
||||
|
||||
let mut remaining = chunk.length;
|
||||
let mut received_bytes = 0u64;
|
||||
|
||||
while let Some(bytes) = rx.receive().await? {
|
||||
file.write_all(&bytes).await?;
|
||||
received_bytes += bytes.len() as u64;
|
||||
|
||||
if remaining == 0 {
|
||||
continue;
|
||||
}
|
||||
remaining = remaining.saturating_sub(bytes.len() as u64);
|
||||
if remaining == 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Verify we received the expected amount of data
|
||||
if chunk.length > 0 && received_bytes != chunk.length {
|
||||
eyre::bail!(
|
||||
"Incomplete chunk download: expected {} bytes, received {} bytes for file {} at offset {}",
|
||||
chunk.length,
|
||||
received_bytes,
|
||||
chunk.relative_path,
|
||||
chunk.offset
|
||||
);
|
||||
}
|
||||
|
||||
file.flush().await?;
|
||||
|
||||
// Verify file integrity by checking the file size
|
||||
verify_chunk_integrity(&validated_path, chunk.offset, chunk.length).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn download_version_ini_chunk(
|
||||
rx: &mut s2n_quic::stream::ReceiveStream,
|
||||
chunk: &DownloadChunk,
|
||||
buffer: &VersionIniBuffer,
|
||||
) -> eyre::Result<()> {
|
||||
let mut received = Vec::new();
|
||||
while let Some(bytes) = rx.receive().await? {
|
||||
received.extend_from_slice(&bytes);
|
||||
}
|
||||
|
||||
if chunk.length > 0 && u64::try_from(received.len())? != chunk.length {
|
||||
eyre::bail!(
|
||||
"Incomplete version.ini chunk download: expected {} bytes, received {} bytes at offset {}",
|
||||
chunk.length,
|
||||
received.len(),
|
||||
chunk.offset
|
||||
);
|
||||
}
|
||||
|
||||
buffer.write_at(chunk.offset, &received).await
|
||||
}
|
||||
|
||||
/// Verifies that a chunk was written correctly.
|
||||
async fn verify_chunk_integrity(
|
||||
file_path: &Path,
|
||||
offset: u64,
|
||||
expected_length: u64,
|
||||
) -> eyre::Result<()> {
|
||||
if expected_length == 0 {
|
||||
return Ok(()); // Skip verification for whole files or zero-length chunks
|
||||
}
|
||||
|
||||
let metadata = tokio::fs::metadata(file_path).await?;
|
||||
let file_size = metadata.len();
|
||||
|
||||
if file_size < offset + expected_length {
|
||||
eyre::bail!(
|
||||
"File integrity check failed: file size {} is less than expected {} (offset: {})",
|
||||
file_size,
|
||||
offset + expected_length,
|
||||
offset
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Downloads a whole file from a peer.
|
||||
async fn download_whole_file(
|
||||
conn: &mut s2n_quic::Connection,
|
||||
base_dir: &Path,
|
||||
desc: &GameFileDescription,
|
||||
) -> eyre::Result<()> {
|
||||
use futures::SinkExt;
|
||||
use lanspread_proto::{Message, Request};
|
||||
|
||||
let stream = conn.open_bidirectional_stream().await?;
|
||||
let (mut rx, tx) = stream.split();
|
||||
let mut framed_tx = FramedWrite::new(tx, LengthDelimitedCodec::new());
|
||||
|
||||
framed_tx
|
||||
.send(Request::GetGameFileData(desc.clone()).encode())
|
||||
.await?;
|
||||
framed_tx.close().await?;
|
||||
|
||||
// Validate the path to prevent directory traversal
|
||||
let validated_path = validate_game_file_path(base_dir, &desc.relative_path)?;
|
||||
let mut file = OpenOptions::new()
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.write(true)
|
||||
.open(&validated_path)
|
||||
.await?;
|
||||
file.seek(std::io::SeekFrom::Start(0)).await?;
|
||||
|
||||
while let Some(bytes) = rx.receive().await? {
|
||||
file.write_all(&bytes).await?;
|
||||
}
|
||||
|
||||
file.flush().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Downloads all assigned chunks and files from a single peer.
|
||||
pub(super) async fn download_from_peer(
|
||||
peer_addr: SocketAddr,
|
||||
game_id: &str,
|
||||
plan: PeerDownloadPlan,
|
||||
games_folder: PathBuf,
|
||||
cancel_token: &CancellationToken,
|
||||
version_buffer: Option<Arc<VersionIniBuffer>>,
|
||||
) -> eyre::Result<Vec<ChunkDownloadResult>> {
|
||||
if plan.chunks.is_empty() && plan.whole_files.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
ensure_download_not_cancelled(cancel_token, game_id)?;
|
||||
|
||||
let mut conn = connect_to_peer(peer_addr).await?;
|
||||
conn.keep_alive(true)?;
|
||||
conn.keep_alive(true)?;
|
||||
|
||||
let base_dir = games_folder;
|
||||
let mut results = Vec::new();
|
||||
|
||||
// Download chunks with error handling
|
||||
for chunk in &plan.chunks {
|
||||
ensure_download_not_cancelled(cancel_token, game_id)?;
|
||||
|
||||
log::info!(
|
||||
"Downloading chunk {} (offset {}, length {}) from {}",
|
||||
chunk.relative_path,
|
||||
chunk.offset,
|
||||
chunk.length,
|
||||
peer_addr
|
||||
);
|
||||
let result =
|
||||
download_chunk(&mut conn, &base_dir, game_id, chunk, version_buffer.clone()).await;
|
||||
results.push(ChunkDownloadResult {
|
||||
chunk: chunk.clone(),
|
||||
result,
|
||||
peer_addr,
|
||||
});
|
||||
}
|
||||
|
||||
// Download whole files
|
||||
for desc in &plan.whole_files {
|
||||
ensure_download_not_cancelled(cancel_token, game_id)?;
|
||||
|
||||
let chunk = DownloadChunk {
|
||||
relative_path: desc.relative_path.clone(),
|
||||
offset: 0,
|
||||
length: 0, // Indicates whole file
|
||||
retry_count: 0,
|
||||
last_peer: Some(peer_addr),
|
||||
};
|
||||
|
||||
let result = download_whole_file(&mut conn, &base_dir, desc).await;
|
||||
results.push(ChunkDownloadResult {
|
||||
chunk,
|
||||
result,
|
||||
peer_addr,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
Reference in New Issue
Block a user