From 9a314220910e06f83a3d9d420bb930bbd357b983 Mon Sep 17 00:00:00 2001
From: Naim A <227396+naim94a@users.noreply.github.com>
Date: Sun, 21 Oct 2018 23:39:39 +0300
Subject: [PATCH] More values are configurable

---
 Cargo.lock       | 10 ++++++++++
 src/config.rs    | 39 ++++++++++++++++++++++++++++++++++++---
 src/main.rs      | 22 ++++++++++------------
 src/server.rs    | 12 +++++++++---
 src/tracker.rs   |  6 +++---
 src/webserver.rs | 39 +++++++++++++++++++++++++++++++--------
 udpt.toml        | 11 +++++++++++
 7 files changed, 110 insertions(+), 29 deletions(-)
 create mode 100644 udpt.toml

diff --git a/Cargo.lock b/Cargo.lock
index 253d0a2..0ef0f8d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1260,6 +1260,14 @@ dependencies = [
  "tokio-reactor 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
+[[package]]
+name = "toml"
+version = "0.4.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "serde 1.0.80 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
 [[package]]
 name = "tower-service"
 version = "0.1.0"
@@ -1364,6 +1372,7 @@ dependencies = [
  "bincode 1.0.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde 1.0.80 (registry+https://github.com/rust-lang/crates.io-index)",
  "serde_derive 1.0.80 (registry+https://github.com/rust-lang/crates.io-index)",
+ "toml 0.4.8 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
 [[package]]
@@ -1644,6 +1653,7 @@ dependencies = [
 "checksum tokio-timer 0.2.7 (registry+https://github.com/rust-lang/crates.io-index)" = "3a52f00c97fedb6d535d27f65cccb7181c8dd4c6edc3eda9ea93f6d45d05168e"
 "checksum tokio-udp 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "da941144b816d0dcda4db3a1ba87596e4df5e860a72b70783fe435891f80601c"
 "checksum tokio-uds 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "22e3aa6d1fcc19e635418dc0a30ab5bd65d347973d6f43f1a37bf8d9d1335fc9"
+"checksum toml 0.4.8 (registry+https://github.com/rust-lang/crates.io-index)" = "4a2ecc31b0351ea18b3fe11274b8db6e4d82bce861bbb22e6dbed40417902c65"
 "checksum tower-service 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "b32f72af77f1bfe3d3d4da8516a238ebe7039b51dd8637a09841ac7f16d2c987"
 "checksum trust-dns-proto 0.4.3 (registry+https://github.com/rust-lang/crates.io-index)" = "f1525ca4e26f5a09d81b79584f19225e7dba5606ae3a416311c2751c5cea60bb"
 "checksum trust-dns-proto 0.5.0-alpha.3 (registry+https://github.com/rust-lang/crates.io-index)" = "3fabc184ed90d027afee46386e6418b9c953b7be527f62cc37724a1720e07d68"
diff --git a/src/config.rs b/src/config.rs
index 423a6d2..545d6f4 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -7,7 +7,17 @@ pub use tracker::TrackerMode;
 #[derive(Deserialize)]
 pub struct UDPConfig {
     bind_address: String,
-    mode: TrackerMode,
+    announce_interval: u32,
+}
+
+impl UDPConfig {
+    pub fn get_address(&self) -> &str {
+        self.bind_address.as_str()
+    }
+
+    pub fn get_announce_interval(&self) -> u32 {
+        self.announce_interval
+    }
 }
 
 #[derive(Deserialize)]
@@ -16,8 +26,19 @@ pub struct HTTPConfig {
     access_tokens: HashMap<String, String>,
 }
 
+impl HTTPConfig {
+    pub fn get_address(&self) -> &str {
+        self.bind_address.as_str()
+    }
+
+    pub fn get_access_tokens(&self) -> &HashMap<String, String> {
+        &self.access_tokens
+    }
+}
+
 #[derive(Deserialize)]
 pub struct Configuration {
+    mode: TrackerMode,
     udp: UDPConfig,
     http: Option<HTTPConfig>,
 }
@@ -39,7 +60,6 @@ impl std::fmt::Display for ConfigError {
 }
 impl std::error::Error for ConfigError {}
 
-
 impl Configuration {
     pub fn load(data: &[u8]) -> Result<Configuration, toml::de::Error> {
         toml::from_slice(data)
@@ -56,14 +76,27 @@ impl Configuration {
             }
         }
     }
+
+    pub fn get_mode(&self) -> &TrackerMode {
+        &self.mode
+    }
+
+    pub fn get_udp_config(&self) -> &UDPConfig {
+        &self.udp
+    }
+
+    pub fn get_http_config(&self) -> &Option<HTTPConfig> {
+        &self.http
+    }
 }
 
 impl Default for Configuration {
     fn default() -> Configuration {
         Configuration{
+            mode: TrackerMode::DynamicMode,
             udp: UDPConfig{
+                announce_interval: 120,
                 bind_address: String::from("0.0.0.0:6969"),
-                mode: TrackerMode::DynamicMode,
             },
             http: None,
         }
diff --git a/src/main.rs b/src/main.rs
index 33adca5..0156fc8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -14,27 +14,25 @@ use config::Configuration;
 
 fn main() {
     let cfg = match Configuration::load_file("udpt.toml") {
-        Ok(v) => v,
+        Ok(v) => std::sync::Arc::new(v),
         Err(e) => {
             eprintln!("failed to open configuration: {}", e);
             return;
         }
     };
 
-
-    let tracker = std::sync::Arc::new(tracker::TorrentTracker::new());
+    let tracker = std::sync::Arc::new(tracker::TorrentTracker::new(cfg.get_mode().clone()));
 
     // start http server:
-    let mut access_tokens = std::collections::HashMap::new();
-    access_tokens.insert(String::from("MySpecialToken"), String::from("username"));
+    if let Some(http_cfg) = cfg.get_http_config() {
+        let http_tracker_ref = tracker.clone();
+        let cfg_ref = cfg.clone();
+        std::thread::spawn(move || {
+            webserver::WebServer::new(http_tracker_ref, cfg_ref);
+        });
+    }
 
-    let http_tracker_ref = tracker.clone();
-    std::thread::spawn(move || {
-        webserver::WebServer::new(http_tracker_ref);
-    });
-
-    let addr = "0.0.0.0:1212";
-    let s = std::sync::Arc::new(server::UDPTracker::new(addr, tracker.clone()).unwrap());
+    let s = std::sync::Arc::new(server::UDPTracker::new(cfg, tracker.clone()).unwrap());
 
     loop {
         match s.accept_packet() {
diff --git a/src/server.rs b/src/server.rs
index 3236680..8c3a336 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1,4 +1,5 @@
 use std;
+use std::sync::Arc;
 use std::net::{SocketAddr, UdpSocket};
 use std::io::Write;
 
@@ -7,6 +8,7 @@ use serde::{Serialize, Deserialize};
 
 use tracker;
 use stackvec::StackVec;
+use config::Configuration;
 
 // maximum MTU is usually 1500, but our stack allows us to allocate the maximum - so why not?
 const MAX_PACKET_SIZE: usize = 0xffff;
@@ -109,11 +111,14 @@ struct UDPAnnounceResponse {
 pub struct UDPTracker {
     server: std::net::UdpSocket,
     tracker: std::sync::Arc<tracker::TorrentTracker>,
+    config: Arc<Configuration>,
 }
 
 impl UDPTracker {
-    pub fn new<T: std::net::ToSocketAddrs>(bind_address: T, tracker: std::sync::Arc<tracker::TorrentTracker>) -> Result<UDPTracker, std::io::Error> {
-        let server = match UdpSocket::bind(bind_address) {
+    pub fn new(config: Arc<Configuration>, tracker: std::sync::Arc<tracker::TorrentTracker>) -> Result<UDPTracker, std::io::Error> {
+        let cfg = config.clone();
+
+        let server = match UdpSocket::bind(cfg.get_udp_config().get_address()) {
             Ok(s) => s,
             Err(e) => {
                 return Err(e);
@@ -123,6 +128,7 @@ impl UDPTracker {
         Ok(UDPTracker{
             server,
             tracker,
+            config: cfg,
         })
     }
 
@@ -216,7 +222,7 @@ impl UDPTracker {
                         transaction_id: packet.header.transaction_id,
                     },
                     seeders,
-                    interval: 20,
+                    interval: self.config.get_udp_config().get_announce_interval(),
                     leechers,
                 }) {
                     Ok(_) => {},
diff --git a/src/tracker.rs b/src/tracker.rs
index ac84042..6d5cefd 100644
--- a/src/tracker.rs
+++ b/src/tracker.rs
@@ -2,7 +2,7 @@ use std;
 
 use server::Events;
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Clone)]
 pub enum TrackerMode {
 
     /// In static mode torrents are tracked only if they were added ahead of time.
@@ -179,9 +179,9 @@ pub enum TorrentStats {
 }
 
 impl TorrentTracker {
-    pub fn new() -> TorrentTracker {
+    pub fn new(mode: TrackerMode) -> TorrentTracker {
         TorrentTracker{
-            mode: TrackerMode::DynamicMode,
+            mode,
             database: TorrentDatabase{
                 torrent_peers: std::sync::RwLock::new(std::collections::BTreeMap::new()),
             }
diff --git a/src/webserver.rs b/src/webserver.rs
index 9bf0add..67b3b17 100644
--- a/src/webserver.rs
+++ b/src/webserver.rs
@@ -4,6 +4,7 @@ use std::sync::Arc;
 use actix_web;
 use binascii;
 
+use config;
 use tracker;
 
 const SERVER: &str = concat!("udpt/", env!("CARGO_PKG_VERSION"));
@@ -126,9 +127,25 @@ impl actix_web::middleware::Middleware<UdptState> for UdptMiddleware {
 }
 
 impl WebServer {
-    pub fn new(tracker: Arc<tracker::TorrentTracker>) -> WebServer {
+    fn get_access_tokens(cfg: &config::HTTPConfig, tokens: &mut HashMap<String, String>) {
+        for (user, token) in cfg.get_access_tokens().iter() {
+            tokens.insert(token.clone(), user.clone());
+        }
+    }
+
+    pub fn new(tracker: Arc<tracker::TorrentTracker>, cfg: Arc<config::Configuration>) -> WebServer {
+        let cfg_cp = cfg.clone();
+
         let server = actix_web::server::HttpServer::new(move || {
-            actix_web::App::<UdptState>::with_state(UdptState::new(tracker.clone(), HashMap::new()))
+            let mut access_tokens = HashMap::new();
+
+            if let Some(http_cfg) = cfg_cp.get_http_config() {
+                Self::get_access_tokens(http_cfg, &mut access_tokens);
+            }
+
+            let state = UdptState::new(tracker.clone(), access_tokens);
+
+            actix_web::App::<UdptState>::with_state(state)
                 .middleware(UdptMiddleware)
                 .resource("/t", |r| r.f(Self::view_torrent_list))
                 .scope(r"/t/{info_hash:[\dA-Fa-f]{40,40}}", |scope| {
@@ -141,14 +158,20 @@ impl WebServer {
                 .resource("/", |r| r.method(actix_web::http::Method::GET).f(Self::view_root))
         });
 
-        match server.bind("0.0.0.0:1212") {
-            Ok(v) => {
-                v.run();
-            },
-            Err(_) => {
-                eprintln!("failed to bind server");
+        if let Some(http_cfg) = cfg.get_http_config() {
+            let bind_addr = http_cfg.get_address();
+            match server.bind(bind_addr) {
+                Ok(v) => {
+                    v.run();
+                },
+                Err(_) => {
+                    eprintln!("failed to bind server");
+                }
             }
         }
+        else {
+            unreachable!();
+        }
 
         WebServer{}
     }
diff --git a/udpt.toml b/udpt.toml
new file mode 100644
index 0000000..572ee6e
--- /dev/null
+++ b/udpt.toml
@@ -0,0 +1,11 @@
+mode = "dynamic"
+
+[udp]
+announce_interval = 120         # Two minutes
+bind_address = "0.0.0.0:1212"
+
+[http]
+bind_address = "127.0.0.1:1212"
+
+[http.access_tokens]
+naim = "MyAccessToken"