use bytes::{Buf, Bytes, BytesMut}; use rand::random; use std::cmp::min; use std::convert::TryInto; use std::io::{Read, Write}; use std::net::{SocketAddr, TcpStream}; use std::process::exit; use torment_core::infohash::v1::U160; use torment_core::infohash::InfoHashCapable; use torment_core::metainfo::Torrent; use torment_core::peer_id; use torment_peer::message::{Handshake, Message, SelectionMessage}; use torment_peer::{Peer, PeerProtocol}; use torment_storage::ToStorageMap; pub fn download(torrent: Torrent, peers: Vec) { println!("Starting download of {}", torrent.name()); if peers.len() == 0 { println!("No peers given, stopping"); exit(1); } let peer = peers.first().unwrap(); let mut stream = TcpStream::connect(peer).expect("Failed to open tcp connection to peer"); stream.set_nodelay(true).expect("Failed setting nodelay"); let id = U160::from(peer_id(random())); let our_header = Handshake::new(id, torrent.info_hash()); stream.write(&our_header.to_bytes()).unwrap(); let mut buffer = [0u8; 4096]; let mut done_header = false; #[allow(unused_assignments)] let mut header = Handshake::new(U160::random(), U160::random()); let mut peer = None; let mut message_buffer = BytesMut::with_capacity(4096 * 10); let mut storage = torrent.to_storage_map("/tmp", false); let mut asked = 0; let request_length = 2u32.pow(14); let piece_length = torrent.meta_info().piece_length() as u32; let mut progress = 0; let mut header_buffer = BytesMut::new(); loop { let size = stream.read(&mut buffer).expect("Failed to receive data"); if done_header == false { header_buffer.extend_from_slice(&buffer[..size]); if header_buffer.len() < 68 { continue; } header = Handshake::from_bytes(Bytes::from(buffer[..size].to_vec())) .expect("Failed to parse header"); if header.info_hash() != torrent.info_hash() { panic!( "Peer wants to serve different torrent ({}) :(", header.info_hash() ); } println!( "Peer[{}] connected", String::from_utf8_lossy(&header.peer_id().to_bytes()) ); peer = Some(Peer::new( stream.peer_addr().unwrap(), PeerProtocol::TCP, header, torrent.meta_info(), )); done_header = true; // stream.write_all(&*Message::Request(SelectionMessage::new(0, 0, 1024)).to_bytes()); continue; } let peer = peer.as_mut().unwrap(); message_buffer.extend_from_slice(&buffer[..size]); while message_buffer.len() >= 4 { let length = u32::from_be_bytes(message_buffer[..4].try_into().unwrap()); if length == 0 { message_buffer.advance(4); continue; } if message_buffer.len() < (length + 4) as usize { break; } let message = message_buffer.split_to((4 + length) as usize).freeze(); let msg = Message::from_bytes(message.slice(4..)).expect("Failed parsing message"); // println!("=> {:?}", msg); if msg == Message::Unchoke { stream .write_all(&Message::Interested.to_length_prefixed_bytes()) .unwrap(); } peer.process(msg); } let mut messages = vec![]; let was_above_90 = peer.count_open_requests() > 90; while !peer.is_choked() && peer.count_open_requests() < 500 && !was_above_90 && asked < torrent.meta_info().pieces() { let curr_piece_length = if asked + 1 == torrent.meta_info().pieces() { piece_length - (storage.size() as u32 % piece_length) } else { piece_length }; let size = min(request_length, curr_piece_length - progress); let sel = SelectionMessage::new(asked as u32, progress, size); peer.requested(sel); let msg = Message::Request(sel); // println!("<= {:?}", msg); messages.extend(msg.to_length_prefixed_bytes()); progress += size; if progress >= curr_piece_length { asked += 1; progress = 0; } } if messages.len() > 0 { stream .write_all(&messages) .expect("Failed writing to network socket"); } while let Some(piece) = peer.next_piece() { // println!("> {} {} {}", piece.index(), piece.offset(), piece.length()); storage .write( piece.index() as usize, piece.offset() as usize, piece.piece(), ) .expect("Failed writing piece to storage"); } if peer.count_open_requests() == 0 && asked > 1 { println!("done"); break; } } }