106 lines
3.5 KiB
Rust
106 lines
3.5 KiB
Rust
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)
|
|
}
|
|
}
|
|
} |