Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 21 additions & 13 deletions crates/bit_rev/src/file.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,9 @@ use serde_bencode::de;
use serde_bencode::ser;
use serde_bytes::ByteBuf;
use std::fmt::Write;
use std::{error::Error, io::Read};
use std::io::Read;

use anyhow::Result;

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Node(String, i64);
Expand Down Expand Up @@ -95,15 +97,15 @@ impl TorrentMeta {
}
}

pub fn from_filename(filename: &str) -> Result<TorrentMeta, Box<dyn Error>> {
pub fn from_filename(filename: &str) -> Result<TorrentMeta> {
let mut file = std::fs::File::open(filename)?;
let mut content = Vec::new();
file.read_to_end(&mut content)?;
let torrent = de::from_bytes::<TorrentFile>(&content)?;
Ok(TorrentMeta::new(torrent))
}

pub fn url_encode_bytes(content: &[u8]) -> Result<String, Box<dyn Error>> {
pub fn url_encode_bytes(content: &[u8]) -> Result<String> {
let mut out: String = String::new();

for byte in content.iter() {
Expand All @@ -121,21 +123,27 @@ pub fn build_tracker_url(
peer_id: &[u8],
port: u16,
tracker_url: &str,
) -> String {
) -> Result<String> {
// let announce_url = torrent_meta.torrent_file.announce.as_ref().unwrap();
let info_hash_encoded = url_encode_bytes(torrent_meta.info_hash.as_ref()).unwrap();
let peer_id_encoded = url_encode_bytes(peer_id).unwrap();
let info_hash_encoded = url_encode_bytes(torrent_meta.info_hash.as_ref())?;
let peer_id_encoded = url_encode_bytes(peer_id)?;
// let info_hash_encoded = urlencoding::encode_binary(&torrent_meta.info_hash);
// let peer_id_encoded = urlencoding::encode_binary(&peer_id);

format!(
let total_length = if let Some(length) = torrent_meta.torrent_file.info.length {
length
} else if let Some(files) = &torrent_meta.torrent_file.info.files {
files.iter().map(|f| f.length).sum()
} else {
return Err(anyhow::anyhow!(
"Invalid torrent file: missing length information"
));
};

Ok(format!(
// "{}?info_hash={}&peer_id={}&port={}&uploaded=0&downloaded=0&compact=1&left={}&event=started?supportcrypto=1&numwant=80&key=DF45C574",
"{}?info_hash={}&peer_id={}&port={}&uploaded=0&downloaded=0&compact=1&left={}",
tracker_url,
info_hash_encoded,
peer_id_encoded,
port,
torrent_meta.torrent_file.info.length.as_ref().unwrap()
tracker_url, info_hash_encoded, peer_id_encoded, port, total_length
)
.to_string()
.to_string())
}
2 changes: 1 addition & 1 deletion crates/bit_rev/src/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ impl Session {
&self,
add_torrent: AddTorrentOptions,
) -> anyhow::Result<AddTorrentResult> {
let torrent = Torrent::new(&add_torrent.torrent_meta.clone());
let torrent = Torrent::new(&add_torrent.torrent_meta.clone())?;
let torrent_meta = add_torrent.torrent_meta.clone();
let (pr_tx, pr_rx) = flume::bounded::<PieceResult>(torrent.piece_hashes.len());
let have_broadcast = Arc::new(tokio::sync::broadcast::channel(128).0);
Expand Down
55 changes: 50 additions & 5 deletions crates/bit_rev/src/torrent.rs
Original file line number Diff line number Diff line change
@@ -1,20 +1,65 @@
use crate::file::TorrentMeta;

use anyhow::Result;

#[derive(Debug, Clone, PartialEq)]
pub struct Torrent {
pub info_hash: [u8; 20],
pub piece_hashes: Vec<[u8; 20]>,
pub piece_length: i64,
pub length: i64,
pub files: Vec<TorrentFileInfo>,
pub name: String,
}

#[derive(Debug, Clone, PartialEq)]
pub struct TorrentFileInfo {
pub path: Vec<String>,
pub length: i64,
pub offset: i64,
}

impl Torrent {
pub fn new(torrent_meta: &TorrentMeta) -> Torrent {
Torrent {
pub fn new(torrent_meta: &TorrentMeta) -> Result<Torrent> {
let info = &torrent_meta.torrent_file.info;

let (total_length, files) = if let Some(file_list) = &info.files {
// Multi-file torrent
let mut total = 0i64;
let mut torrent_files = Vec::new();

for file in file_list {
torrent_files.push(TorrentFileInfo {
path: file.path.clone(),
length: file.length,
offset: total,
});
total += file.length;
}

(total, torrent_files)
} else if let Some(length) = info.length {
// Single-file torrent
let single_file = TorrentFileInfo {
path: vec![info.name.clone()],
length,
offset: 0,
};

(length, vec![single_file])
} else {
return Err(anyhow::anyhow!(
"Invalid torrent file: missing length information"
));
};

Ok(Torrent {
info_hash: torrent_meta.info_hash,
piece_hashes: torrent_meta.piece_hashes.clone(),
piece_length: torrent_meta.torrent_file.info.piece_length,
length: torrent_meta.torrent_file.info.length.unwrap(),
}
piece_length: info.piece_length,
length: total_length,
files,
name: info.name.clone(),
})
}
}
7 changes: 6 additions & 1 deletion crates/bit_rev/src/tracker_peers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,12 @@ impl TrackerPeers {
let have_broadcast = have_broadcast.clone();
let torrent_downloaded_state = torrent_downloaded_state.clone();
tokio::spawn(async move {
let url = file::build_tracker_url(&torrent_meta, &peer_id, 6881, &tracker);
let url = file::build_tracker_url(&torrent_meta, &peer_id, 6881, &tracker)
.map_err(|e| {
error!("Failed to build tracker URL for {}: {}", tracker, e);
e
})
.unwrap();

match request_peers(&url).await {
Ok(request_peers_res) => {
Expand Down
43 changes: 42 additions & 1 deletion crates/bit_rev/src/utils.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::torrent::Torrent;
use crate::torrent::{Torrent, TorrentFileInfo};
use rand::Rng;

const BLOCK_SIZE: u32 = 16384;
Expand Down Expand Up @@ -42,3 +42,44 @@ pub fn generate_peer_id() -> [u8; 20] {
.try_into()
.unwrap()
}

#[derive(Debug, Clone)]
pub struct PieceFileMapping {
pub file_index: usize,
pub file_offset: usize,
pub length: usize,
}

pub fn map_piece_to_files(torrent: &Torrent, piece_index: usize) -> Vec<PieceFileMapping> {
let (piece_start, piece_end) = calculate_bounds_for_piece(torrent, piece_index);
let mut mappings = Vec::new();

for (file_index, file) in torrent.files.iter().enumerate() {
let file_start = file.offset as usize;
let file_end = file_start + file.length as usize;

// Check if piece overlaps with this file
if piece_start < file_end && piece_end > file_start {
let overlap_start = piece_start.max(file_start);
let overlap_end = piece_end.min(file_end);
let file_offset = overlap_start - file_start;
let length = overlap_end - overlap_start;

mappings.push(PieceFileMapping {
file_index,
file_offset,
length,
});
}
}

mappings
}

pub fn get_full_file_path(torrent: &Torrent, file_info: &TorrentFileInfo) -> std::path::PathBuf {
let mut path = std::path::PathBuf::from(&torrent.name);
for component in &file_info.path {
path.push(component);
}
path
}
83 changes: 66 additions & 17 deletions crates/cli/src/main.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
use indicatif::{ProgressBar, ProgressState, ProgressStyle};
use std::{
collections::HashMap,
fmt::Write,
io::SeekFrom,
sync::{atomic::AtomicU64, Arc},
};
use tokio::{
fs::File,
fs::{create_dir_all, File},
io::{AsyncSeekExt, AsyncWriteExt},
};
use tracing::trace;
Expand Down Expand Up @@ -33,7 +34,7 @@ pub async fn download_file(filename: &str, out_file: Option<String>) -> anyhow::

let add_torrent_result = session.add_torrent(filename.into()).await?;
let torrent = add_torrent_result.torrent.clone();
let torrent_meta = add_torrent_result.torrent_meta;
let _torrent_meta = add_torrent_result.torrent_meta;

let total_size = torrent.length as u64;
let pb = ProgressBar::new(total_size);
Expand All @@ -47,13 +48,41 @@ pub async fn download_file(filename: &str, out_file: Option<String>) -> anyhow::
).progress_chars("#>-")
);

let out_filename = match out_file {
// Determine the output directory
let output_dir = match out_file {
Some(name) => name,
None => torrent_meta.clone().torrent_file.info.name.clone(),
None => torrent.name.clone(),
};
let mut file = File::create(out_filename).await?;

// File
// Create output directory and prepare file handles
let mut file_handles: HashMap<usize, File> = HashMap::new();

// Create directories and prepare files for multi-file torrents
for (file_index, file_info) in torrent.files.iter().enumerate() {
let file_path = if torrent.files.len() == 1 {
// Single file torrent - use output_dir as filename
std::path::PathBuf::from(&output_dir)
} else {
// Multi-file torrent - create subdirectory structure
let mut path = std::path::PathBuf::from(&output_dir);
for component in &file_info.path {
path.push(component);
}
path
};

// Create parent directories if needed
if let Some(parent) = file_path.parent() {
create_dir_all(parent).await?;
}

// Create the file
let file = File::create(&file_path).await?;
file_handles.insert(file_index, file);

trace!("Created file: {:?}", file_path);
}

let total_downloaded = Arc::new(AtomicU64::new(0));
let total_downloaded_clone = total_downloaded.clone();

Expand All @@ -71,21 +100,41 @@ pub async fn download_file(filename: &str, out_file: Option<String>) -> anyhow::
let pr = add_torrent_result.pr_rx.recv_async().await?;

hashset.insert(pr.index);
let (start, end) = utils::calculate_bounds_for_piece(&torrent, pr.index as usize);
trace!(
"index: {}, start: {}, end: {} len {}",
pr.index,
start,
end,
pr.length
);
file.seek(SeekFrom::Start(start as u64)).await?;
file.write_all(pr.buf.as_slice()).await?;

// Map piece to files and write data accordingly
let file_mappings = utils::map_piece_to_files(&torrent, pr.index as usize);
let mut piece_offset = 0;

for mapping in file_mappings {
let file = file_handles.get_mut(&mapping.file_index).ok_or_else(|| {
anyhow::anyhow!("File handle not found for index {}", mapping.file_index)
})?;

// Seek to correct position in file
file.seek(SeekFrom::Start(mapping.file_offset as u64))
.await?;

// Write the portion of the piece that belongs to this file
let piece_data = &pr.buf[piece_offset..piece_offset + mapping.length];
file.write_all(piece_data).await?;

piece_offset += mapping.length;

trace!(
"Wrote {} bytes to file {} at offset {}",
mapping.length,
mapping.file_index,
mapping.file_offset
);
}

total_downloaded.fetch_add(pr.length as u64, std::sync::atomic::Ordering::Relaxed);
}

file.sync_all().await?;
// Sync all files
for (_, file) in file_handles {
file.sync_all().await?;
}

Ok(())
}