this now actually works, and can be used. TODO: cleanup and make it pretty

1 files changed, 35 insertions, 11 deletions

diff --git a/src/main.rs b/src/main.rs
index 7a60081..439f8fd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -10,13 +10,13 @@ use rand::distributions::Alphanumeric;
 
 
 // Get URL from database
-fn get_url(dburl:String, short: &str) -> Result<String, postgres::Error> {
+fn get_url(dburl:&str, short:&str, update:bool) -> 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 {
+    if row.len() == 1 && update {
         client.execute("UPDATE shorts SET count = count + 1, last_visited = now() WHERE short = $1;", &[&short])?;
     }
     client.close()?;
@@ -95,12 +95,20 @@ fn gen_short() -> String {
 }
 
 
-
 // Do the dirty
 fn main() -> std::io::Result<()> {
+
+    // 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 dburl from .env-file
     let dburl = dotenv::var("DATABASE_URL").unwrap();
-    let prefix = dotenv::var("PREFIX").unwrap();
+    let shortprefix = dotenv::var("SHORTPREFIX").unwrap();
 
     // Gather all request data from the environment and stdin.
     let req = Request::new().unwrap();
@@ -124,7 +132,7 @@ fn main() -> std::io::Result<()> {
             }
 
             // Fetch URL from postgres, and redirect or return 404
-            match get_url(dburl, docuri) {
+            match get_url(&dburl, docuri, true) {
                 Ok(url) => redirect(301, &url),
                 Err(_e) => error(404, "Not Found"),
             };
@@ -143,13 +151,29 @@ fn main() -> std::io::Result<()> {
                     }
 
                     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");
+                    let mut short:String = gen_short();
+
+                    for i in 1..5 {
+                        match get_url(&dburl, &short, false) {
+                            Ok(_url) => short = gen_short(), // If Ok, then short is already in
+                                                             // use. Try a new one.
+                            Err(_e) => break, // we assume that we found a unique short if get_url
+                                              // returns an error
+                        }
+
+                        // Throw error if we couldn't create a unique short in fire tries
+                        if i == 5 { error(500, "Could not find unique short"); }
                     }
+
+                    match insert_short(dburl, url, &short, user) {
+                        Ok(_v) => error(200, &format!("{}{}", shortprefix, short)),
+                        Err(_e) => error(400, "looool"),
+                    };
+                    //if n == 1 {
+                    //    error(200, &res);
+                    //} else {
+                    //    error(400, "Bad Request\n\nDunno wtf happened");
+                    //}
                     exit(0);
 
                 },