From 12ed0e7c03794a9658428d0aacdffb7d642189c7 Mon Sep 17 00:00:00 2001
From: int 80h <int@80h.dev>
Date: Fri, 22 May 2020 15:12:28 -0400
Subject: [PATCH] Timeout for request and cgi

---
 Cargo.toml  |  2 +-
 src/cgi.rs  | 25 ++++++++++++++++++++++---
 src/main.rs | 22 +++++++++++++---------
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index d3335d0..4e096a0 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,7 +7,7 @@ edition = "2018"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-tokio = { version = "0.2", features = [ "net", "io-util", "rt-threaded" ] }
+tokio = { version = "0.2", features = [ "time", "fs", "process", "net", "io-util", "rt-threaded" ] }
 openssl = "0.10"
 tokio-openssl = "0.4"
 futures-util = "0.3"
diff --git a/src/cgi.rs b/src/cgi.rs
index 31e795a..c7a5939 100644
--- a/src/cgi.rs
+++ b/src/cgi.rs
@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 use std::io;
 use std::path::PathBuf;
-use std::process::Command;
+use tokio::process::Command;
 use url::Url;
 
 use crate::config;
@@ -44,8 +44,27 @@ pub async fn cgi(
     let cmd = Command::new(path.to_str().unwrap())
         .env_clear()
         .envs(&envs)
-        .output()
-        .unwrap();
+        .output();
+    
+    let cmd = match tokio::time::timeout(tokio::time::Duration::from_secs(5), cmd).await {
+        Ok(c) => {
+            match c {
+                Ok(cc) => cc,
+
+                Err(_) => {
+                    logger::logger(con.peer_addr, Status::CGIError, url.as_str());
+                    con.send_status(Status::CGIError, None).await?;
+                    return Ok(());
+                },
+            }
+        },
+        Err(_) => {
+            logger::logger(con.peer_addr, Status::CGIError, url.as_str());
+            con.send_status(Status::CGIError, None).await?;
+            return Ok(());
+        },
+    };
+
     if !cmd.status.success() {
         logger::logger(con.peer_addr, Status::CGIError, url.as_str());
         con.send_status(Status::CGIError, None).await?;
diff --git a/src/main.rs b/src/main.rs
index 613ea14..9ba04a9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -66,10 +66,10 @@ async fn get_binary(mut con: conn::Connection, path: PathBuf, meta: String) -> i
     Ok(())
 }
 
-fn get_content(path: PathBuf, u: url::Url) -> Result<String, io::Error> {
-    let meta = fs::metadata(&path).expect("Unable to read metadata");
+async fn get_content(path: PathBuf, u: url::Url) -> Result<String, io::Error> {
+    let meta = tokio::fs::metadata(&path).await?;
     if meta.is_file() {
-        return Ok(std::fs::read_to_string(path).expect("Unable to read file"));
+        return Ok(tokio::fs::read_to_string(path).await?);
     }
 
     let mut list = String::from("# Directory Listing\r\n\r\n");
@@ -104,7 +104,11 @@ async fn handle_connection(
         None => "index.gemini".to_string(),
     };
     let mut buffer = [0; 1024];
-    con.stream.read(&mut buffer).await?;
+    if let Err(_) = tokio::time::timeout(tokio::time::Duration::from_secs(5), con.stream.read(&mut buffer)).await {
+        logger::logger(con.peer_addr, Status::BadRequest, "");
+        con.send_status(Status::BadRequest, None).await?;
+        return Ok(());
+    }
     let mut request = match String::from_utf8(buffer[..].to_vec()) {
         Ok(request) => request,
         Err(_) => {
@@ -201,7 +205,7 @@ async fn handle_connection(
         return Ok(());
     }
 
-    let mut meta = fs::metadata(&path).expect("Unable to read metadata");
+    let mut meta = tokio::fs::metadata(&path).await?;
     let mut perm = meta.permissions();
 
     // TODO fix me
@@ -218,19 +222,18 @@ async fn handle_connection(
         }
         if path.join(&index).exists() {
             path.push(index);
-            meta = fs::metadata(&path).expect("Unable to read metadata");
+            meta = tokio::fs::metadata(&path).await?;
             perm = meta.permissions();
             if perm.mode() & 0o0444 != 0o444 {
                 let mut p = path.clone();
                 p.pop();
                 path.push(format!("{}/", p.display()));
-                meta = fs::metadata(&path).expect("Unable to read metadata");
+                meta = tokio::fs::metadata(&path).await?;
                 perm = meta.permissions();
             }
         }
     }
 
-    // TODO add timeout
     match &srv.server.cgi {
         Some(c) => {
             if c.trim_end_matches("/") == path.parent().unwrap().to_str().unwrap() {
@@ -259,7 +262,7 @@ async fn handle_connection(
         get_binary(con, path, mime).await?;
         return Ok(());
     }
-    let content = get_content(path, url)?;
+    let content = get_content(path, url).await?;
     con.send_body(status::Status::Success, Some(&mime), Some(content))
         .await?;
     logger::logger(con.peer_addr, Status::Success, &request);
@@ -293,6 +296,7 @@ fn main() -> io::Result<()> {
     let mut runtime = runtime::Builder::new()
         .threaded_scheduler()
         .enable_io()
+        .enable_time()
         .build()?;
 
     let handle = runtime.handle().clone();