1 files changed, 34 insertions, 10 deletions

diff --git a/src/main.rs b/src/main.rs
index a6bca8a..c041886 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,34 +4,58 @@ use std::net::SocketAddrV4;
 use igd::aio::search_gateway;
 use igd::PortMappingProtocol;
 
-struct DiplonatConfig<'a> {
-   private_ip: &'a str,
-   refresh_time: u16,
-   expiration_time: u16
+struct DiplonatConfig {
+   private_ip: String,
+   refresh_time: u32,
+   expiration_time: u32
 }
 
-fn fetch_configuration {
+fn fetch_configuration() -> Result<DiplonatConfig, String> {
   let env_private_ip = "DIPLONAT_PRIVATE_IP";
   let private_ip = match env::var(env_private_ip) {
       Ok(val) => val,
-      Err(e) => return println!("Unable to fetch {} environment variable: {}", env_private_ip, e),
+      Err(e) => return Err(format!("unable to fetch environment variable {}: {}", env_private_ip, e)),
   };
 
+  let env_refresh_time = "DIPLONAT_REFRESH_TIME";
+  let refresh_time: u32 = match env::var(env_refresh_time) {
+      Ok(val) => val.parse().unwrap(),
+      Err(e) => return Err(format!("unable to fetch environment variable {}: {}", env_refresh_time,e))
+  };
+  
+  let env_expiration_time = "DIPLONAT_EXPIRATION_TIME";
+  let expiration_time: u32 = match env::var(env_expiration_time) {
+      Ok(val) => val.parse().unwrap(),
+      Err(e) => return Err(format!("unable to fetch environment variable {}: {}", env_expiration_time,e))
+  };
+
+  if refresh_time * 2 > expiration_time {
+    return Err(format!("Expiration time (currently: {}s) must be twice bigger than refresh time (currently: {}s)", expiration_time, refresh_time))
+  }
+
+  let config = DiplonatConfig { private_ip: private_ip, refresh_time: refresh_time, expiration_time: expiration_time };
+  println!("\tPrivate IP address: {}", config.private_ip);
+  println!("\tRefresh time: {} seconds", config.refresh_time);
+  println!("\tExpiration time: {} seconds", config.expiration_time);
+  return Ok(config);
 }
 
 #[tokio::main]
 async fn main() {
+    let config = match fetch_configuration() {
+        Ok(val) => val,
+        Err(e) => return println!("unable to build configuration: {}", e),
+    };
 
-     println!("Private IP address: {}", private_ip);
 
     let gateway = match search_gateway(Default::default()).await {
         Ok(g) => g,
         Err(err) => return println!("Faild to find IGD: {}", err),
     };
 
-    let service = format!("{}:{}", private_ip, 1234);
-    let service: SocketAddrV4 = private_ip.parse().expect("Invalid socket address");
-    match gateway.add_port(PortMappingProtocol::TCP, 1234, service, 120, "diplonat").await {
+    let service = format!("{}:{}", config.private_ip, 1234);
+    let service: SocketAddrV4 = service.parse().expect("Invalid socket address");
+    match gateway.add_port(PortMappingProtocol::TCP, 1234, service, config.expiration_time, "diplonat").await {
         Ok(_) => (),
         Err(e) => return println!("Unable to insert port 1234: {}", e),
     };