fp/apps/aquatic/crates/http_protocol/src/response.rs

use std::borrow::Cow;
use std::io::Write;
use std::net::{Ipv4Addr, Ipv6Addr};

use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

use super::common::*;
use super::utils::*;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
pub struct ResponsePeer<I: Eq> {
    pub ip_address: I,
    pub port: u16,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(transparent)]
pub struct ResponsePeerListV4(
    #[serde(
        serialize_with = "serialize_response_peers_ipv4",
        deserialize_with = "deserialize_response_peers_ipv4"
    )]
    pub Vec<ResponsePeer<Ipv4Addr>>,
);

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(transparent)]
pub struct ResponsePeerListV6(
    #[serde(
        serialize_with = "serialize_response_peers_ipv6",
        deserialize_with = "deserialize_response_peers_ipv6"
    )]
    pub Vec<ResponsePeer<Ipv6Addr>>,
);

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScrapeStatistics {
    pub complete: usize,
    pub incomplete: usize,
    pub downloaded: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AnnounceResponse {
    #[serde(rename = "interval")]
    pub announce_interval: usize,
    pub complete: usize,
    pub incomplete: usize,
    #[serde(default)]
    pub peers: ResponsePeerListV4,
    #[serde(default)]
    pub peers6: ResponsePeerListV6,
    // Serialize as string if Some, otherwise skip
    #[serde(
        rename = "warning message",
        skip_serializing_if = "Option::is_none",
        serialize_with = "serialize_optional_string"
    )]
    pub warning_message: Option<String>,
}

impl AnnounceResponse {
    pub fn write_bytes<W: Write>(&self, output: &mut W) -> ::std::io::Result<usize> {
        let mut bytes_written = 0usize;

        bytes_written += output.write(b"d8:completei")?;
        bytes_written += output.write(itoa::Buffer::new().format(self.complete).as_bytes())?;

        bytes_written += output.write(b"e10:incompletei")?;
        bytes_written += output.write(itoa::Buffer::new().format(self.incomplete).as_bytes())?;

        bytes_written += output.write(b"e8:intervali")?;
        bytes_written += output.write(
            itoa::Buffer::new()
                .format(self.announce_interval)
                .as_bytes(),
        )?;

        bytes_written += output.write(b"e5:peers")?;
        bytes_written += output.write(
            itoa::Buffer::new()
                .format(self.peers.0.len() * 6)
                .as_bytes(),
        )?;
        bytes_written += output.write(b":")?;
        for peer in self.peers.0.iter() {
            bytes_written += output.write(&u32::from(peer.ip_address).to_be_bytes())?;
            bytes_written += output.write(&peer.port.to_be_bytes())?;
        }

        bytes_written += output.write(b"6:peers6")?;
        bytes_written += output.write(
            itoa::Buffer::new()
                .format(self.peers6.0.len() * 18)
                .as_bytes(),
        )?;
        bytes_written += output.write(b":")?;
        for peer in self.peers6.0.iter() {
            bytes_written += output.write(&u128::from(peer.ip_address).to_be_bytes())?;
            bytes_written += output.write(&peer.port.to_be_bytes())?;
        }

        if let Some(ref warning_message) = self.warning_message {
            let message_bytes = warning_message.as_bytes();

            bytes_written += output.write(b"15:warning message")?;
            bytes_written +=
                output.write(itoa::Buffer::new().format(message_bytes.len()).as_bytes())?;
            bytes_written += output.write(b":")?;
            bytes_written += output.write(message_bytes)?;
        }

        bytes_written += output.write(b"e")?;

        Ok(bytes_written)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScrapeResponse {
    /// BTreeMap instead of HashMap since keys need to be serialized in order
    pub files: BTreeMap<InfoHash, ScrapeStatistics>,
}

impl ScrapeResponse {
    pub fn write_bytes<W: Write>(&self, output: &mut W) -> ::std::io::Result<usize> {
        let mut bytes_written = 0usize;

        bytes_written += output.write(b"d5:filesd")?;

        for (info_hash, statistics) in self.files.iter() {
            bytes_written += output.write(b"20:")?;
            bytes_written += output.write(&info_hash.0)?;
            bytes_written += output.write(b"d8:completei")?;
            bytes_written +=
                output.write(itoa::Buffer::new().format(statistics.complete).as_bytes())?;
            bytes_written += output.write(b"e10:downloadedi0e10:incompletei")?;
            bytes_written +=
                output.write(itoa::Buffer::new().format(statistics.incomplete).as_bytes())?;
            bytes_written += output.write(b"ee")?;
        }

        bytes_written += output.write(b"ee")?;

        Ok(bytes_written)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FailureResponse {
    #[serde(rename = "failure reason")]
    pub failure_reason: Cow<'static, str>,
}

impl FailureResponse {
    pub fn new<S: Into<Cow<'static, str>>>(reason: S) -> Self {
        Self {
            failure_reason: reason.into(),
        }
    }

    pub fn write_bytes<W: Write>(&self, output: &mut W) -> ::std::io::Result<usize> {
        let mut bytes_written = 0usize;

        let reason_bytes = self.failure_reason.as_bytes();

        bytes_written += output.write(b"d14:failure reason")?;
        bytes_written += output.write(itoa::Buffer::new().format(reason_bytes.len()).as_bytes())?;
        bytes_written += output.write(b":")?;
        bytes_written += output.write(reason_bytes)?;
        bytes_written += output.write(b"e")?;

        Ok(bytes_written)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Response {
    Announce(AnnounceResponse),
    Scrape(ScrapeResponse),
    Failure(FailureResponse),
}

impl Response {
    pub fn write_bytes<W: Write>(&self, output: &mut W) -> ::std::io::Result<usize> {
        match self {
            Response::Announce(r) => r.write_bytes(output),
            Response::Failure(r) => r.write_bytes(output),
            Response::Scrape(r) => r.write_bytes(output),
        }
    }
    pub fn parse_bytes(bytes: &[u8]) -> Result<Self, ::serde_bencode::Error> {
        ::serde_bencode::from_bytes(bytes)
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for ResponsePeer<Ipv4Addr> {
    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
        Self {
            ip_address: Ipv4Addr::arbitrary(g),
            port: u16::arbitrary(g),
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for ResponsePeer<Ipv6Addr> {
    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
        Self {
            ip_address: Ipv6Addr::arbitrary(g),
            port: u16::arbitrary(g),
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for ResponsePeerListV4 {
    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
        Self(Vec::arbitrary(g))
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for ResponsePeerListV6 {
    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
        Self(Vec::arbitrary(g))
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for ScrapeStatistics {
    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
        Self {
            complete: usize::arbitrary(g),
            incomplete: usize::arbitrary(g),
            downloaded: 0,
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for AnnounceResponse {
    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
        Self {
            announce_interval: usize::arbitrary(g),
            complete: usize::arbitrary(g),
            incomplete: usize::arbitrary(g),
            peers: ResponsePeerListV4::arbitrary(g),
            peers6: ResponsePeerListV6::arbitrary(g),
            warning_message: quickcheck::Arbitrary::arbitrary(g),
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for ScrapeResponse {
    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
        Self {
            files: BTreeMap::arbitrary(g),
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for FailureResponse {
    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
        Self {
            failure_reason: String::arbitrary(g).into(),
        }
    }
}

#[cfg(test)]
mod tests {
    use quickcheck_macros::*;

    use super::*;

    #[quickcheck]
    fn test_announce_response_to_bytes(response: AnnounceResponse) -> bool {
        let reference = bendy::serde::to_bytes(&Response::Announce(response.clone())).unwrap();

        let mut hand_written = Vec::new();

        response.write_bytes(&mut hand_written).unwrap();

        let success = hand_written == reference;

        if !success {
            println!("reference:    {}", String::from_utf8_lossy(&reference));
            println!("hand_written: {}", String::from_utf8_lossy(&hand_written));
        }

        success
    }

    #[quickcheck]
    fn test_scrape_response_to_bytes(response: ScrapeResponse) -> bool {
        let reference = bendy::serde::to_bytes(&Response::Scrape(response.clone())).unwrap();

        let mut hand_written = Vec::new();

        response.write_bytes(&mut hand_written).unwrap();

        let success = hand_written == reference;

        if !success {
            println!("reference: {}", String::from_utf8_lossy(&reference));
            println!("hand_written: {}", String::from_utf8_lossy(&hand_written));
        }

        success
    }

    #[quickcheck]
    fn test_failure_response_to_bytes(response: FailureResponse) -> bool {
        let reference = bendy::serde::to_bytes(&Response::Failure(response.clone())).unwrap();

        let mut hand_written = Vec::new();

        response.write_bytes(&mut hand_written).unwrap();

        let success = hand_written == reference;

        if !success {
            println!("reference:    {}", String::from_utf8_lossy(&reference));
            println!("hand_written: {}", String::from_utf8_lossy(&hand_written));
        }

        success
    }
}