torment/torment-tracker/src/lib.rs

use reqwest::blocking::Client;
use reqwest::{Method, Url};
use std::collections::HashMap;
use std::convert::TryInto;
use std::net::{IpAddr, SocketAddr};
use std::ops::Add;
use std::str::FromStr;
use std::time::{Duration, Instant};
use torment_bencode::BencodeValue;
use torment_core::infohash::v1::U160;
use torment_core::infohash::InfoHashCapable;
use torment_core::peer_storage::{PeerStorageHolder, PeerStorageSource};
use torment_core::CompactContact;

pub struct Tracker {
    url: Url,
    is_udp: bool,
    // last_scrape: Option<Instant>,
    last_announce: HashMap<U160, Instant>,
    requested_interval: Duration,
    failure: HashMap<U160, String>,
    peer_storage: PeerStorageHolder,
    // supports_scrape: bool,
}

pub enum TrackerAnnounceRequest {
    HTTP(Url, U160),
    UDP(Url, U160),
}

#[derive(Debug)]
pub struct TrackerAnnounceResponse {
    info_hash: U160,
    // url: Url,
    interval: Option<Duration>,
    source: PeerStorageSource,
    result: Result<Vec<SocketAddr>, String>,
}

impl Tracker {
    pub fn url(&self) -> &Url {
        &self.url
    }

    pub fn new(url: Url, peer_storage: PeerStorageHolder) -> Tracker {
        Tracker {
            is_udp: url.scheme() == "udp",
            url,
            last_announce: HashMap::new(),
            requested_interval: Duration::from_secs(30 * 60),
            failure: HashMap::new(),
            peer_storage,
        }
    }

    pub fn needs_update(&self, info_hash: U160) -> bool {
        self.last_announce
            .get(&info_hash)
            .map(|inst| inst.add(self.requested_interval) < Instant::now())
            .unwrap_or(true)
    }

    pub fn process(&mut self, response: TrackerAnnounceResponse) {
        match response.result {
            Ok(peers) => {
                self.peer_storage
                    .add_peers(response.info_hash, peers, response.source);
            }

            Err(err) => {
                self.failure.insert(response.info_hash, err);
            }
        }

        if let Some(interval) = response.interval {
            self.requested_interval = interval;
        } else {
            self.requested_interval = Duration::from_secs(60 * 30);
        }
    }

    pub fn create_update_request(
        &mut self,
        info_hash: U160,
        peer_id: U160,
        ip: IpAddr,
        port: u16,
        uploaded: usize,
        downloaded: usize,
        left: i64,
    ) -> TrackerAnnounceRequest {
        if self.is_udp {
            return TrackerAnnounceRequest::UDP(self.url.clone(), info_hash);
        }

        let mut url = self.url.clone();
        info_hash.to_byte_array();
        {
            let mut query = url.query_pairs_mut();
            query.append_pair("ip", &format!("{}", ip));
            query.append_pair("port", &format!("{}", port));
            query.append_pair("uploaded", &format!("{}", uploaded));
            query.append_pair("downloaded", &format!("{}", downloaded));
            query.append_pair("left", &format!("{}", left));
        }

        let encoded_info_hash = url::form_urlencoded::byte_serialize(&info_hash.to_bytes()).fold(
            String::new(),
            |mut coll, curr| {
                coll.push_str(curr);
                coll
            },
        );

        let encoded_peer_id = url::form_urlencoded::byte_serialize(&peer_id.to_bytes()).fold(
            String::new(),
            |mut coll, curr| {
                coll.push_str(curr);
                coll
            },
        );
        let query = url.query().unwrap_or("");
        url.set_query(Some(&format!(
            "{}&info_hash={}&peer_id={}",
            query, encoded_info_hash, encoded_peer_id
        )));

        self.last_announce.insert(info_hash, Instant::now());

        TrackerAnnounceRequest::HTTP(url, info_hash)
    }
}

pub fn block_do_announce(
    request: TrackerAnnounceRequest,
    client: &Client,
) -> TrackerAnnounceResponse {
    let (url, info_hash) = match request {
        TrackerAnnounceRequest::HTTP(url, info_hash) => (url, info_hash),
        TrackerAnnounceRequest::UDP(url, info_hash) => {
            return TrackerAnnounceResponse {
                info_hash,
                interval: None,
                source: PeerStorageSource::Tracker(url.clone()),
                // url,
                result: Err(format!("UDP trackers not supported")),
            };
        }
    };

    println!("Announcing {} on {}", info_hash, url);

    let res = block_do_http_response_parse(url.clone(), client);

    TrackerAnnounceResponse {
        info_hash,
        source: PeerStorageSource::Tracker(url.clone()),
        // url,
        interval: res.as_ref().map(|x| x.0).ok(),
        result: res.map(|x| x.1),
    }
}

pub fn block_do_http_response_parse(
    url: Url,
    http_client: &Client,
) -> Result<(Duration, Vec<SocketAddr>), String> {
    let resp = match http_client.request(Method::GET, url).send() {
        Ok(resp) => resp,

        Err(err) => {
            return Err(format!("{}", err));
        }
    };

    let body = resp.bytes();
    if let Err(err) = body {
        return Err(format!("{}", err));
    }

    let body = body.unwrap();
    let body = BencodeValue::decode(body);
    if let Err(err) = body {
        return Err(format!("Bencoding error: {}", err));
    }

    let body = body.unwrap();
    let body = body.dict();
    if body.is_none() {
        return Err(format!("Tracker gave invalid response"));
    }

    let body = body.unwrap();

    if let Some(failure) = body.get(b"failure reason") {
        let failure = failure.string().and_then(|x| x.ok()).unwrap_or(format!(
            "Tracker failed to respond with valid failure, but did respond with a failure"
        ));

        return Err(failure);
    }

    let interval = body
        .get(b"interval")
        .and_then(|item| item.int())
        .unwrap_or(30 * 60);

    let interval = Duration::from_secs(interval as u64);

    let peers = body.get(b"peers");

    if peers.is_none() {
        return Ok((interval, vec![]));
    }

    let peers = peers.unwrap();

    let peers: Vec<SocketAddr> = match peers {
        BencodeValue::List(old) => {
            let mut buffer = vec![];
            for item in old.iter() {
                let dict = if let Some(dict) = item.dict() {
                    dict
                } else {
                    continue;
                };

                let ip = if let Some(BencodeValue::Bytes(ip)) = dict.get(b"ip") {
                    ip.clone()
                } else {
                    continue;
                };

                let ip = if let Some(ip) = IpAddr::from_str(&String::from_utf8_lossy(&ip)).ok() {
                    ip
                } else {
                    continue;
                };

                let port: u16 = if let Some(BencodeValue::Int(nr)) = dict.get(b"port") {
                    if *nr < u16::MAX as i64 {
                        if let Ok(port) = (*nr).try_into() {
                            port
                        } else {
                            continue;
                        }
                    } else {
                        continue;
                    }
                } else {
                    continue;
                };

                buffer.push(SocketAddr::new(ip, port));
            }

            buffer
        }
        BencodeValue::Bytes(compact) => compact
            .chunks(6)
            .filter_map(|x| SocketAddr::from_compact_contact(x).ok())
            .collect(),
        _ => {
            return Err(format!("Tracker failed to respond with valid response"));
        }
    };

    Ok((interval, peers))
}