PURL

//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:String, short: &str) -> 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 { 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 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<()> { // Get dburl from .env-file let dburl = dotenv::var("DATABASE_URL").unwrap(); let prefix = dotenv::var("PREFIX").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) { 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 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"); } exit(0); }, Query::Err(_e) => error(400,"Bad Request\n\nError reading query string"), } }, _ => error(400, "Bad Request"), } Ok(()) }