format main.rs

Dockerfile

@@ -0,0 +1,10 @@
+FROM rust:latest as build
+WORKDIR /usr/src/app
+COPY Cargo.toml Cargo.lock secret.txt postgres_password.txt .
+COPY src src
+RUN rustup target add x86_64-unknown-linux-musl
+RUN cargo build --target=x86_64-unknown-linux-musl --release
+
+FROM scratch
+COPY --from=build /usr/src/app/target/x86_64-unknown-linux-musl/release/shortener /
+CMD ["/shortener"]

src/db.rs

@@ -1,3 +1,5 @@
+use std::time::{Duration, SystemTime};
+
 use rocket::futures::lock::Mutex;
 use tokio_postgres::{Client, Error, Statement};
 
@@ -5,26 +7,65 @@ use crate::GlobalState;
 
 pub const STATEMENT_GET_LINK: usize = 0;
 pub const STATEMENT_ADD_LINK: usize = 1;
-pub const N_STATEMENTS: usize = 2;
+pub const STATEMENT_ADD_TEMP_LINK: usize = 2;
+pub const STATEMENT_DELETE_EXPIRED: usize = 3;
+pub const N_STATEMENTS: usize = 4;
 
 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.statements[STATEMENT_GET_LINK], &[&link]).await?.map(|row| row.get(0)))
+    Ok(lock
+        .db_client
+        .query_opt(&lock.statements[STATEMENT_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.statements[STATEMENT_ADD_LINK], &[&link, &url]).await?;
+    lock.db_client
+        .execute(&lock.statements[STATEMENT_ADD_LINK], &[&link, &url])
+        .await?;
+    Ok(())
+}
+
+pub async fn add_temporary_link(
+    state: &Mutex<GlobalState>,
+    link: &str,
+    url: &str,
+    ttl: u64,
+) -> Result<(), Error> {
+    let time = SystemTime::now() + Duration::from_secs(ttl);
+    let lock = state.lock().await;
+    lock.db_client
+        .execute(
+            &lock.statements[STATEMENT_ADD_TEMP_LINK],
+            &[&link, &url, &time],
+        )
+        .await?;
+    Ok(())
+}
+
+pub async fn delete_expired_links(state: &Mutex<GlobalState>) -> Result<(), Error> {
+    let lock = state.lock().await;
+    lock.db_client
+        .execute(&lock.statements[STATEMENT_DELETE_EXPIRED], &[])
+        .await?;
     Ok(())
 }
 
 pub async fn prepare_statements(db: &Client) -> Result<[Statement; N_STATEMENTS], Error> {
-    Ok([db.prepare("SELECT url FROM links WHERE id = $1").await?,
-        db.prepare("INSERT INTO links (id, url) VALUES ($1, $2)").await?,
+    Ok([
+        db.prepare("SELECT url FROM links WHERE id = $1").await?,
+        db.prepare("INSERT INTO links (id, url) VALUES ($1, $2)")
+            .await?,
+        db.prepare("INSERT INTO links (id, url, valid_until) VALUES ($1, $2, $3)")
+            .await?,
+        db.prepare("DELETE FROM links WHERE valid_until < NOW()")
+            .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?;
+    db.execute("CREATE TABLE IF NOT EXISTS links (id TEXT PRIMARY KEY, url TEXT NOT NULL, valid_until TIMESTAMP NOT NULL)", &[]).await?;
     Ok(())
-}
+}

src/linkgen.rs

@@ -5,10 +5,15 @@ pub fn generate_link_id(length: u32) -> String {
     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());
+        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

@@ -1,13 +1,23 @@
 use db::N_STATEMENTS;
-use rocket::{futures::lock::Mutex, get, fs::FileServer, http::Status, response::Redirect, response::content::RawHtml, routes, State};
+use rocket::{
+    fs::FileServer, futures::lock::Mutex, get, http::Status, response::content::RawHtml,
+    response::Redirect, routes, State,
+};
+use std::{
+    env,
+    fs::File,
+    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
+    str::FromStr,
+    sync::Arc,
+    time::Duration,
+};
 use tokio_postgres::{Client, NoTls, Statement};
-use std::fs::File;
 
 mod db;
 mod linkgen;
 
 const DEFAULT_LINK_LENGTH: u32 = 6;
-const LINK_PREFIX: &str = "https://slavasil.ru/";
+const LINK_PREFIX: &str = "https://xurl.pw/";
 
 pub struct GlobalState {
     db_client: Client,
@@ -16,8 +26,15 @@ pub struct GlobalState {
 }
 #[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();
+    let postgres_config = "host=localhost user=shortener password=".to_owned()
+        + include_str!("../postgres_password.txt").trim();
+    let (client, conn) = match tokio_postgres::connect(&postgres_config, NoTls).await {
+        Ok((client, conn)) => (client, conn),
+        Err(e) => {
+            eprintln!("Failed to connect to PostgreSQL: {}", e);
+            return;
+        }
+    };
     tokio::spawn(async move {
         if let Err(e) = conn.await {
             eprintln!("postgresql error: {}", e);
@@ -25,19 +42,50 @@ async fn main() {
     });
     db::prepare_tables(&client).await.unwrap();
     let statements = db::prepare_statements(&client).await.unwrap();
+
+    let default_listen_addr = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 3020).into();
+    let listen_addr = env::var("SERVER_ADDR")
+        .ok()
+        .map(|v| SocketAddr::from_str(&v).ok())
+        .flatten()
+        .unwrap_or(default_listen_addr);
     let mut config = rocket::Config::default();
-    config.port = 3020;
-    let state = GlobalState { db_client: client, statements };
+    config.address = listen_addr.ip();
+    config.port = listen_addr.port();
+
+    let state = GlobalState {
+        db_client: client,
+        statements,
+    };
+    let state_mutex = Arc::new(Mutex::new(state));
+    {
+        let state_mutex = state_mutex.clone();
+        tokio::spawn(async move {
+            loop {
+                delete_expired(&state_mutex).await.unwrap();
+                tokio::time::sleep(Duration::from_secs(2 * 3600)).await;
+            }
+        });
+    }
+
     rocket::build()
         .mount("/", routes![create, go_to_link])
-        .manage(Mutex::new(state))
+        .manage(state_mutex.clone())
         .configure(config)
         .launch()
-        .await.unwrap();
+        .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) {
+#[get("/create?<url>&<secret>&<length>&<link>&<ttl>")]
+async fn create(
+    state: &State<Arc<Mutex<GlobalState>>>,
+    url: &str,
+    secret: Option<&str>,
+    length: Option<u32>,
+    link: Option<&str>,
+    ttl: Option<u64>,
+) -> (Status, String) {
     let mut allow_secret_options = false;
     if let Some(secret) = secret {
         if secret == include_str!("../secret.txt").trim() {
@@ -52,7 +100,11 @@ async fn create(state: &State<Mutex<GlobalState>>, url: &str, secret: Option<&st
             custom_link = true;
         }
     }
-    let actual_len = if allow_secret_options && length.is_some() { length.unwrap() } else { DEFAULT_LINK_LENGTH };
+    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());
     }
@@ -70,7 +122,7 @@ async fn create(state: &State<Mutex<GlobalState>>, url: &str, secret: Option<&st
                     if custom_link {
                         return (Status::BadRequest, "link already exists".to_owned());
                     }
-                },
+                }
                 None => {
                     actual_link = Some(link_id.clone());
                     break;
@@ -83,7 +135,11 @@ async fn create(state: &State<Mutex<GlobalState>>, url: &str, secret: Option<&st
         }
     }
 
-    match db::add_link(&state, actual_link.as_ref().unwrap(), url).await {
+    let res = match ttl {
+        Some(ttl) => db::add_temporary_link(&state, actual_link.as_ref().unwrap(), url, ttl).await,
+        None => db::add_link(&state, actual_link.as_ref().unwrap(), url).await,
+    };
+    match res {
         Ok(_) => (),
         Err(e) => {
             eprintln!("add link error:  {}", e);
@@ -96,7 +152,10 @@ async fn create(state: &State<Mutex<GlobalState>>, url: &str, secret: Option<&st
 }
 
 #[get("/<link>")]
-pub async fn go_to_link(state: &State<Mutex<GlobalState>>, link: &str) -> Result<Option<Redirect>, Status> {
+pub async fn go_to_link(
+    state: &State<Arc<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))),
@@ -108,3 +167,7 @@ pub async fn go_to_link(state: &State<Mutex<GlobalState>>, link: &str) -> Result
         }
     }
 }
+
+async fn delete_expired(state: &Mutex<GlobalState>) -> Result<(), tokio_postgres::Error> {
+    db::delete_expired_links(state).await
+}