diff options
| -rw-r--r-- | src/main.rs | 207 |
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 |