WIP: announce and DB better implementation
This commit is contained in:
Naim A
parent
dd04df33f6
commit
53eb629f50
657
src/server.rs
657
src/server.rs
|
|
@ -1,300 +1,357 @@
|
|||
use std;
|
||||
use std::net::{SocketAddr, UdpSocket};
|
||||
|
||||
use bincode;
|
||||
use serde::{Serialize, Deserialize, Serializer, Deserializer};
|
||||
|
||||
use tracker;
|
||||
|
||||
// 2000 should be enough, MTU is usually 1500
|
||||
const MAX_PACKET_SIZE: usize = 2000;
|
||||
|
||||
// protocol contants
|
||||
const PROTOCOL_ID: u64 = 0x0000041727101980;
|
||||
|
||||
#[repr(u32)]
|
||||
#[derive(Serialize, Deserialize)]
|
||||
enum Actions {
|
||||
Connect = 0,
|
||||
Announce = 1,
|
||||
Scrape = 2,
|
||||
Error = 3,
|
||||
}
|
||||
|
||||
#[repr(u32)]
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub enum Events {
|
||||
None = 0,
|
||||
Complete = 1,
|
||||
Started = 2,
|
||||
Stopped = 3,
|
||||
}
|
||||
|
||||
fn pack<T: Serialize>(data: &T) -> Option<Vec<u8>> {
|
||||
let mut bo = bincode::config();
|
||||
bo.big_endian();
|
||||
|
||||
match bo.serialize(data) {
|
||||
Ok(bytes) => Some(bytes),
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn unpack<'a, T: Deserialize<'a>>(data: &'a [u8]) -> Option<T> {
|
||||
let mut bo = bincode::config();
|
||||
bo.big_endian();
|
||||
|
||||
match bo.deserialize(data) {
|
||||
Ok(obj) => Some(obj),
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct UDPRequestHeader {
|
||||
connection_id: u64,
|
||||
action: Actions,
|
||||
transaction_id: u32,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct UDPResponseHeader {
|
||||
action: Actions,
|
||||
transaction_id: u32,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct UDPConnectionResponse {
|
||||
header: UDPResponseHeader,
|
||||
connection_id: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct UDPAnnounceRequest {
|
||||
header: UDPRequestHeader,
|
||||
|
||||
info_hash: [u8; 20],
|
||||
peer_id: [u8; 20],
|
||||
downloaded: u64,
|
||||
left: u64,
|
||||
uploaded: u64,
|
||||
event: Events,
|
||||
ip_address: u32,
|
||||
key: u32,
|
||||
num_want: i32,
|
||||
port: u16,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct UDPAnnounceResponse {
|
||||
header: UDPResponseHeader,
|
||||
|
||||
interval: u32,
|
||||
leechers: u32,
|
||||
seeders: u32,
|
||||
}
|
||||
|
||||
pub struct UDPTracker<'a> {
|
||||
server: std::net::UdpSocket,
|
||||
tracker: &'a mut tracker::TorrentTracker,
|
||||
}
|
||||
|
||||
impl<'a> UDPTracker<'a> {
|
||||
pub fn new<T: std::net::ToSocketAddrs>(bind_address: T, tracker: &mut tracker::TorrentTracker) -> Result<UDPTracker, std::io::Error> {
|
||||
let server = match UdpSocket::bind(bind_address) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
Ok(UDPTracker{
|
||||
server,
|
||||
tracker,
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_packet(&mut self, remote_address: &SocketAddr, payload: &[u8]) {
|
||||
let header : UDPRequestHeader = match unpack(payload) {
|
||||
Some(val) => val,
|
||||
None => {
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
match header.action {
|
||||
Actions::Connect => self.handle_connect(remote_address, &header, payload),
|
||||
Actions::Announce => self.handle_announce(remote_address, &header, payload),
|
||||
Actions::Scrape => self.handle_scrape(remote_address, &header, payload),
|
||||
_ => {
|
||||
// someone is playing around... ignore request.
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_connect(&self, remote_addr: &SocketAddr, header: &UDPRequestHeader, _payload: &[u8]) {
|
||||
if header.connection_id != PROTOCOL_ID {
|
||||
return;
|
||||
}
|
||||
|
||||
// send response...
|
||||
let conn_id = self.get_connection_id(remote_addr);
|
||||
|
||||
let response = UDPConnectionResponse{
|
||||
header: UDPResponseHeader{
|
||||
transaction_id: header.transaction_id,
|
||||
action: Actions::Connect,
|
||||
},
|
||||
connection_id: conn_id,
|
||||
};
|
||||
|
||||
if let Some(payload) = pack(&response) {
|
||||
let _ = self.send_packet(remote_addr, payload.as_slice());
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_announce(&mut self, remote_addr: &SocketAddr, header: &UDPRequestHeader, payload: &[u8]) {
|
||||
if header.connection_id != self.get_connection_id(remote_addr) {
|
||||
return;
|
||||
}
|
||||
|
||||
let announce_packet: Option<UDPAnnounceRequest> = unpack(payload);
|
||||
match announce_packet {
|
||||
Some(packet) => {
|
||||
let plen = bincode::serialized_size(&packet).unwrap() as usize;
|
||||
|
||||
println!("payload len={}, announce len={}", payload.len(), plen);
|
||||
|
||||
if payload.len() > plen {
|
||||
let bep41_payload = &payload[std::mem::size_of::<UDPAnnounceRequest>()..];
|
||||
println!("bep41: {:?}", bep41_payload);
|
||||
}
|
||||
|
||||
if packet.ip_address != 0 {
|
||||
// TODO: allow configurability of ip address
|
||||
// for now, ignore request.
|
||||
return;
|
||||
}
|
||||
|
||||
let client_addr = SocketAddr::new(remote_addr.ip(), packet.port);
|
||||
|
||||
self.tracker.update_torrent_peer(&packet.info_hash, &packet.peer_id, &client_addr, packet.uploaded, packet.downloaded, packet.left, packet.event);
|
||||
|
||||
let stats = self.tracker.get_stats(&packet.info_hash).unwrap_or((0, 0, 0));
|
||||
let results = self.tracker.get_peers(&packet.info_hash, &client_addr);
|
||||
|
||||
if let Some(mut payload) = pack(&UDPAnnounceResponse{
|
||||
header: UDPResponseHeader{
|
||||
action: Actions::Announce,
|
||||
transaction_id: packet.header.transaction_id,
|
||||
},
|
||||
seeders: stats.0 as u32,
|
||||
interval: 20,
|
||||
leechers: stats.2 as u32,
|
||||
}) {
|
||||
for peer in results {
|
||||
match client_addr {
|
||||
SocketAddr::V4(ipv4) => {
|
||||
payload.extend(&ipv4.ip().octets());
|
||||
},
|
||||
SocketAddr::V6(ipv6) => {
|
||||
payload.extend(&ipv6.ip().octets());
|
||||
}
|
||||
};
|
||||
|
||||
let port_hton = client_addr.port().to_be();
|
||||
payload.extend(&[ (port_hton & 0xff) as u8, ((port_hton >> 8) & 0xff) as u8 ]);
|
||||
}
|
||||
|
||||
} else {
|
||||
self.send_error(&remote_addr, &packet.header, "internal error");
|
||||
}
|
||||
},
|
||||
None => {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_scrape(&self, remote_addr: &SocketAddr, header: &UDPRequestHeader, payload: &[u8]) {
|
||||
if header.connection_id != self.get_connection_id(remote_addr) {
|
||||
return;
|
||||
}
|
||||
|
||||
self.send_error(remote_addr, header, "scrape not yet implemented");
|
||||
}
|
||||
|
||||
fn get_connection_id(&self, remote_address: &SocketAddr) -> u64 {
|
||||
match std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH) {
|
||||
Ok(duration) => {
|
||||
(duration.as_secs() / 3600) | ((remote_address.port() as u64) << 36)
|
||||
},
|
||||
Err(_) => {
|
||||
0x8000000000000000
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn send_packet(&self, remote_addr: &SocketAddr, payload: &[u8]) -> Result<usize, std::io::Error> {
|
||||
self.server.send_to(payload, remote_addr)
|
||||
}
|
||||
|
||||
fn send_error(&self, remote_addr: &SocketAddr, header: &UDPRequestHeader, error_msg: &str) {
|
||||
if let Some(mut payload) = pack(&UDPResponseHeader{
|
||||
transaction_id: header.transaction_id,
|
||||
action: Actions::Error,
|
||||
}) {
|
||||
let msg_bytes = Vec::from(error_msg.as_bytes());
|
||||
payload.extend(msg_bytes);
|
||||
|
||||
let _ = self.send_packet(remote_addr, payload.as_slice());
|
||||
}
|
||||
}
|
||||
|
||||
pub fn accept_packet(&mut self) -> Result<(), std::io::Error> {
|
||||
let mut packet = [0u8; MAX_PACKET_SIZE];
|
||||
match self.server.recv_from(&mut packet) {
|
||||
Ok((size, remote_address)) => {
|
||||
self.handle_packet(&remote_address, &packet[..size]);
|
||||
|
||||
Ok(())
|
||||
},
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[test]
|
||||
fn pack() {
|
||||
let mystruct = super::UDPRequestHeader {
|
||||
connection_id: 200,
|
||||
action: super::Actions::Connect,
|
||||
transaction_id: 77771,
|
||||
};
|
||||
match super::pack(&mystruct) {
|
||||
Some(data) => {
|
||||
println!("serialized data = [{}, {:?}]", data.len(), data);
|
||||
},
|
||||
None => {
|
||||
assert!(false);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unpack() {
|
||||
let buf = [0u8, 0, 0, 0, 0, 0, 0, 200, 0, 0, 0, 1, 0, 1, 47, 203];
|
||||
match super::unpack(&buf) {
|
||||
Some(obj) => {
|
||||
let x : super::UDPResponseHeader = obj;
|
||||
println!("conn_id={}", x.action as u32);
|
||||
},
|
||||
None => {
|
||||
assert!(false);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
use std;
|
||||
use std::net::{SocketAddr, UdpSocket};
|
||||
|
||||
use bincode;
|
||||
use serde::{Serialize, Deserialize};
|
||||
|
||||
use tracker;
|
||||
|
||||
// 2000 should be enough, MTU is usually 1500
|
||||
const MAX_PACKET_SIZE: usize = 2000;
|
||||
|
||||
// protocol contants
|
||||
const PROTOCOL_ID: u64 = 0x0000041727101980;
|
||||
|
||||
#[repr(u32)]
|
||||
#[derive(Serialize, Deserialize)]
|
||||
enum Actions {
|
||||
Connect = 0,
|
||||
Announce = 1,
|
||||
Scrape = 2,
|
||||
Error = 3,
|
||||
}
|
||||
|
||||
#[repr(u32)]
|
||||
#[derive(Serialize, Deserialize, Clone, Copy)]
|
||||
pub enum Events {
|
||||
None = 0,
|
||||
Complete = 1,
|
||||
Started = 2,
|
||||
Stopped = 3,
|
||||
}
|
||||
|
||||
fn pack<T: Serialize>(data: &T) -> Option<Vec<u8>> {
|
||||
let mut bo = bincode::config();
|
||||
bo.big_endian();
|
||||
|
||||
match bo.serialize(data) {
|
||||
Ok(bytes) => Some(bytes),
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn unpack<'a, T: Deserialize<'a>>(data: &'a [u8]) -> Option<T> {
|
||||
let mut bo = bincode::config();
|
||||
bo.big_endian();
|
||||
|
||||
match bo.deserialize(data) {
|
||||
Ok(obj) => Some(obj),
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct UDPRequestHeader {
|
||||
connection_id: u64,
|
||||
action: Actions,
|
||||
transaction_id: u32,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct UDPResponseHeader {
|
||||
action: Actions,
|
||||
transaction_id: u32,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct UDPConnectionResponse {
|
||||
header: UDPResponseHeader,
|
||||
connection_id: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct UDPAnnounceRequest {
|
||||
header: UDPRequestHeader,
|
||||
|
||||
info_hash: [u8; 20],
|
||||
peer_id: [u8; 20],
|
||||
downloaded: u64,
|
||||
left: u64,
|
||||
uploaded: u64,
|
||||
event: Events,
|
||||
ip_address: u32,
|
||||
key: u32,
|
||||
num_want: i32,
|
||||
port: u16,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct UDPAnnounceResponse {
|
||||
header: UDPResponseHeader,
|
||||
|
||||
interval: u32,
|
||||
leechers: u32,
|
||||
seeders: u32,
|
||||
}
|
||||
|
||||
pub struct UDPTracker<'a> {
|
||||
server: std::net::UdpSocket,
|
||||
tracker: &'a mut tracker::TorrentTracker,
|
||||
}
|
||||
|
||||
impl<'a> UDPTracker<'a> {
|
||||
pub fn new<T: std::net::ToSocketAddrs>(bind_address: T, tracker: &mut tracker::TorrentTracker) -> Result<UDPTracker, std::io::Error> {
|
||||
let server = match UdpSocket::bind(bind_address) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
Ok(UDPTracker{
|
||||
server,
|
||||
tracker,
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_packet(&mut self, remote_address: &SocketAddr, payload: &[u8]) {
|
||||
let header : UDPRequestHeader = match unpack(payload) {
|
||||
Some(val) => val,
|
||||
None => {
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
match header.action {
|
||||
Actions::Connect => self.handle_connect(remote_address, &header, payload),
|
||||
Actions::Announce => self.handle_announce(remote_address, &header, payload),
|
||||
Actions::Scrape => self.handle_scrape(remote_address, &header, payload),
|
||||
_ => {
|
||||
// someone is playing around... ignore request.
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_connect(&self, remote_addr: &SocketAddr, header: &UDPRequestHeader, _payload: &[u8]) {
|
||||
if header.connection_id != PROTOCOL_ID {
|
||||
return;
|
||||
}
|
||||
|
||||
// send response...
|
||||
let conn_id = self.get_connection_id(remote_addr);
|
||||
|
||||
let response = UDPConnectionResponse{
|
||||
header: UDPResponseHeader{
|
||||
transaction_id: header.transaction_id,
|
||||
action: Actions::Connect,
|
||||
},
|
||||
connection_id: conn_id,
|
||||
};
|
||||
|
||||
if let Some(payload) = pack(&response) {
|
||||
let _ = self.send_packet(remote_addr, payload.as_slice());
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_announce(&mut self, remote_addr: &SocketAddr, header: &UDPRequestHeader, payload: &[u8]) {
|
||||
if header.connection_id != self.get_connection_id(remote_addr) {
|
||||
return;
|
||||
}
|
||||
|
||||
let packet: UDPAnnounceRequest = match unpack(payload) {
|
||||
Some(v) => v,
|
||||
None => {
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let plen = bincode::serialized_size(&packet).unwrap() as usize;
|
||||
|
||||
println!("payload len={}, announce len={}", payload.len(), plen);
|
||||
|
||||
if payload.len() > plen {
|
||||
let bep41_payload = &payload[std::mem::size_of::<UDPAnnounceRequest>()..];
|
||||
println!("bep41: {:?}", bep41_payload);
|
||||
}
|
||||
|
||||
if packet.ip_address != 0 {
|
||||
// TODO: allow configurability of ip address
|
||||
// for now, ignore request.
|
||||
return;
|
||||
}
|
||||
|
||||
let client_addr = SocketAddr::new(remote_addr.ip(), packet.port);
|
||||
match self.tracker.get_torrent(&packet.info_hash, |torrent: &mut tracker::TorrentEntry | -> Result<Vec<u8>, &str> {
|
||||
if torrent.is_flagged() {
|
||||
return Err("Torrent has been flagged.");
|
||||
}
|
||||
torrent.update_peer(&packet.peer_id, &client_addr, packet.uploaded, packet.downloaded, packet.left, packet.event);
|
||||
|
||||
let stats = torrent.get_stats();
|
||||
let peers = torrent.get_peers(&client_addr);
|
||||
if let Some(mut payload) = pack(&UDPAnnounceResponse {
|
||||
header: UDPResponseHeader {
|
||||
action: Actions::Announce,
|
||||
transaction_id: packet.header.transaction_id,
|
||||
},
|
||||
seeders: stats.0 as u32,
|
||||
interval: 20,
|
||||
leechers: stats.2 as u32,
|
||||
}) {
|
||||
for peer in peers {
|
||||
match peer {
|
||||
SocketAddr::V4(ipv4) => {
|
||||
payload.extend(&ipv4.ip().octets());
|
||||
},
|
||||
SocketAddr::V6(ipv6) => {
|
||||
payload.extend(&ipv6.ip().octets());
|
||||
}
|
||||
};
|
||||
|
||||
let port_hton = client_addr.port().to_be();
|
||||
payload.extend(&[(port_hton & 0xff) as u8, ((port_hton >> 8) & 0xff) as u8]);
|
||||
|
||||
return Ok(payload);
|
||||
}
|
||||
}
|
||||
return Err("");
|
||||
}) {
|
||||
Some(error_message) => {
|
||||
match error_message {
|
||||
Err(err) => {
|
||||
if err.len() > 0 {
|
||||
self.send_error(remote_addr, header, err);
|
||||
}
|
||||
},
|
||||
Ok(payload) => {
|
||||
// send packet...
|
||||
let _ = self.send_packet(remote_addr, payload.as_slice());
|
||||
}
|
||||
}
|
||||
},
|
||||
None => {
|
||||
self.send_error(remote_addr, header, "Unregistered torrent");
|
||||
return;
|
||||
}
|
||||
};
|
||||
/*
|
||||
if torrent.is_flagged() {
|
||||
error = Some("Torrent was flagged");
|
||||
} else {
|
||||
torrent.update_peer(&packet.peer_id, &client_addr, packet.uploaded, packet.downloaded, packet.left, packet.event);
|
||||
|
||||
let stats = torrent.get_stats();
|
||||
let peers = torrent.get_peers(&client_addr);
|
||||
|
||||
if let Some(mut payload) = pack(&UDPAnnounceResponse {
|
||||
header: UDPResponseHeader {
|
||||
action: Actions::Announce,
|
||||
transaction_id: packet.header.transaction_id,
|
||||
},
|
||||
seeders: stats.0 as u32,
|
||||
interval: 20,
|
||||
leechers: stats.2 as u32,
|
||||
}) {
|
||||
for peer in peers {
|
||||
match peer {
|
||||
SocketAddr::V4(ipv4) => {
|
||||
payload.extend(&ipv4.ip().octets());
|
||||
},
|
||||
SocketAddr::V6(ipv6) => {
|
||||
payload.extend(&ipv6.ip().octets());
|
||||
}
|
||||
};
|
||||
|
||||
let port_hton = client_addr.port().to_be();
|
||||
payload.extend(&[(port_hton & 0xff) as u8, ((port_hton >> 8) & 0xff) as u8]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(error_message) = error {
|
||||
self.send_error(remote_addr, &packet.header, error_message);
|
||||
}*/
|
||||
}
|
||||
|
||||
fn handle_scrape(&self, remote_addr: &SocketAddr, header: &UDPRequestHeader, payload: &[u8]) {
|
||||
if header.connection_id != self.get_connection_id(remote_addr) {
|
||||
return;
|
||||
}
|
||||
|
||||
self.send_error(remote_addr, header, "scrape not yet implemented");
|
||||
}
|
||||
|
||||
fn get_connection_id(&self, remote_address: &SocketAddr) -> u64 {
|
||||
match std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH) {
|
||||
Ok(duration) => {
|
||||
(duration.as_secs() / 3600) | ((remote_address.port() as u64) << 36)
|
||||
},
|
||||
Err(_) => {
|
||||
0x8000000000000000
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn send_packet(&self, remote_addr: &SocketAddr, payload: &[u8]) -> Result<usize, std::io::Error> {
|
||||
self.server.send_to(payload, remote_addr)
|
||||
}
|
||||
|
||||
fn send_error(&self, remote_addr: &SocketAddr, header: &UDPRequestHeader, error_msg: &str) {
|
||||
if let Some(mut payload) = pack(&UDPResponseHeader{
|
||||
transaction_id: header.transaction_id,
|
||||
action: Actions::Error,
|
||||
}) {
|
||||
let msg_bytes = Vec::from(error_msg.as_bytes());
|
||||
payload.extend(msg_bytes);
|
||||
|
||||
let _ = self.send_packet(remote_addr, payload.as_slice());
|
||||
}
|
||||
}
|
||||
|
||||
pub fn accept_packet(&mut self) -> Result<(), std::io::Error> {
|
||||
let mut packet = [0u8; MAX_PACKET_SIZE];
|
||||
match self.server.recv_from(&mut packet) {
|
||||
Ok((size, remote_address)) => {
|
||||
self.handle_packet(&remote_address, &packet[..size]);
|
||||
|
||||
Ok(())
|
||||
},
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[test]
|
||||
fn pack() {
|
||||
let mystruct = super::UDPRequestHeader {
|
||||
connection_id: 200,
|
||||
action: super::Actions::Connect,
|
||||
transaction_id: 77771,
|
||||
};
|
||||
match super::pack(&mystruct) {
|
||||
Some(data) => {
|
||||
println!("serialized data = [{}, {:?}]", data.len(), data);
|
||||
},
|
||||
None => {
|
||||
assert!(false);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unpack() {
|
||||
let buf = [0u8, 0, 0, 0, 0, 0, 0, 200, 0, 0, 0, 1, 0, 1, 47, 203];
|
||||
match super::unpack(&buf) {
|
||||
Some(obj) => {
|
||||
let x : super::UDPResponseHeader = obj;
|
||||
println!("conn_id={}", x.action as u32);
|
||||
},
|
||||
None => {
|
||||
assert!(false);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
320
src/tracker.rs
320
src/tracker.rs
|
|
@ -1,149 +1,171 @@
|
|||
use std;
|
||||
|
||||
use server::Events;
|
||||
|
||||
pub enum TrackerMode {
|
||||
|
||||
/// In static mode torrents are tracked only if they were added ahead of time.
|
||||
StaticMode,
|
||||
|
||||
/// In dynamic mode, torrents are tracked being added ahead of time.
|
||||
DynamicMode,
|
||||
|
||||
/// Tracker will only serve authenticated peers.
|
||||
PrivateMode,
|
||||
}
|
||||
|
||||
struct TorrentPeer {
|
||||
ip: std::net::SocketAddr,
|
||||
last_connection_id: u64,
|
||||
uploaded: u64,
|
||||
downloaded: u64,
|
||||
left: u64,
|
||||
event: Events,
|
||||
updated: std::time::SystemTime,
|
||||
}
|
||||
|
||||
type PeerId = [u8; 20];
|
||||
type InfoHash = [u8; 20];
|
||||
|
||||
struct TorrentEntry {
|
||||
is_flagged: bool,
|
||||
peers: std::collections::BTreeMap<PeerId, TorrentPeer>,
|
||||
}
|
||||
|
||||
struct TorrentDatabase {
|
||||
torrent_peers: std::collections::BTreeMap<InfoHash, TorrentEntry>,
|
||||
}
|
||||
|
||||
pub struct TorrentTracker {
|
||||
mode: TrackerMode,
|
||||
database: TorrentDatabase,
|
||||
}
|
||||
|
||||
impl TorrentTracker {
|
||||
pub fn new() -> TorrentTracker {
|
||||
TorrentTracker{
|
||||
mode: TrackerMode::DynamicMode,
|
||||
database: TorrentDatabase{
|
||||
torrent_peers: std::collections::BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Adding torrents is not relevant to dynamic trackers.
|
||||
pub fn add_torrent(&mut self, info_hash: &InfoHash) {
|
||||
use std::collections::BTreeMap;
|
||||
self.database.torrent_peers.entry(*info_hash).or_insert(TorrentEntry{
|
||||
is_flagged: false,
|
||||
peers: std::collections::BTreeMap::new(),
|
||||
});
|
||||
}
|
||||
|
||||
/// If the torrent is flagged, it will not be removed unless force is set to true.
|
||||
pub fn remove_torrent(&mut self, info_hash: &InfoHash, force: bool) {
|
||||
if !force {
|
||||
if let Some(entry) = self.database.torrent_peers.get(info_hash) {
|
||||
if entry.is_flagged {
|
||||
// torrent is flagged, ignore request.
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
// torrent not found, no point looking for it again...
|
||||
return;
|
||||
}
|
||||
}
|
||||
self.database.torrent_peers.remove(info_hash);
|
||||
}
|
||||
|
||||
/// flagged torrents will result in a tracking error. This is to allow enforcement against piracy.
|
||||
pub fn set_torrent_flag(&mut self, info_hash: &InfoHash, is_flagged: bool) {
|
||||
if let Some(mut entry) = self.database.torrent_peers.get_mut(info_hash) {
|
||||
if is_flagged && !entry.is_flagged {
|
||||
// empty peer list.
|
||||
entry.peers.clear();
|
||||
}
|
||||
entry.is_flagged = is_flagged;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn update_torrent_peer(&mut self, info_hash: &InfoHash, peer_id: &PeerId, remote_address: &std::net::SocketAddr, uploaded: u64, downloaded: u64, left: u64, event: Events) {
|
||||
if let Some(mut torrent_entry) = self.database.torrent_peers.get_mut(info_hash) {
|
||||
torrent_entry.peers.insert(*peer_id, TorrentPeer{
|
||||
updated: std::time::SystemTime::now(),
|
||||
left,
|
||||
downloaded,
|
||||
uploaded,
|
||||
ip: *remote_address,
|
||||
event,
|
||||
last_connection_id: 0,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// returns a list of peers with the same type of address of the remote_addr (IP v4/v6)
|
||||
pub fn get_peers(&self, info_hash: &InfoHash, remote_addr: &std::net::SocketAddr) -> Vec<std::net::SocketAddr> {
|
||||
let mut list = Vec::new();
|
||||
if let Some(entry) = self.database.torrent_peers.get(info_hash) {
|
||||
for (_, peer) in entry.peers.iter().filter(|e| e.1.ip.is_ipv4() == remote_addr.is_ipv4()) {
|
||||
if peer.ip == *remote_addr {
|
||||
continue;
|
||||
}
|
||||
|
||||
list.push(peer.ip);
|
||||
|
||||
if list.len() >= 74 {
|
||||
// 74 is maximum peers supported by the protocol.
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
list
|
||||
}
|
||||
|
||||
pub fn get_stats(&self, info_hash: &InfoHash) -> Option<(i32, i32, i32)> {
|
||||
if let Some(torrent_entry) = self.database.torrent_peers.get(info_hash) {
|
||||
|
||||
// TODO: store stats in temporary location...
|
||||
let mut seeders = 0;
|
||||
let mut leechers = 0;
|
||||
|
||||
for (_, peer) in torrent_entry.peers.iter() {
|
||||
if peer.left == 0 {
|
||||
seeders += 1;
|
||||
}
|
||||
else {
|
||||
leechers += 1;
|
||||
}
|
||||
}
|
||||
|
||||
Some((seeders, -1, leechers))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub fn cleanup(&mut self) {
|
||||
|
||||
}
|
||||
}
|
||||
use std;
|
||||
|
||||
use server::Events;
|
||||
|
||||
pub enum TrackerMode {
|
||||
|
||||
/// In static mode torrents are tracked only if they were added ahead of time.
|
||||
StaticMode,
|
||||
|
||||
/// In dynamic mode, torrents are tracked being added ahead of time.
|
||||
DynamicMode,
|
||||
|
||||
/// Tracker will only serve authenticated peers.
|
||||
PrivateMode,
|
||||
}
|
||||
|
||||
struct TorrentPeer {
|
||||
ip: std::net::SocketAddr,
|
||||
uploaded: u64,
|
||||
downloaded: u64,
|
||||
left: u64,
|
||||
event: Events,
|
||||
updated: std::time::SystemTime,
|
||||
}
|
||||
|
||||
type PeerId = [u8; 20];
|
||||
type InfoHash = [u8; 20];
|
||||
|
||||
pub struct TorrentEntry {
|
||||
is_flagged: bool,
|
||||
peers: std::collections::BTreeMap<PeerId, TorrentPeer>,
|
||||
|
||||
completed: u32,
|
||||
seeders: u32,
|
||||
}
|
||||
|
||||
impl TorrentEntry {
|
||||
pub fn is_flagged(&self) -> bool {
|
||||
self.is_flagged
|
||||
}
|
||||
|
||||
pub fn update_peer(&mut self, peer_id: &PeerId, remote_address: &std::net::SocketAddr, uploaded: u64, downloaded: u64, left: u64, event: Events) {
|
||||
let is_seeder = left == 0 && uploaded > 0;
|
||||
let mut was_seeder = false;
|
||||
let mut is_completed = left == 0 && (event as u32) == (Events::Complete as u32);
|
||||
if let Some(prev) = self.peers.insert(*peer_id, TorrentPeer{
|
||||
updated: std::time::SystemTime::now(),
|
||||
left,
|
||||
downloaded,
|
||||
uploaded,
|
||||
ip: *remote_address,
|
||||
event,
|
||||
}) {
|
||||
was_seeder = prev.left == 0 && prev.uploaded > 0;
|
||||
|
||||
if is_completed && (prev.event as u32) == (Events::Complete as u32) {
|
||||
// don't update count again. a torrent should only be updated once per peer.
|
||||
is_completed = false;
|
||||
}
|
||||
}
|
||||
|
||||
if is_seeder && !was_seeder {
|
||||
self.seeders += 1;
|
||||
} else if was_seeder && !is_seeder {
|
||||
self.seeders -= 1;
|
||||
}
|
||||
|
||||
if is_completed {
|
||||
self.completed += 1;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_peers(&self, remote_addr: &std::net::SocketAddr) -> Vec<std::net::SocketAddr> {
|
||||
let mut list = Vec::new();
|
||||
for (_, peer) in self.peers.iter().filter(|e| e.1.ip.is_ipv4() == remote_addr.is_ipv4()) {
|
||||
if peer.ip == *remote_addr {
|
||||
continue;
|
||||
}
|
||||
|
||||
list.push(peer.ip);
|
||||
|
||||
if list.len() >= 74 {
|
||||
// 74 is maximum peers supported by the protocol.
|
||||
break;
|
||||
}
|
||||
}
|
||||
list
|
||||
}
|
||||
|
||||
pub fn get_stats(&self) -> (u32, u32, u32) {
|
||||
let leechers = (self.peers.len() as u32) - self.seeders;
|
||||
(self.seeders, self.completed, leechers)
|
||||
}
|
||||
}
|
||||
|
||||
struct TorrentDatabase {
|
||||
torrent_peers: std::collections::BTreeMap<InfoHash, TorrentEntry>,
|
||||
}
|
||||
|
||||
pub struct TorrentTracker {
|
||||
mode: TrackerMode,
|
||||
database: TorrentDatabase,
|
||||
}
|
||||
|
||||
impl TorrentTracker {
|
||||
pub fn new() -> TorrentTracker {
|
||||
TorrentTracker{
|
||||
mode: TrackerMode::DynamicMode,
|
||||
database: TorrentDatabase{
|
||||
torrent_peers: std::collections::BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Adding torrents is not relevant to dynamic trackers.
|
||||
pub fn add_torrent(&mut self, info_hash: &InfoHash) {
|
||||
self.database.torrent_peers.entry(*info_hash).or_insert(TorrentEntry{
|
||||
is_flagged: false,
|
||||
peers: std::collections::BTreeMap::new(),
|
||||
seeders: 0,
|
||||
completed: 0,
|
||||
});
|
||||
}
|
||||
|
||||
/// If the torrent is flagged, it will not be removed unless force is set to true.
|
||||
pub fn remove_torrent(&mut self, info_hash: &InfoHash, force: bool) {
|
||||
if !force {
|
||||
if let Some(entry) = self.database.torrent_peers.get(info_hash) {
|
||||
if entry.is_flagged {
|
||||
// torrent is flagged, ignore request.
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
// torrent not found, no point looking for it again...
|
||||
return;
|
||||
}
|
||||
}
|
||||
self.database.torrent_peers.remove(info_hash);
|
||||
}
|
||||
|
||||
/// flagged torrents will result in a tracking error. This is to allow enforcement against piracy.
|
||||
pub fn set_torrent_flag(&mut self, info_hash: &InfoHash, is_flagged: bool) {
|
||||
if let Some(mut entry) = self.database.torrent_peers.get_mut(info_hash) {
|
||||
if is_flagged && !entry.is_flagged {
|
||||
// empty peer list.
|
||||
entry.peers.clear();
|
||||
}
|
||||
entry.is_flagged = is_flagged;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_torrent<F, R>(&mut self, info_hash: &InfoHash, action: F) -> Option<R>
|
||||
where F: Fn(&mut TorrentEntry) -> R
|
||||
{
|
||||
if let Some(torrent_entry) = self.database.torrent_peers.get_mut(info_hash) {
|
||||
Some(action(torrent_entry))
|
||||
} else {
|
||||
match self.mode {
|
||||
TrackerMode::StaticMode => None,
|
||||
TrackerMode::PrivateMode => None,
|
||||
TrackerMode::DynamicMode => {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn cleanup(&mut self) {
|
||||
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue