Add stats

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

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "zlnk"
-version = "0.3.3"
+version = "0.4.0"
 authors = ["Jonas Kuche <Jonas.kuche@gmail.com>"]
 
 [dependencies]
@@ -12,3 +12,8 @@ r2d2_redis = "0.7.0"
 dotenv = "0.11.0"
 regex = "0.2.10"
 rand = "0.4.2"
+url = "1.7.0"
+woothee = "0.8.0"
+maxminddb = "0.9.0"
+serde_json = "1.0.13"
+

src/env_loader.rs

@@ -8,7 +8,9 @@ pub struct Env {
     pub url_regex: Regex,
     pub short_length: usize,
     pub short_alphabet: String,
-    pub bad_request_message: String
+    pub bad_request_message: String,
+    pub mmdb_path: String,
+    pub trust_proxy: bool
 }
 
 const URL_REGEX: &'static str = r"^(https?://)?([\da-z\.-]+)\.([a-z\.]{2,6})([/\w \.-]*)*/?$";
@@ -20,6 +22,8 @@ pub fn init() -> Env {
         url_regex: Regex::new(&env::var("URL_REGEX").unwrap_or(URL_REGEX.to_string())).unwrap(),
         short_length: usize::from_str(&env::var("SHORT_LENGTH").unwrap_or("5".to_string())).unwrap(),
         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())
+        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()
     }
 }

src/geo_locate_ip.rs

@@ -0,0 +1,27 @@
+use maxminddb::{geoip2, Reader};
+use std::net::IpAddr;
+use std::str::FromStr;
+
+pub struct GeoLocateIP {
+    reader: Reader
+}
+
+impl GeoLocateIP {
+    pub fn new(path: String) -> Self {
+        let reader = Reader::open(&path).unwrap();
+        GeoLocateIP {reader: reader}
+    }
+    pub fn locate(&self, ip_string: String) -> Option<String> {
+        let ip: IpAddr = FromStr::from_str(&ip_string).unwrap();
+        let country = self.reader.lookup(ip);
+        match country {
+            Ok(result) => {
+                let country: geoip2::Country = result;
+                Some(country.country.unwrap().iso_code.unwrap())
+            }
+            Err(_err) => {
+                None
+            }
+        }
+    }
+}

src/main.rs

@@ -7,9 +7,16 @@ extern crate rand;
 extern crate r2d2;
 extern crate r2d2_redis;
 extern crate redis;
+extern crate url;
+extern crate woothee;
+extern crate maxminddb;
+#[macro_use]
+extern crate serde_json;
 
 mod env_loader;
 mod shortener;
+mod geo_locate_ip;
+mod stats;
 #[cfg(test)]
 mod tests;
 
@@ -23,6 +30,8 @@ use rocket::State;
 use r2d2::Pool;
 use r2d2_redis::RedisConnectionManager;
 use env_loader::Env;
+use geo_locate_ip::GeoLocateIP;
+use stats::Stats;
 use shortener::{short, long};
 
 macro_rules! static_file {
@@ -54,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>>) -> Option<Redirect> {
+fn longen(short: String, pool: State<Pool<RedisConnectionManager>>, stats: Stats) -> Option<Redirect> {
     let connection = &pool.get().unwrap();
-    let long = long(short, connection);
+    let long = long(short, connection, stats);
     match long {
         Some(long) => {
             Some(Redirect::to(&long))
@@ -67,6 +76,20 @@ fn longen(short: String, pool: State<Pool<RedisConnectionManager>>) -> Option<Re
     }
 }
 
+#[get("/<short>/stats")]
+fn stats(short: String, pool: State<Pool<RedisConnectionManager>>) -> Option<Content<String>> {
+    let connection = &pool.get().unwrap();
+    let stats = Stats::stats_as_json(short, connection);
+    match stats {
+        Ok(stats) => {
+            Some(Content(ContentType::JSON, stats.to_string()))
+        }
+        Err(_stats) => {
+            None
+        }
+    }
+}
+
 #[error(404)]
 fn not_found() -> Content<Vec<u8>> {
     Content(ContentType::HTML, include_bytes!("../static/404.html").to_vec())
@@ -80,12 +103,13 @@ fn bad_request(request: &Request) -> String {
 
 fn main() {
     let env = env_loader::init();
+    let geo_locate_ip = GeoLocateIP::new(env.mmdb_path.clone());
     let manager = RedisConnectionManager::new(env.redis_url.as_str()).unwrap();
     let pool = r2d2::Pool::builder()
         .build(manager)
         .unwrap();
     rocket::ignite()
-        .mount("/", routes![index, js, jquery, bootstrap_css, bootstrap_js, shorten, longen])
-        .catch(errors![not_found, bad_request]).manage(env).manage(pool)
+        .mount("/", routes![index, js, jquery, bootstrap_css, bootstrap_js, shorten, longen, stats])
+        .catch(errors![not_found, bad_request]).manage(env).manage(geo_locate_ip).manage(pool)
         .launch();
 }