path: root/src/main.rs

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#"<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>
"#;

    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;
}