torrust-tracker/src/response.rs at main · cornservant/torrust-tracker · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
use super::common::*;
use crate::TorrentPeer;
use byteorder::{NetworkEndian, WriteBytesExt};
use std;
use std::io;
use std::io::Write;
use std::net::SocketAddr;

#[derive(PartialEq, Eq, Clone, Debug)]
pub enum UdpResponse {
    Connect(UdpConnectionResponse),
    Announce(UdpAnnounceResponse),
    Scrape(UdpScrapeResponse),
    Error(UdpErrorResponse),
}

#[derive(PartialEq, Eq, Clone, Debug)]
pub struct UdpConnectionResponse {
    pub action: Actions,
    pub transaction_id: TransactionId,
    pub connection_id: ConnectionId,
}

#[derive(PartialEq, Eq, Clone, Debug)]
pub struct UdpAnnounceResponse {
    pub action: Actions,
    pub transaction_id: TransactionId,
    pub interval: u32,
    pub leechers: u32,
    pub seeders: u32,
    pub peers: Vec<TorrentPeer>,
}

#[derive(PartialEq, Eq, Clone, Debug)]
pub struct UdpScrapeResponse {
    pub action: Actions,
    pub transaction_id: TransactionId,
    pub torrent_stats: Vec<UdpScrapeResponseEntry>,
}

#[derive(PartialEq, Eq, Clone, Debug)]
pub struct UdpScrapeResponseEntry {
    pub seeders: i32,
    pub completed: i32,
    pub leechers: i32,
}

#[derive(PartialEq, Eq, Clone, Debug)]
pub struct UdpErrorResponse {
    pub action: Actions,
    pub transaction_id: TransactionId,
    pub message: String,
}

impl From<UdpConnectionResponse> for UdpResponse {
    fn from(r: UdpConnectionResponse) -> Self {
        Self::Connect(r)
    }
}

impl From<UdpAnnounceResponse> for UdpResponse {
    fn from(r: UdpAnnounceResponse) -> Self {
        Self::Announce(r)
    }
}

impl From<UdpScrapeResponse> for UdpResponse {
    fn from(r: UdpScrapeResponse) -> Self {
        Self::Scrape(r)
    }
}

impl From<UdpErrorResponse> for UdpResponse {
    fn from(r: UdpErrorResponse) -> Self {
        Self::Error(r)
    }
}

impl UdpResponse {
    pub fn write_to_bytes(self, bytes: &mut impl Write) -> Result<(), io::Error> {
        match self {
            UdpResponse::Connect(r) => {
                bytes.write_i32::<NetworkEndian>(0)?; // 0 = connect
                bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
                bytes.write_i64::<NetworkEndian>(r.connection_id.0)?;
            }
            UdpResponse::Announce(r) => {
                bytes.write_i32::<NetworkEndian>(1)?; // 1 = announce
                bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
                bytes.write_u32::<NetworkEndian>(r.interval)?;
                bytes.write_u32::<NetworkEndian>(r.leechers)?;
                bytes.write_u32::<NetworkEndian>(r.seeders)?;

                for peer in r.peers {
                    match peer.peer_addr {
                        SocketAddr::V4(socket_addr) => {
                            bytes.write_all(&socket_addr.ip().octets())?;
                            bytes.write_u16::<NetworkEndian>(socket_addr.port())?;
                        }
                        SocketAddr::V6(socket_addr) => {
                            bytes.write_all(&socket_addr.ip().octets())?;
                            bytes.write_u16::<NetworkEndian>(socket_addr.port())?;
                        }
                    }
                }
            }
            UdpResponse::Scrape(r) => {
                bytes.write_i32::<NetworkEndian>(2)?; // 2 = scrape
                bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;

                for torrent_stat in r.torrent_stats {
                    bytes.write_i32::<NetworkEndian>(torrent_stat.seeders)?;
                    bytes.write_i32::<NetworkEndian>(torrent_stat.completed)?;
                    bytes.write_i32::<NetworkEndian>(torrent_stat.leechers)?;
                }
            }
            UdpResponse::Error(r) => {
                bytes.write_i32::<NetworkEndian>(3)?;
                bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
                bytes.write_all(r.message.as_bytes())?;
            }
        }

        Ok(())
    }
}