//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<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 {
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<u64, postgres::Error> {
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#"<html>
<head>
<title>PURL</title>
</head>
<body>
<form method="post" action="/purl.cgi" enctype="multipart/form-data" name="main">
Username: $user<br>
URL to shorten: <input type="text" name="url" size="50"><br>
Custom short: <input type="text" name="short"><br>
<input type="hidden" name="form" value="html">
<input type="submit" value="Submit">
</form>
</body>
</html>
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(())
}
