arch/udpt
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

Changed InfoHash to a real struct.

Browse source
This commit is contained in:
Naim A 2018-10-28 01:59:35 +03:00
parent 703dfbb57d
commit 22bbb9d533
No known key found for this signature in database
GPG key ID: FD7948915D9EF8B9
3 changed files with 43 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
src/server.rs
View file

@ -199,10 +199,11 @@ impl UDPTracker {
}
let client_addr = SocketAddr::new(remote_addr.ip(), packet.port);
let info_hash = packet.info_hash.into();
match self.tracker.update_torrent_and_get_stats(&packet.info_hash, &packet.peer_id, &client_addr, packet.uploaded, packet.downloaded, packet.left, packet.event) {
match self.tracker.update_torrent_and_get_stats(&info_hash, &packet.peer_id, &client_addr, packet.uploaded, packet.downloaded, packet.left, packet.event) {
tracker::TorrentStats::Stats {leechers, complete: _, seeders} => {
let peers = match self.tracker.get_torrent_peers(&packet.info_hash, &client_addr) {
let peers = match self.tracker.get_torrent_peers(&info_hash, &client_addr) {
Some(v) => v,
None => {
return;

68
src/tracker.rs
View file

@ -1,6 +1,5 @@
use std;
use binascii;
use serde::{self, Serialize, Deserialize};
use server::Events;
@ -29,8 +28,36 @@ struct TorrentPeer {
updated: std::time::SystemTime,
}
#[derive(Ord, PartialEq, Eq, Clone)]
pub struct InfoHash {
info_hash: [u8; 20],
}
impl std::cmp::PartialOrd<InfoHash> for InfoHash {
fn partial_cmp(&self, other: &InfoHash) -> Option<std::cmp::Ordering> {
self.info_hash.partial_cmp(&other.info_hash)
}
}
impl std::convert::Into<InfoHash> for [u8; 20] {
fn into(self) -> InfoHash {
InfoHash{
info_hash: self,
}
}
}
impl serde::ser::Serialize for InfoHash {
fn serialize<S: serde::ser::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
let mut buffer = [0u8; 40];
let bytes_out = binascii::bin2hex(&self.info_hash, &mut buffer).ok().unwrap();
let str_out = std::str::from_utf8(bytes_out).unwrap();
serializer.serialize_str(str_out)
}
}
pub type PeerId = [u8; 20];
pub type InfoHash = [u8; 20];
#[derive(Serialize, Deserialize)]
pub struct TorrentEntry {
@ -120,37 +147,6 @@ impl Default for TorrentDatabase {
}
}
impl Serialize for TorrentDatabase {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
use serde::ser::SerializeMap;
use binascii::bin2hex;
let db = match self.torrent_peers.read() {
Ok(v) => v,
Err(err) => {
error!("Failed to lock database for reading. {}", err);
panic!("fatal error. check logs.");
}
};
let mut info_hash_buffer = [0u8; 40];
let mut map = serializer.serialize_map(None)?;
for entry in db.iter() {
let info_hash = match bin2hex(entry.0, &mut info_hash_buffer) {
Ok(v) => std::str::from_utf8(v).unwrap(),
Err(err) => {
error!("Failed to serialize database key.");
panic!("fatal error. check logs.")
}
};
map.serialize_entry(info_hash, entry.1)?;
}
map.end()
}
}
pub struct TorrentTracker {
mode: TrackerMode,
database: TorrentDatabase,
@ -179,7 +175,7 @@ impl TorrentTracker {
/// Adding torrents is not relevant to dynamic trackers.
pub fn add_torrent(&self, info_hash: &InfoHash) -> Result<(), ()> {
let mut write_lock = self.database.torrent_peers.write().unwrap();
match write_lock.entry(*info_hash) {
match write_lock.entry(info_hash.clone()) {
std::collections::btree_map::Entry::Vacant(ve) => {
ve.insert(TorrentEntry::new());
return Ok(());
@ -194,7 +190,7 @@ impl TorrentTracker {
pub fn remove_torrent(&self, info_hash: &InfoHash, force: bool) -> Result<(), ()> {
use std::collections::btree_map::Entry;
let mut entry_lock = self.database.torrent_peers.write().unwrap();
let torrent_entry = entry_lock.entry(*info_hash);
let torrent_entry = entry_lock.entry(info_hash.clone());
match torrent_entry {
Entry::Vacant(_) => {
// no entry, nothing to do...
@ -236,7 +232,7 @@ impl TorrentTracker {
pub fn update_torrent_and_get_stats(&self, info_hash: &InfoHash, peer_id: &PeerId, remote_address: &std::net::SocketAddr, uploaded: u64, downloaded: u64, left: u64, event: Events) -> TorrentStats {
use std::collections::btree_map::Entry;
let mut torrent_peers = self.database.torrent_peers.write().unwrap();
let torrent_entry = match torrent_peers.entry(*info_hash) {
let torrent_entry = match torrent_peers.entry(info_hash.clone()) {
Entry::Vacant(vacant) => {
match self.mode {
TrackerMode::DynamicMode => {

36
src/webserver.rs
View file

@ -12,9 +12,7 @@ const SERVER: &str = concat!("udpt/", env!("CARGO_PKG_VERSION"));
pub struct WebServer;
mod http_responses {
use std;
use binascii;
use serde;
use tracker::InfoHash;
#[derive(Serialize)]
pub struct TorrentInfo {
@ -29,8 +27,7 @@ mod http_responses {
pub offset: u32,
pub length: u32,
pub total: u32,
#[serde(serialize_with = "infohash_as_str")]
pub torrents: Vec<[u8; 20]>,
pub torrents: Vec<InfoHash>,
}
#[derive(Serialize)]
@ -40,23 +37,6 @@ mod http_responses {
TorrentList(TorrentList),
TorrentInfo(TorrentInfo),
}
fn infohash_as_str<S: serde::Serializer>(field: &Vec<[u8; 20]>, serializer: S) -> Result<S::Ok, S::Error> {
use serde::ser::SerializeSeq;
let mut output_str = [0u8; 40];
let mut seq = serializer.serialize_seq(Some(field.len()))?;
for infohash in field.iter() {
let _ = binascii::bin2hex(infohash, &mut output_str);
let mystr = std::str::from_utf8(&output_str).unwrap();
seq.serialize_element(mystr)?;
}
seq.end()
}
}
struct UdptState {
@ -265,7 +245,7 @@ impl WebServer {
let app_state: &UdptState = req.state();
let db = app_state.tracker.get_database();
let entry = match db.get(&info_hash) {
let entry = match db.get(&info_hash.into()) {
Some(v) => v,
None => {
return actix_web::HttpResponse::build(actix_web::http::StatusCode::NOT_FOUND)
@ -309,7 +289,7 @@ impl WebServer {
let path: actix_web::Path<String> = match actix_web::Path::extract(req) {
Ok(v) => v,
Err(err) => {
Err(_err) => {
return actix_web::HttpResponse::build(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR)
.json(http_responses::APIResponse::Error(String::from("internal_error")));
}
@ -324,19 +304,19 @@ impl WebServer {
match action.as_str() {
"flag" => {
app_state.tracker.set_torrent_flag(&info_hash, true);
app_state.tracker.set_torrent_flag(&info_hash.into(), true);
info!("Flagged {}", info_hash_str.as_str());
return actix_web::HttpResponse::build(actix_web::http::StatusCode::OK)
.body("")
},
"unflag" => {
app_state.tracker.set_torrent_flag(&info_hash, false);
app_state.tracker.set_torrent_flag(&info_hash.into(), false);
info!("Unflagged {}", info_hash_str.as_str());
return actix_web::HttpResponse::build(actix_web::http::StatusCode::OK)
.body("")
},
"add" => {
let success = app_state.tracker.add_torrent(&info_hash).is_ok();
let success = app_state.tracker.add_torrent(&info_hash.into()).is_ok();
info!("Added {}, success={}", info_hash_str.as_str(), success);
let code = if success { actix_web::http::StatusCode::OK } else { actix_web::http::StatusCode::INTERNAL_SERVER_ERROR };
@ -344,7 +324,7 @@ impl WebServer {
.body("")
},
"remove" => {
let success = app_state.tracker.remove_torrent(&info_hash, true).is_ok();
let success = app_state.tracker.remove_torrent(&info_hash.into(), true).is_ok();
info!("Removed {}, success={}", info_hash_str.as_str(), success);
let code = if success { actix_web::http::StatusCode::OK } else { actix_web::http::StatusCode::INTERNAL_SERVER_ERROR };