torrent-rs/src/lib.rs
src/lib.rsRaw
1use async_std::net::{SocketAddr, ToSocketAddrs, UdpSocket, TcpStream};
2use async_std::prelude::*;
3use async_std::io;
4use async_std::task::spawn;
5// use std::io::prelude::*;
6// use std::net::TcpStream;
7use std::error::Error;
8use std::net::Ipv4Addr;
9use std::convert::TryInto;
10use std::str;
11
12use byteorder::{BigEndian, ByteOrder};
13use rand::Rng;
14use sha1::{Sha1, Digest};
15
16extern crate serde;
17extern crate serde_bencode;
18#[macro_use]
19extern crate serde_derive;
20extern crate serde_bytes;
21
22use serde_bencode::{de, ser};
23use serde_bytes::ByteBuf;
24
25// use rustc_hex::ToHex;
26
27const CONNECT_MSG: u32 = 0;
28const ANNOUNCE_MSG: u32 = 1;
29const BLOCK_LEN: u32 = 16384; // 2**14
30
31#[derive(Debug, Deserialize)]
32struct Node(String, i64);
33
34#[derive(Debug, Serialize, Deserialize)]
35struct File {
36 path: Vec<String>,
37 length: i64,
38 #[serde(default)]
39 md5sum: Option<String>,
40}
41
42#[derive(Debug, Serialize, Deserialize)]
43struct Info {
44 name: String,
45 pieces: ByteBuf,
46 #[serde(rename = "piece length")]
47 piece_length: i64,
48 #[serde(default)]
49 md5sum: Option<String>,
50 #[serde(default)]
51 length: Option<i64>,
52 #[serde(default)]
53 files: Option<Vec<File>>,
54 #[serde(default)]
55 private: Option<u8>,
56 #[serde(default)]
57 path: Option<Vec<String>>,
58 #[serde(default)]
59 #[serde(rename = "root hash")]
60 root_hash: Option<String>,
61}
62
63#[derive(Debug, Deserialize)]
64struct Torrent {
65 info: Info,
66 #[serde(default)]
67 announce: Option<String>,
68 #[serde(default)]
69 nodes: Option<Vec<Node>>,
70 #[serde(default)]
71 encoding: Option<String>,
72 #[serde(default)]
73 httpseeds: Option<Vec<String>>,
74 #[serde(default)]
75 #[serde(rename = "announce-list")]
76 announce_list: Option<Vec<Vec<String>>>,
77 #[serde(default)]
78 #[serde(rename = "creation date")]
79 creation_date: Option<i64>,
80 #[serde(rename = "comment")]
81 comment: Option<String>,
82 #[serde(default)]
83 #[serde(rename = "created by")]
84 created_by: Option<String>,
85}
86
87impl Torrent {
88 pub fn open(path: &str) -> Torrent {
89 let b = std::fs::read(path).unwrap();
90 de::from_bytes::<Torrent>(&b).unwrap()
91 }
92 pub fn size(&self) -> u64 {
93 let files = self.info.files.as_ref().unwrap();
94 let mut size: u64 = 0;
95 for i in 0..files.len() {
96 size = size + files[i].length as u64;
97 }
98 size
99 }
100 pub fn info_hash(&self) -> Vec<u8> {
101 let info_bytes: Vec<u8> = ser::to_bytes(&self.info).unwrap();
102 let mut hasher = Sha1::new();
103 hasher.update(info_bytes);
104 let result = hasher.finalize();
105 result.to_vec()
106 }
107 pub fn piece_len(&self, piece_index: u32) -> u32 {
108 let total_len = self.size() as u32;
109 let piece_len = self.info.piece_length as u32;
110
111 let last_piece_len = total_len % piece_len;
112 let last_piece_index = ((total_len / piece_len) as f32).floor() as u32;
113 if last_piece_index == piece_index {
114 return last_piece_len;
115 }
116 piece_len
117 }
118 pub fn blocks_per_piece(&self, piece_index: u32) -> u32 {
119 let piece_len = self.piece_len(piece_index);
120 ((piece_len as u32 / BLOCK_LEN) as f32).floor() as u32
121 }
122 pub fn block_len(&self, piece_index: u32, block_index: u32) -> u32 {
123 let piece_len = self.piece_len(piece_index);
124 let last_piece_len = piece_len % BLOCK_LEN;
125 let last_piece_index = ((piece_len / BLOCK_LEN) as f32).floor() as u32;
126 if block_index==last_piece_index {
127 return last_piece_len;
128 }
129 BLOCK_LEN
130 }
131
132 pub fn announce(&self) -> Option<String> {
133 let mut a: String = "".to_string();
134 match &self.announce {
135 Some(s) => {
136 a = s.to_string();
137 }
138 _ => (),
139 };
140 match &self.announce_list { // TODO refactor approach
141 Some(list) => {
142 a = list[0][0].clone();
143 }
144 _ => (),
145 };
146 Some(a)
147 }
148 pub fn protocol(&self) -> Option<String> {
149 let announce = &self.announce().unwrap();
150 let aux: Vec<&str> = announce.split(":").collect();
151 let protocol = aux[0];
152 Some(protocol.to_string())
153 }
154 pub async fn get_peers(&self) -> Result<Vec<Peer>, Box<dyn Error>> {
155 // TODO timming system to resend request if no answer in X seconds
156 if self.protocol().unwrap() != "udp" {
157 panic!("not udp: {:?}", self.protocol().unwrap()); // TODO remove panic, return error
158 }
159
160 let peer = self.announce.clone().unwrap().replace("udp://", "");
161 println!("peer {:?}", peer);
162 let socket = UdpSocket::bind("0.0.0.0:6681").await?;
163
164 let conn_req = msg_conn_req().to_vec();
165 println!("SENDING conn_req {:?}", conn_req);
166 socket.send_to(&conn_req, &peer).await?;
167
168 let mut buf = vec![0u8; 1024];
169
170 let peers = loop {
171 let (n, src) = socket.recv_from(&mut buf).await?;
172
173 let typ = msg_resp_type(&buf);
174 if typ == CONNECT_MSG {
175 println!("TYPE: CONNECT: {:?}", CONNECT_MSG);
176 // println!("HEX {:?}", &buf[0..n].to_hex());
177 let conn_resp = msg_parse_connect_resp(&buf[0..n].to_vec());
178 let announce_req = msg_ann_req(conn_resp.connection_id, self, 6681);
179 socket.send_to(&announce_req[..], &src).await?;
180 } else if typ==ANNOUNCE_MSG {
181 println!("TYPE: ANNOUNCE: {:?}", ANNOUNCE_MSG);
182 // println!("HEX {:?}", &buf[0..n].to_hex());
183 let ann_resp = msg_parse_announce_resp(&buf[0..n].to_vec());
184 break ann_resp.peers;
185 }
186 };
187
188 Ok(peers)
189 }
190
191}
192
193
194struct Block {
195 length: u32,
196 bytes: Vec<u8>,
197}
198struct Payload {
199 index: u32,
200 begin: u32,
201 length: u32,
202 block: Block,
203}
204struct Msg {
205 size: u32,
206 id: u8,
207 payload: Payload,
208}
209
210fn msg_parse(b: Vec<u8>) -> Result<Msg, &'static str> {
211 if b.len()<4 {
212 return Err("No id in msg.");
213 }
214 let id = b[4];
215 if b.len()<5 {
216 return Err("No payload in msg.");
217 }
218 let mut payload: Payload = Payload{
219 index: 0,
220 begin: 0,
221 length: 0,
222 block: Block{
223 length: 0,
224 bytes: Vec::new()
225 }
226 };
227 match id {
228 6 | 7 | 8 => {
229 payload.index = BigEndian::read_u32(&b[5..9]);
230 payload.begin = BigEndian::read_u32(&b[9..13]);
231 if id==7 {
232 payload.length = BigEndian::read_u32(&b[13..]);
233 }
234 }
235 _ => ()
236 };
237 Ok(Msg{
238 size: BigEndian::read_u32(&b[0..4]),
239 id: id,
240 payload: payload,
241 })
242}
243
244async fn msg_handler<'a>(
245 b: Vec<u8>,
246 mut stream: TcpStream,
247 torrent: &Torrent,
248 pieces: &mut Pieces,
249 queue: &mut Queue<'a>,
250 file: String,
251 ) -> Result<(), Box<dyn Error>> {
252 if is_handshake(&b) {
253 stream.write(&msg_interested()).await?;
254 println!("IS HANDSHAKE");
255 return Ok(())
256 }
257 let m = msg_parse(b).unwrap();
258 match m.id {
259 0 => handler_choke(stream),
260 1 => handler_unchoke(stream, pieces, queue),
261 4 => handler_have(stream, pieces, queue, m.payload),
262 5 => handler_bitfield(stream, pieces, queue, m.payload),
263 7 => handler_piece(stream, pieces, queue, torrent, file, m.payload),
264 _ => {},
265 }
266 Ok(())
267}
268
269fn is_handshake(b: &Vec<u8>) -> bool {
270 b.len()==(b[0] as usize) +49 && b.len()==68 && &b[1..20] == b"BitTorrent protocol"
271}
272
273fn handler_choke(mut stream: TcpStream) {
274 // stream.shutdown()
275}
276fn handler_unchoke(mut stream: TcpStream, pieces: &mut Pieces, queue: &mut Queue) {
277 queue.chocked = false;
278 request_piece(stream, pieces, queue);
279
280}
281fn handler_have(mut stream: TcpStream, pieces: &mut Pieces, queue: &mut Queue, payload: Payload) {
282
283}
284fn handler_bitfield(mut stream: TcpStream, pieces: &mut Pieces, queue: &mut Queue, payload: Payload) {
285
286}
287fn handler_piece(mut stream: TcpStream, pieces: &mut Pieces, queue: &mut Queue, torrent: &Torrent, file: String, payload: Payload) {
288
289}
290async fn request_piece<'a>(mut stream: TcpStream, pieces: &mut Pieces, queue: &mut Queue<'a>) -> Result<(), Box<dyn Error>> {
291 if queue.chocked {
292 // return
293 }
294 while queue.length()>0 {
295 let piece_block = queue.get_last();
296 if pieces.needed(&piece_block) {
297 stream.write(&msg_request(&piece_block)).await?;
298 pieces.add_requested(&piece_block);
299 break;
300 }
301 };
302
303 Ok(())
304}
305
306fn msg_conn_req() -> [u8; 16] {
307 let mut b: [u8; 16] = [0; 16];
308 BigEndian::write_u64(&mut b[0..8], 0x41727101980); // connection_id
309 let random_bytes = rand::thread_rng().gen::<[u8; 4]>();
310 b[12..16].clone_from_slice(&random_bytes[..]);
311 b
312}
313
314fn msg_ann_req(connection_id: u64, torrent: &Torrent, port: u16) -> Vec<u8> {
315 let mut b: [u8; 98] = [0; 98];
316 BigEndian::write_u64(&mut b[0..8], connection_id);
317 BigEndian::write_u32(&mut b[8..12], ANNOUNCE_MSG); // action
318 let random_bytes = rand::thread_rng().gen::<[u8; 4]>();
319 b[12..16].clone_from_slice(&random_bytes[..]);
320 b[16..36].clone_from_slice(&torrent.info_hash());
321 println!("info_hash: {:?}", &b[16..36]);
322 // TODO [36..56] peerId b[36..56].clone_from_slice("todo".as_bytes());
323 // TODO [56..64] downloaded
324 println!("torrent.size(): {:?}", torrent.size());
325 BigEndian::write_u64(&mut b[64..72], torrent.size()); // [64..72] left
326
327 // TODO [72..80] uploaded
328 // TODO [80..84] event (0:none, 1:completed, 2:started, 3:stopped)
329 // TODO [84..88] ip address
330 let random_bytes = rand::thread_rng().gen::<[u8; 4]>();
331 b[88..92].clone_from_slice(&random_bytes[..]); // [88..92] key (random)
332 BigEndian::write_i32(&mut b[92..96], -1); // [92..96] num_want (default: -1)
333 BigEndian::write_u16(&mut b[96..98], port); // [96..98] port
334 b.to_vec()
335}
336
337fn msg_resp_type(b: &Vec<u8>) -> u32 {
338 let action = BigEndian::read_u32(&b[0..4]);
339 if action == 0 {
340 return CONNECT_MSG;
341 } else {
342 return ANNOUNCE_MSG;
343 }
344}
345
346#[derive(Debug)]
347struct ConnResp {
348 action: u32,
349 transaction_id: u32,
350 connection_id: u64,
351}
352fn msg_parse_connect_resp(b: &Vec<u8>) -> ConnResp {
353 ConnResp {
354 action: BigEndian::read_u32(&b[0..4]),
355 transaction_id: BigEndian::read_u32(&b[4..8]),
356 connection_id: BigEndian::read_u64(&b[8..]),
357 }
358}
359
360#[derive(Debug)]
361struct Peer {
362 ip: String,
363 port: u16,
364}
365
366#[derive(Debug)]
367struct AnnResp {
368 action: u32,
369 transaction_id: u32,
370 interval: u32,
371 leechers: u32,
372 seeders: u32,
373 peers: Vec<Peer>,
374}
375fn msg_parse_announce_resp(b: &Vec<u8>) -> AnnResp {
376 let mut peers: Vec<Peer> = Vec::new();
377 let n_peers = (b.len() - 20) / 6;
378 for i in 0..n_peers {
379 let peer: Peer = Peer {
380 // ip: BigEndian::read_u32(&b[20+(6*i)..24+(6*i)]),
381 ip: bytes_to_ip(
382 &b[20 + (6 * i)..24 + (6 * i)]
383 .try_into()
384 .expect("err parsing peer ip"),
385 ),
386 port: BigEndian::read_u16(&b[24 + (6 * i)..26 + (6 * i)]),
387 };
388 peers.push(peer);
389 }
390
391 let ann_resp: AnnResp = AnnResp {
392 action: BigEndian::read_u32(&b[0..4]),
393 transaction_id: BigEndian::read_u32(&b[4..8]),
394 interval: BigEndian::read_u32(&b[8..12]),
395 leechers: BigEndian::read_u32(&b[12..16]),
396 seeders: BigEndian::read_u32(&b[16..20]),
397 peers: peers,
398 };
399
400 ann_resp
401}
402
403fn bytes_to_ip(b: &[u8; 4]) -> String {
404 Ipv4Addr::new(b[0], b[1], b[2], b[3]).to_string()
405}
406
407
408
409#[cfg(test)]
410mod tests {
411 use super::*;
412 use rustc_hex::ToHex;
413
414 #[test]
415 fn torrent_structs() {
416 let t = Torrent::open("test.torrent");
417 assert_eq!(t.info.name, "Big Buck Bunny");
418 assert_eq!(t.announce().unwrap(), "udp://tracker.leechers-paradise.org:6969");
419 assert_eq!(t.info_hash().to_hex(), "dd8255ecdc7ca55fb0bbf81323d87062db1f6d1c");
420 assert_eq!(t.size(), 276_445_467);
421 }
422
423 #[test]
424 fn msgs() {
425 let t = Torrent::open("test.torrent");
426
427 let conn_req = msg_conn_req().to_vec();
428 assert_eq!(conn_req[0..12].to_hex(), "000004172710198000000000"); // [12..16] is random
429 let conn_resp_buf = b"00000000d4c575c2c078f9a83418bebc";
430 let conn_resp = msg_parse_connect_resp(&conn_resp_buf.to_vec());
431 println!("conn_resp {:?}", conn_resp);
432 assert_eq!(conn_resp.action, 808464432);
433 assert_eq!(conn_resp.transaction_id, 808464432);
434 assert_eq!(conn_resp.connection_id, 7220505182792409906);
435
436 let announce_req = msg_ann_req(conn_resp.connection_id, &t, 6681);
437 assert_eq!(announce_req[0..12].to_hex(), "643463353735633200000001");
438 assert_eq!(announce_req[16..72].to_hex(), "dd8255ecdc7ca55fb0bbf81323d87062db1f6d1c0000000000000000000000000000000000000000000000000000000000000000107a391b");
439 assert_eq!(announce_req[92..98].to_hex(), "ffffffff1a19");
440
441 let announce_resp_buf = b"000000017a638a60000006e8000000030000002429d412a11ae12d53dcbd1a19d048c0e51ae1d8241960c8d5d5c36e71c8d5c5b974a31ae2c11ef348d4aabc5f370ec8d5bc4e16321ae1b9cee13b2327b94187b66f7cb94186b1c8d5b75308331ae1b068c0e0c263ac6f857911389d27ac821ae19088c2c01ae172fe2468c00a699aebe6720f68a2f99a1ae167fce2141ae163582ee41ae1634921a5d2725b72aaf5c8d559482e8fc9d257fb394cc8d55740e9ebd166524044e11aff51b6dcb6c8d54e2ebe61c3574b87860fc8d54b6e11026dff473afc7a1ec944cd47cd1ae131249b051ae124ff6b521ae105b75c13c8d50551153a1ae105024e09cd52";
442 let ann_resp = msg_parse_announce_resp(&announce_resp_buf.to_vec());
443 assert_eq!(ann_resp.peers.len(), 81);
444 assert_eq!(ann_resp.peers[0].ip, "48.54.101.56");
445 assert_eq!(ann_resp.peers[0].port, 12336);
446 assert_eq!(ann_resp.peers[80].ip, "52.101.48.57");
447 assert_eq!(ann_resp.peers[80].port, 25444);
448 }
449
450 #[test]
451 fn pieces() {
452 let t = Torrent::open("test.torrent");
453 let mut p = Pieces::new(&t);
454 // println!("{:?}", p);
455
456 let payload: Payload = Payload{
457 index: 0,
458 begin: 0,
459 length: 0,
460 block: Block{
461 length: 0,
462 bytes: Vec::new()
463 }
464 };
465 assert_eq!(true, p.needed(&payload));
466 p.add_requested(&payload);
467 p.add_received(&payload);
468 assert_eq!(false, p.needed(&payload));
469 assert_eq!(false, p.is_done());
470 }
471
472 #[async_std::test]
473 async fn get_peers() {
474 let t = Torrent::open("test.torrent");
475 // println!("{:?}", t);
476 println!("torrent: {:?}", t.info.name);
477 println!("protocol: {:?}", t.protocol().unwrap());
478
479 println!("msg_conn_req: {:?}", msg_conn_req());
480
481 let peers_w = t.get_peers().await;
482 println!("peers getted");
483 let peers = peers_w.unwrap();
484 // println!("get_peers peers: {:?}", peers);
485
486 }
487}