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, distributions::Alphanumeric}; // 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" // Connect to db let dburl:String = dotenv::var("DATABASE_URL").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"); if action == "form" { 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"); } // 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) { respond(400, "Not a valid 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, &[&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")); }, } } 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() { respond(400, "Invalid url"); } let user:&str = req.var("REMOTE_USER").unwrap_or("none"); let mut short:String = gen_short(); 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 } // Throw error if we couldn't create a unique short in fire tries if i == 5 { respond(500, "Could not find unique short"); } } 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)), Err(_e) => respond(500, "Could not save shortened url to database"), }; exit(0); } else { respond(400, ""); } } // send cgi-response fn respond(code:u16, body:&str) { // initiate response let mut r = EmptyResponse::new(code) .with_content_type("text/plain"); // is 300 <= status < 400, and body is a url, perform redirection if (300..399).contains(&code) && Url::parse(body).is_ok() { r.add_header("Location", body); } // if status >= 300, we want to include status text in the response if code >= 300 { let mut toptext = format!("{} {}\n", code, status(code)); if !body.is_empty() { toptext.push_str("\n"); } write!(&mut r, "{toptext}").unwrap(); } if body != "" { // write body write!(&mut r, "{}\n", body).unwrap(); } // respond 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 { 200 => "OK", 201 => "Created", 202 => "Accepted", 204 => "No Content", 301 => "Moved Permanently", 302 => "Found", 304 => "Not Modified", 307 => "Temporary Redirect", 308 => "Permanent Redirect", 400 => "Bad Request", 401 => "Unauthorized", 402 => "Payment Required", 403 => "Forbidden", 404 => "Not Found", 405 => "Method Not Allowed", 406 => "Not Acceptableo", 410 => "Gone", 414 => "URI Too Long", 500 => "Internal Server Error", 501 => "Not Implemented", 503 => "Service Unavailable", 508 => "Loop Detected", _ => "Not Found", }; } // Print form fn print_form() { let body = r#" PURL
Username: $user
URL to shorten:
Custom short:
"#; respond(200, body); } // 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; }