arnaucube
arnaucube/torrent-rsPublic Self-hosted repository
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
407fn msg_handshake(info_hash: Vec<u8>) -> Vec<u8> {
408    let mut b: [u8; 68] = [0; 68];
409    b[0] = 19;
410    b[1..20].clone_from_slice(b"BitTorrent protocol"); // TODO use constant for the string
411
412    b[28..48].clone_from_slice(&info_hash);
413
414    b.to_vec()
415}
416
417fn msg_choke() -> Vec<u8> {
418    let mut b: [u8; 5] = [0; 5];
419    BigEndian::write_u32(&mut b[0..4], 1);
420    b.to_vec()
421}
422fn msg_unchoke() -> Vec<u8> {
423    let mut b: [u8; 5] = [0; 5];
424    BigEndian::write_u32(&mut b[0..4], 1);
425    b[4] = 1;
426    b.to_vec()
427}
428fn msg_interested() -> Vec<u8> {
429    let mut b: [u8; 5] = [0; 5];
430    BigEndian::write_u32(&mut b[0..4], 1);
431    b[4] = 2;
432    b.to_vec()
433}
434fn msg_notinterested() -> Vec<u8> {
435    let mut b: [u8; 5] = [0; 5];
436    BigEndian::write_u32(&mut b[0..4], 1);
437    b[4] = 3;
438    b.to_vec()
439}
440fn msg_have(payload: Payload) -> Vec<u8> {
441    let mut b: [u8; 9] = [0; 9];
442    BigEndian::write_u32(&mut b[0..4], 5);
443    b[4] = 4;
444    BigEndian::write_u32(&mut b[5..9], payload.index); // TODO REVIEW CHECK payload.index
445    b.to_vec()
446}
447fn msg_bitfield(bitfield: Vec<u8>) -> Vec<u8> {
448    let mut b: [u8; 14] = [0; 14];
449    BigEndian::write_u32(&mut b[0..4], bitfield.len() as u32);
450    b[4] = 5;
451    b[5..].clone_from_slice(&bitfield);
452    b.to_vec()
453}
454fn msg_request(payload: &Payload) -> Vec<u8> {
455    let mut b: [u8; 17] = [0; 17];
456    BigEndian::write_u32(&mut b[0..4], 13);
457    b[4] = 6;
458    BigEndian::write_u32(&mut b[5..9], payload.index);
459    BigEndian::write_u32(&mut b[9..13], payload.begin);
460    BigEndian::write_u32(&mut b[13..17], payload.length);
461    b.to_vec()
462}
463fn msg_piece(payload: Payload) -> Vec<u8> {
464    let mut b: [u8; 13] = [0; 13]; // TODO payload.block.length
465    BigEndian::write_u32(&mut b[0..4], payload.block.length+9);
466    b[4] = 7;
467    BigEndian::write_u32(&mut b[5..9], payload.index);
468    BigEndian::write_u32(&mut b[9..13], payload.begin);
469    b[13..].clone_from_slice(&payload.block.bytes);
470    b.to_vec()
471}
472fn msg_cancel(payload: Payload) -> Vec<u8> {
473    let mut b: [u8; 17] = [0; 17];
474    BigEndian::write_u32(&mut b[0..4], 13);
475    b[4] = 8;
476    BigEndian::write_u32(&mut b[5..9], payload.index);
477    BigEndian::write_u32(&mut b[9..13], payload.begin);
478    BigEndian::write_u32(&mut b[13..17], payload.length);
479    b.to_vec()
480}
481fn msg_port(port: u16) -> Vec<u8> {
482    let mut b: [u8; 7] = [0; 7];
483    BigEndian::write_u32(&mut b[0..4], 3);
484    b[4] = 9;
485    BigEndian::write_u16(&mut b[5..7], port);
486    b.to_vec()
487}
488
489
490#[cfg(test)]
491mod tests {
492    use super::*;
493    use rustc_hex::ToHex;
494
495    #[test]
496    fn torrent_structs() {
497        let t = Torrent::open("test.torrent");
498        assert_eq!(t.info.name, "Big Buck Bunny");
499        assert_eq!(t.announce().unwrap(), "udp://tracker.leechers-paradise.org:6969");
500        assert_eq!(t.info_hash().to_hex(), "dd8255ecdc7ca55fb0bbf81323d87062db1f6d1c");
501        assert_eq!(t.size(), 276_445_467);
502    }
503
504    #[test]
505    fn msgs() {
506        let t = Torrent::open("test.torrent");
507
508        let conn_req = msg_conn_req().to_vec();
509        assert_eq!(conn_req[0..12].to_hex(), "000004172710198000000000"); // [12..16] is random
510        let conn_resp_buf = b"00000000d4c575c2c078f9a83418bebc";
511        let conn_resp = msg_parse_connect_resp(&conn_resp_buf.to_vec());
512        println!("conn_resp {:?}", conn_resp);
513        assert_eq!(conn_resp.action, 808464432);
514        assert_eq!(conn_resp.transaction_id, 808464432);
515        assert_eq!(conn_resp.connection_id, 7220505182792409906);
516
517        let announce_req = msg_ann_req(conn_resp.connection_id, &t, 6681);
518        assert_eq!(announce_req[0..12].to_hex(), "643463353735633200000001");
519        assert_eq!(announce_req[16..72].to_hex(), "dd8255ecdc7ca55fb0bbf81323d87062db1f6d1c0000000000000000000000000000000000000000000000000000000000000000107a391b");
520        assert_eq!(announce_req[92..98].to_hex(), "ffffffff1a19");
521
522        let announce_resp_buf = b"000000017a638a60000006e8000000030000002429d412a11ae12d53dcbd1a19d048c0e51ae1d8241960c8d5d5c36e71c8d5c5b974a31ae2c11ef348d4aabc5f370ec8d5bc4e16321ae1b9cee13b2327b94187b66f7cb94186b1c8d5b75308331ae1b068c0e0c263ac6f857911389d27ac821ae19088c2c01ae172fe2468c00a699aebe6720f68a2f99a1ae167fce2141ae163582ee41ae1634921a5d2725b72aaf5c8d559482e8fc9d257fb394cc8d55740e9ebd166524044e11aff51b6dcb6c8d54e2ebe61c3574b87860fc8d54b6e11026dff473afc7a1ec944cd47cd1ae131249b051ae124ff6b521ae105b75c13c8d50551153a1ae105024e09cd52";
523        let ann_resp = msg_parse_announce_resp(&announce_resp_buf.to_vec());
524        assert_eq!(ann_resp.peers.len(), 81);
525        assert_eq!(ann_resp.peers[0].ip, "48.54.101.56");
526        assert_eq!(ann_resp.peers[0].port, 12336);
527        assert_eq!(ann_resp.peers[80].ip, "52.101.48.57");
528        assert_eq!(ann_resp.peers[80].port, 25444);
529    }
530
531    #[test]
532    fn pieces() {
533        let t = Torrent::open("test.torrent");
534        let mut p = Pieces::new(&t);
535        // println!("{:?}", p);
536
537        let payload: Payload = Payload{
538            index: 0,
539            begin: 0,
540            length: 0,
541            block: Block{
542                length: 0,
543                bytes: Vec::new()
544            }
545        };
546        assert_eq!(true, p.needed(&payload));
547        p.add_requested(&payload);
548        p.add_received(&payload);
549        assert_eq!(false, p.needed(&payload));
550        assert_eq!(false, p.is_done());
551    }
552    
553    #[async_std::test]
554    async fn get_peers() {
555        let t = Torrent::open("test.torrent");
556        // println!("{:?}", t);
557        println!("torrent: {:?}", t.info.name);
558        println!("protocol: {:?}", t.protocol().unwrap());
559    
560        println!("msg_conn_req: {:?}", msg_conn_req());
561    
562        let peers_w = t.get_peers().await;
563        println!("peers getted");
564        let peers = peers_w.unwrap();
565        // println!("get_peers peers: {:?}", peers);
566
567    }
568}