1 files changed, 165 insertions, 0 deletions
diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..7a60081
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,165 @@
+//use std::io::Write;
+use dumb_cgi::{Request, EmptyResponse, Query};
+use postgres::{Client, NoTls};
+use std::process::exit;
+use regex::Regex;
+use dotenv;
+use url::Url;
+use rand::{thread_rng, Rng};
+use rand::distributions::Alphanumeric;
+
+
+// Get URL from database
+fn get_url(dburl:String, short: &str) -> Result<String, postgres::Error> {
+    // Connect to db
+    let mut client = Client::connect(&dburl, NoTls)?;
+
+    let row = client.query_one("SELECT url FROM shorts WHERE short = $1 LIMIT 1", &[&short])?;
+
+    if row.len() == 1 {
+        client.execute("UPDATE shorts SET count = count + 1, last_visited = now() WHERE short = $1;", &[&short])?;
+    }
+    client.close()?;
+
+    let url: String = row.get("url");
+
+    Ok(url)
+}
+
+fn insert_short(dburl:String, url: &str, short: &str, user: &str) -> Result<u64, postgres::Error> {
+    let mut client = Client::connect(&dburl, NoTls)?;
+
+    let n = client.execute("INSERT INTO shorts (url, short, created_by) VALUES ($1, $2, $3);", &[&url, &short, &user])?;
+    client.close()?;
+    Ok(n)
+}
+// Do the redirect
+// return std::io::Result<()> so we can use ? -
+// https://doc.rust-lang.org/std/result/#the-question-mark-operator-
+fn redirect(code:u16, url:&str) {
+    EmptyResponse::new(code)
+        .with_header("Location", url)
+        .with_content_type("text/plain")
+        .with_body(
+            format!("Redirecting to {}\n", url))
+        .respond().unwrap();
+    exit(0);
+}
+
+fn error(code:u16, msg:&str) {
+    EmptyResponse::new(code)
+        .with_content_type("text/plain")
+        .with_body(
+            format!("{} {}\n", code, msg))
+        .respond().unwrap();
+    exit(0);
+}
+
+// Print form
+fn print_form() {
+    let html = r#"<html>
+<head>
+    <title>PURL</title>
+</head>
+<body>
+    <form method="post" action="/purl.cgi" enctype="multipart/form-data" name="main">
+    Username: $user<br>
+    URL to shorten: <input type="text" name="url" size="50"><br>
+    Custom short: <input type="text" name="short"><br>
+    <input type="hidden" name="form" value="html">
+    <input type="submit" value="Submit">
+    </form>
+
+</body>
+</html>
+HTML
+"#;
+
+    EmptyResponse::new(200)
+        .with_content_type("text/html")
+        .with_body(html)
+        .respond().unwrap();
+    exit(0);
+}
+
+
+// Generate random short
+fn gen_short() -> String {
+    let rand_string: String = thread_rng()
+        .sample_iter(&Alphanumeric)
+        .take(thread_rng().gen_range(2..6))
+        .map(char::from)
+        .collect();
+
+    return rand_string;
+}
+
+
+
+// Do the dirty
+fn main() -> std::io::Result<()> {
+    // Get dburl from .env-file
+    let dburl = dotenv::var("DATABASE_URL").unwrap();
+    let prefix = dotenv::var("PREFIX").unwrap();
+
+    // Gather all request data from the environment and stdin.
+    let req = Request::new().unwrap();
+
+    // Find out what action we're performing
+    let action:&str = req.var("ACTION").unwrap_or("none");
+
+    match action {
+        "form" => print_form(),
+
+        "redirect" => {
+            // get DOCUMENT_URI
+            let docuri:&str = req.var("DOCUMENT_URI").unwrap_or("none");
+            // Trim / and + from start
+            let docuri = docuri.trim_start_matches(&['/', '+']);
+
+            // Make SURE docuri is a valid short
+            let shortregex = Regex::new(r"^[a-zA-Z0-9_-]+$").unwrap();
+            if !shortregex.is_match(docuri) {
+                error(400, "Bad Request");
+            }
+
+            // Fetch URL from postgres, and redirect or return 404
+            match get_url(dburl, docuri) {
+                Ok(url) => redirect(301, &url),
+                Err(_e) => error(404, "Not Found"),
+            };
+        },
+
+        "create" => {
+            match req.query() {
+                Query::None => redirect(303, "/purl.cgi"),
+                Query::Some(map) => {
+                    if map.iter().count() != 1 {
+                        error(400, "Bad Request\n\nIncorrect number of query items");
+                    }
+                    let url:&str = &map["url"];
+                    if !Url::parse(url).is_ok() {
+                        error(400, "Bad Request\n\nInvalid url");
+                    }
+
+                    let user:&str = req.var("REMOTE_USER").unwrap_or("none");
+                    let short:&str = &gen_short();
+                    let res:String = format!("{}{}", prefix, short);
+                    if insert_short(dburl, url, short, user).unwrap() > 0 {
+                        error(200, &res);
+                    } else {
+                        error(400, "Bad Request\n\nDunno wtf happened");
+                    }
+                    exit(0);
+
+                },
+                Query::Err(_e) => error(400,"Bad Request\n\nError reading query string"),
+            }
+        },
+
+        _ => error(400, "Bad Request"),
+    }
+
+
+    Ok(())
+}