More cleaning! And more

Added ability to add custom short
Added function to check if short is valid
Moved a lot of stuff around to make everything seem cleaner

1 files changed, 132 insertions, 75 deletions

diff --git a/src/main.rs b/src/main.rs
index 3af9e48..380cc2f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,99 +1,104 @@
+/* purl-rs
+ *
+ * Author   : Dennis Eriksen <d@ennis.no>
+ * Project  : purlrs
+ * File     : src/main.rs
+ * Created  : 2023-06-27
+ *
+ */
+
+// standard imports
+use std::collections::HashMap;
 use std::io::Write;
-use dumb_cgi::{Request, EmptyResponse, Query};
-use postgres::{Client, NoTls};
 use std::process::exit;
-use regex::Regex;
+
+// imports from external crates
 use dotenv;
-use url::Url;
+use postgres::{Client, NoTls};
+use dumb_cgi::{Request, EmptyResponse, Query};
 use rand::{thread_rng, Rng, distributions::Alphanumeric};
-
+use regex::Regex;
+use url::Url;
 
 
 // Do the dirty
 fn main() {
 
-    // short = ID for shortened url
-    // url = url to be shortened / that has been shortened
-    // shorturl = full shortened url
-    // shortprefix = part of url that comes before short id
-    // docuri = DOCUMENT_URI from http request. This is everything after the domain.
-    //          Example: "/hey" in "https://example.com/hey"
-
+    // Get variables from dotenv
+    let dburl:&str              = &dotenv::var("DATABASE_URL").unwrap();
+    let create_uri:&str         = &dotenv::var("CREATE_URI").unwrap_or("/create".to_string());
+    let form_uri:&str           = &dotenv::var("FORM_URI").unwrap_or("/form".to_string());
+    let short_uri_prefix:&str   = &dotenv::var("SHORT_URI_PREFIX").unwrap_or("/".to_string());
 
     // Connect to db
-    let dburl:String = dotenv::var("DATABASE_URL").unwrap();
-    let mut db = Client::connect(&dburl, NoTls).unwrap();
+    let mut db = Client::connect(dburl, NoTls).unwrap();
     // TODO: Close connection when done.
 
-    //
-    let shortprefix = dotenv::var("SHORTPREFIX").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");
+    // Get some variable from the current request
+    let scheme:&str         = req.var("REQUEST_SCHEME").unwrap_or("https");
+    let server_name:&str    = req.var("SERVER_NAME").unwrap_or("example.com");
+    let docuri:&str         = req.var("DOCUMENT_URI").unwrap_or("");
 
-    if action == "form" {
+
+    //
+    // Form
+    //
+    if docuri == form_uri {
         print_form();
     }
 
-    else if action == "redirect" {
 
-        // get DOCUMENT_URI
-        let docuri:&str = req.var("DOCUMENT_URI").unwrap_or("");
-        if docuri == "" {
-            respond(400, "No short provided");
-        }
+    //
+    // Create
+    //
+    else if docuri == create_uri {
 
-        // Trim / and + from start
-        let docuri = docuri.trim_start_matches(&['/', '+']);
+        // temporary empty map in case there is something wrong with the query
+        let empty:&HashMap<String, String> = &HashMap::new();
 
-        // Make SURE docuri is a valid short
-        let shortregex = Regex::new(r"^[a-zA-Z0-9_-]+$").unwrap();
-        if !shortregex.is_match(docuri) {
-            respond(400, "Not a valid short");
-        }
+        let query = match req.query() {
+            Query::Some(map) => map,
+            _ => empty, // empty map used here, if no map is found, or there is an error
+        };
 
-        // Fetch URL from postgres, and redirect or return 404
-        let sql = "SELECT url FROM shorts WHERE short = $1 LIMIT 1";
-        match db.query_opt(sql, &[&docuri]).unwrap() {
-            None => respond(404, ""),
-            Some(row) => {
-                let sql = "UPDATE shorts SET count = count + 1, last_visited = now() WHERE short = $1;";
-                db.execute(sql, &[&docuri]).unwrap();
-                respond(301, row.get("url"));
-            },
+        // Check for url in query. It is obligatory.
+        if ! query.contains_key("url") {
+            respond(400, "Error, no url in query");
         }
-    }
 
-    else if action == "create" {
-        let url:&str = match req.query() {
-            Query::None => "Error, no url in query",
-            Query::Err(_e) => "Error reading query string",
-            Query::Some(map) => {
-                if !map.contains_key("url") {
-                    "Error, no url in query"
-                } else {
-                    &map["url"]
-                }
-            },
-        };
-        if url.starts_with("Error") {
-            respond(400, url);
-        } else if !Url::parse(url).is_ok() {
+        // Get url from query
+        let url:&str = &query["url"];
+
+        // Get short from query. Use supplied short if it exists, else generate random.
+        // Use mut String, since we will have to change it if it already exists
+        let mut short:String = if query.contains_key("short") {
+                 query["short"].to_string()
+            } else {
+                gen_short()
+            };
+
+        // Get user from proxy, if user is logged in. Else use "nobody".
+        let user:&str = req.var("REMOTE_USER").unwrap_or("nobody");
+
+        // Check that url and short is valid
+        if ! Url::parse(url).is_ok() {
             respond(400, "Invalid url");
+        } else if ! check_short(&short) {
+            respond(400, "Invalid short");
         }
 
-        let user:&str = req.var("REMOTE_USER").unwrap_or("none");
-
-        let mut short:String = gen_short();
 
+        // Check if short already exists.
+        // If it does exist, set new random short. Try this five times. If we can not find a unique
+        // random short i five tries, abort.
         for i in 1..5 {
             let sql = "SELECT url FROM shorts WHERE short = $1 LIMIT 1";
             match db.query_opt(sql, &[&short]).unwrap() {
-                Some(_row) => short = gen_short(), // If a row was returned, the short was not unique. Continue loop
-                None => break,   // If nothing was returned, the short IS unique. Break out of loop
+                Some(_row)  => short = gen_short(), // If a row was returned, the short was not unique. Continue loop
+                None        => break,   // If nothing was returned, the short IS unique. Break out of loop
             }
 
             // Throw error if we couldn't create a unique short in fire tries
@@ -102,21 +107,50 @@ fn main() {
 
         let sql = "INSERT INTO shorts (url, short, created_by) VALUES ($1, $2, $3);";
         match db.execute(sql, &[&url, &short, &user]) {
-            Ok(_v) => respond(200, &format!("{}{}", shortprefix, short)),
+            Ok(_v)  => respond(200, &format!("{}://{}{}{}", scheme, server_name, short_uri_prefix, short)),
             Err(_e) => respond(500, "Could not save shortened url to database"),
         };
         exit(0);
 
     }
 
+
+    //
+    // Redirect
+    //
+    else if docuri.starts_with(short_uri_prefix) {
+
+        // Trim / and + from start
+        let short = docuri.trim_start_matches(&['/', '+']);
+
+        // Make sure short is valid
+        if ! check_short(&short) { respond(400, "Invalid short"); }
+
+        // Fetch URL from postgres and redirect, or return 404
+        let sql = "SELECT url FROM shorts WHERE short = $1 LIMIT 1";
+        match db.query_opt(sql, &[&short]).unwrap() {
+            Some(row) => {
+                let sql = "UPDATE shorts SET count = count + 1, last_visited = now() WHERE short = $1;";
+                db.execute(sql, &[&short]).unwrap();
+                respond(301, row.get("url"));
+            },
+            None => respond(404, ""),
+        }
+    }
+
+
+    //
+    // Else? Oops.
+    //
     else {
-        respond(400, "");
+        respond(500, "Reached end of script. There is probably a config-error somewhere.");
     }
 }
 
 
-
+//
 // send cgi-response
+//
 fn respond(code:u16, body:&str) {
     // initiate response
     let mut r = EmptyResponse::new(code)
@@ -127,26 +161,27 @@ fn respond(code:u16, body:&str) {
         r.add_header("Location", body);
     }
 
-    // if status >= 300, we want to include status text in the response
+    // if status >= 300, we want to include http-status in the response
     if code >= 300 {
-        let mut toptext = format!("{} {}\n", code, status(code));
-        if !body.is_empty() { toptext.push_str("\n"); }
+        let mut toptext:String = format!("{} {}\n", code, status(code));
+        if ! body.is_empty() { toptext.push_str("\n"); } // Add extra linebreak before error-text
         write!(&mut r, "{toptext}").unwrap();
     }
 
-    if body != "" {
-        // write body
+    // write body, unless there is no body
+    if ! body.is_empty() {
         write!(&mut r, "{}\n", body).unwrap();
     }
 
-    // respond
+    // respond and exit
     r.respond().unwrap();
     exit(0);
 }
 
 
-
+//
 // HTTP status codes
+//
 fn status<'a>(code:u16) -> &'a str {
     // I've only implemented statuscodes I *might* use
     return match code {
@@ -177,8 +212,9 @@ fn status<'a>(code:u16) -> &'a str {
 }
 
 
-
+//
 // Print form
+//
 fn print_form() {
     let body = r#"<html>
 <head>
@@ -200,8 +236,9 @@ fn print_form() {
 }
 
 
-
+//
 // Generate random short
+//
 fn gen_short() -> String {
     let rand_string: String = thread_rng()
         .sample_iter(&Alphanumeric)
@@ -211,3 +248,23 @@ fn gen_short() -> String {
 
     return rand_string;
 }
+
+
+//
+// Check if short is valid
+//
+fn check_short(short:&str) -> bool {
+    // Set regex for valid shorts
+    let shortregex = Regex::new(r"^[a-zA-Z0-9_-]+$").unwrap();
+
+    if ! shortregex.is_match(short) {
+        return false; // short contains invalid characters
+    } else if short.chars().count() > 128 {
+        return false; // short is too long
+    }
+
+    return true;
+}
+
+
+// end of file