//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<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 && 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<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<()> {
// 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(())
}
