//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:&str, short:&str, update:bool) -> Result { // 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 && update { 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 { 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#" PURL
Username: $user
URL to shorten:
Custom short:
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<()> { // 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 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"); 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, true) { 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 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); }, Query::Err(_e) => error(400,"Bad Request\n\nError reading query string"), } }, _ => error(400, "Bad Request"), } Ok(()) }