Add disable stats

Signed-off-by: Jonas Kuche <Jonas.kuche@gmail.com>

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "zlnk"
-version = "0.4.1"
+version = "0.4.2"
 authors = ["Jonas Kuche <Jonas.kuche@gmail.com>"]
 
 [dependencies]

src/env_loader.rs

@@ -10,7 +10,8 @@ pub struct Env {
     pub short_alphabet: String,
     pub bad_request_message: String,
     pub mmdb_path: String,
-    pub trust_proxy: bool
+    pub trust_proxy: bool,
+    pub disable_stats: bool
 }
 
 const URL_REGEX: &'static str = r"^(https?://)?([\da-z\.-]+)\.([a-z\.]{2,6})([/\w \.-]*)*/?$";
@@ -24,6 +25,7 @@ pub fn init() -> Env {
         short_alphabet: env::var("SHORT_ALPHABET").unwrap_or("hex".to_string()),
         bad_request_message: env::var("BAD_REQUEST_MESSAGE").unwrap_or("Ungültige URL".to_string()),
         mmdb_path: env::var("MMDB_PATH").unwrap_or("./GeoLite2-Country.mmdb".to_string()),
-        trust_proxy: env::var("TRUST_PROXY").is_ok()
+        trust_proxy: env::var("TRUST_PROXY").is_ok(),
+        disable_stats: env::var("DISABLE_STATS").is_ok()
     }
 }

src/geo_locate_ip.rs

@@ -3,17 +3,29 @@ use std::net::IpAddr;
 use std::str::FromStr;
 
 pub struct GeoLocateIP {
-    reader: Reader
+    reader: Option<Reader>
 }
 
 impl GeoLocateIP {
-    pub fn new(path: String) -> Self {
-        let reader = Reader::open(&path).unwrap();
-        GeoLocateIP {reader: reader}
+    pub fn new(path: String, enabled: bool) -> Self {
+        if enabled {
+            let reader = Reader::open(&path).unwrap();
+            GeoLocateIP {reader: Some(reader)}
+        } else {
+            GeoLocateIP {reader: None}
+        }
     }
     pub fn locate(&self, ip_string: String) -> Option<String> {
         let ip: IpAddr = FromStr::from_str(&ip_string).unwrap();
-        let country = self.reader.lookup(ip);
+        let country;
+        match self.reader {
+            Some(ref reader) => {
+                country = reader.lookup(ip);
+            }
+            None => {
+                panic!("GEOReader not enabled!");
+            }
+        }
         match country {
             Ok(result) => {
                 let country: geoip2::Country = result;

src/main.rs

@@ -63,9 +63,9 @@ fn shorten(long_url: String, env: State<Env>, pool: State<Pool<RedisConnectionMa
 }
 
 #[get("/<short>")]
-fn longen(short: String, pool: State<Pool<RedisConnectionManager>>, stats: Stats) -> Option<Redirect> {
+fn longen(short: String, env: State<Env>, pool: State<Pool<RedisConnectionManager>>, stats: Stats) -> Option<Redirect> {
     let connection = &pool.get().unwrap();
-    let long = long(short, connection, stats);
+    let long = long(short, env.inner(), connection, stats);
     match long {
         Some(long) => {
             Some(Redirect::to(&long))
@@ -77,7 +77,10 @@ fn longen(short: String, pool: State<Pool<RedisConnectionManager>>, stats: Stats
 }
 
 #[get("/<short>/stats")]
-fn stats(short: String, pool: State<Pool<RedisConnectionManager>>) -> Option<Content<String>> {
+fn stats(short: String, env: State<Env>, pool: State<Pool<RedisConnectionManager>>) -> Option<Content<String>> {
+    if env.disable_stats {
+        return None;
+    }
     let connection = &pool.get().unwrap();
     let stats = Stats::stats_as_json(short, connection);
     match stats {
@@ -103,7 +106,7 @@ fn bad_request(request: &Request) -> String {
 
 fn main() {
     let env = env_loader::init();
-    let geo_locate_ip = GeoLocateIP::new(env.mmdb_path.clone());
+    let geo_locate_ip = GeoLocateIP::new(env.mmdb_path.clone(), !env.disable_stats);
     let manager = RedisConnectionManager::new(env.redis_url.as_str()).unwrap();
     let pool = r2d2::Pool::builder()
         .build(manager)

src/shortener.rs

@@ -62,11 +62,13 @@ pub fn short(long_url: String, env: &Env, connection: &r2d2::PooledConnection<r2
     }
 }
 
-pub fn long(short_url: String, connection: &r2d2::PooledConnection<r2d2_redis::RedisConnectionManager>, stats: Stats) -> Option<String> {
+pub fn long(short_url: String, env: &Env, connection: &r2d2::PooledConnection<r2d2_redis::RedisConnectionManager>, stats: Stats) -> Option<String> {
     let long = connection.get(format!("short_{}", &short_url));
     match long {
         Ok(long) => {
-            stats.save(short_url, connection).unwrap();
+            if !env.disable_stats {
+                stats.save(short_url, connection).unwrap();
+            }
             Some(long)
         }
         Err(_err) => {

src/stats.rs

@@ -43,6 +43,15 @@ impl<'a, 'r> FromRequest<'a, 'r> for Stats {
     type Error = ();
 
     fn from_request(request: &'a Request<'r>) -> request::Outcome<Stats, ()> {
+        let env = request.guard::<State<Env>>().unwrap();
+        if env.disable_stats {
+            return Outcome::Success(Stats {
+                refer: "Unknown".to_string(),
+                browser: "Unknown".to_string(),
+                os: "Unknown".to_string(),
+                country: "Unknown".to_string()
+            })
+        }
         let ip = request.guard::<IP>().unwrap();
         let geo_locate_ip = request.guard::<State<GeoLocateIP>>().unwrap();
         let refers: Vec<_> = request.headers().get("Referer").collect();

src/tests.rs

@@ -28,7 +28,7 @@ fn shortener_test() {
     let _:() = redis::cmd("FLUSHDB").query(connection.deref()).unwrap();
     let long_url = "https://zlnk.de".to_string();
     let shorted = short(long_url.clone(), env, connection, None).unwrap();
-    let longed = long(shorted, connection, stats).unwrap();
+    let longed = long(shorted, env, connection, stats).unwrap();
     assert_eq!(longed, long_url);
 }