version 0.1

.gitignore

@@ -0,0 +1,3 @@
+/target
+/secret.txt
+/postgres_password.txt

Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "shortener"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+rocket = "0.5.1"
+tokio = {"version" = "1.40.0", "features" = ["rt"]}
+rand = "0.8.5"
+tokio-postgres = "0.7.12"

src/db.rs

@@ -0,0 +1,25 @@
+use rocket::futures::lock::Mutex;
+use tokio_postgres::{Client, Error, Statement};
+
+use crate::GlobalState;
+
+pub async fn get_link(state: &Mutex<GlobalState>, link: &str) -> Result<Option<String>, Error> {
+    let lock = state.lock().await;
+    Ok(lock.db_client.query_opt(&lock.stmt_get_link, &[&link]).await?.map(|row| row.get(0)))
+}
+
+pub async fn add_link(state: &Mutex<GlobalState>, link: &str, url: &str) -> Result<(), Error> {
+    let lock = state.lock().await;
+    lock.db_client.execute(&lock.stmt_add_link, &[&link, &url]).await?;
+    Ok(())
+}
+
+pub async fn prepare_statements(db: &Client) -> Result<(Statement, Statement), Error> {
+    Ok((db.prepare("SELECT url FROM links WHERE id = $1").await?,
+        db.prepare("INSERT INTO links (id, url) VALUES ($1, $2)").await?))
+}
+
+pub async fn prepare_tables(db: &Client) -> Result<(), Error> {
+    db.execute("CREATE TABLE IF NOT EXISTS links (id TEXT PRIMARY KEY, url TEXT NOT NULL)", &[]).await?;
+    Ok(())
+}

src/linkgen.rs

@@ -0,0 +1,14 @@
+use rand::{distributions::Uniform, Rng};
+
+pub fn generate_link_id(length: u32) -> String {
+    let mut rng = rand::thread_rng();
+    let alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_";
+    let mut link = String::with_capacity(length as usize);
+    for _ in 0..length {
+        link.push(alphabet.chars().nth(rng.sample(Uniform::new(0u8, 64u8)) as usize).unwrap());
+    }
+    if link == "create" || link == "p" {
+        return generate_link_id(length);
+    }
+    return link;
+}

src/main.rs

@@ -0,0 +1,106 @@
+use rocket::{futures::lock::Mutex, get, http::Status, response::Redirect, routes, State};
+use tokio_postgres::{Client, NoTls, Statement};
+
+mod db;
+mod linkgen;
+
+const DEFAULT_LINK_LENGTH: u32 = 6;
+const LINK_PREFIX: &str = "https://slavasil.ru/";
+
+pub struct GlobalState {
+    db_client: Client,
+
+    stmt_get_link: Statement,
+    stmt_add_link: Statement,
+}
+#[tokio::main]
+async fn main() {
+    let postgres_config = "host=localhost user=shortener password=".to_owned() + include_str!("../postgres_password.txt").trim();
+    let (client, conn) = tokio_postgres::connect(&postgres_config, NoTls).await.unwrap();
+    tokio::spawn(async move {
+        if let Err(e) = conn.await {
+            eprintln!("postgresql error: {}", e);
+        }
+    });
+    db::prepare_tables(&client).await.unwrap();
+    let statements = db::prepare_statements(&client).await.unwrap();
+    let state = GlobalState { db_client: client, stmt_get_link: statements.0, stmt_add_link: statements.1 };
+    rocket::build()
+        .mount("/", routes![create, go_to_link])
+        .manage(Mutex::new(state))
+        .launch()
+        .await.unwrap();
+}
+
+#[get("/create?<url>&<secret>&<length>&<link>")]
+async fn create(state: &State<Mutex<GlobalState>>, url: &str, secret: Option<&str>, length: Option<u32>, link: Option<&str>) -> (Status, String) {
+    let mut allow_secret_options = false;
+    if let Some(secret) = secret {
+        if secret == include_str!("../secret.txt").trim() {
+            allow_secret_options = true;
+        }
+    }
+    let mut actual_link = None;
+    let mut custom_link = false;
+    if allow_secret_options {
+        if let Some(link) = link {
+            actual_link = Some(link.to_owned());
+            custom_link = true;
+        }
+    }
+    let actual_len = if allow_secret_options && length.is_some() { length.unwrap() } else { DEFAULT_LINK_LENGTH };
+    if actual_len < 1 || actual_len > 64 {
+        return (Status::BadRequest, "invalid length".to_owned());
+    }
+    
+    loop {
+        let link_id = if !custom_link {
+            linkgen::generate_link_id(actual_len)
+        } else {
+            actual_link.clone().unwrap()
+        };
+
+        match db::get_link(&state, &link_id).await {
+            Ok(url) => match url {
+                Some(_) => {
+                    if custom_link {
+                        return (Status::BadRequest, "link already exists".to_owned());
+                    }
+                },
+                None => {
+                    actual_link = Some(link_id.clone());
+                    break;
+                }
+            },
+            Err(e) => {
+                eprintln!("get link error: {}", e);
+                return (Status::InternalServerError, "server error".to_owned());
+            }
+        }
+    }
+
+    match db::add_link(&state, actual_link.as_ref().unwrap(), url).await {
+        Ok(_) => (),
+        Err(e) => {
+            eprintln!("add link error:  {}", e);
+            return (Status::InternalServerError, "server error".to_owned());
+        }
+    }
+
+    let prefixed_link = LINK_PREFIX.to_owned() + &actual_link.unwrap();
+    (Status::Ok, prefixed_link)
+}
+
+#[get("/<link>")]
+pub async fn go_to_link(state: &State<Mutex<GlobalState>>, link: &str) -> Result<Option<Redirect>, Status> {
+    match db::get_link(&state, link).await {
+        Ok(url) => match url {
+            Some(url) => Ok(Some(Redirect::to(url))),
+            None => Ok(None),
+        },
+        Err(e) => {
+            eprintln!("get link error: {}", e);
+            Err(Status::InternalServerError)
+        }
+    }
+}