Add custom slugs, restore UI, UTM builder, and CLI link tools

This commit is contained in:
thakares committed 2026-06-14 19:11:21 +05:30
1 parent 02a26cfd94
commit c6486763e3
22 files changed
+999 -113

No files matched your search

+32
View File
@@ -0,0 +1,32 @@
use crate::config::Config;
use crate::db::Db;
use std::path::PathBuf;
pub async fn run(
code: String,
data_dir: Option<String>,
mut config: Config,
) -> Result<(), Box<dyn std::error::Error>> {
if let Some(d) = data_dir {
config.data_dir = PathBuf::from(d);
}
let db = Db::init(&config)?;
let normalized_code = code.trim().to_lowercase();
if !crate::utils::validation::validate_redirect_code(&normalized_code) {
return Err("Invalid short code or custom slug format".into());
}
let url_opt = {
let conn = db.content.lock().unwrap();
crate::db::content::get_url_by_code(&conn, &normalized_code)?
};
match url_opt {
Some(url) => {
println!("{}", url.destination);
Ok(())
}
None => Err(format!("Short code not found: {}", normalized_code).into()),
}
}
+19
View File
@@ -3,9 +3,11 @@ use clap::{Parser, Subcommand};
pub mod backup;
pub mod create_admin;
pub mod doctor;
pub mod expand;
pub mod migrate;
pub mod restore;
pub mod serve;
pub mod shorten;
pub mod stats;
pub mod validate;
@@ -72,4 +74,21 @@ pub enum Commands {
#[arg(long)]
data_dir: Option<String>,
},
/// Shorten a URL (Feature 3)
Shorten {
/// The destination URL to shorten
target_url: String,
/// Custom slug (starting with ! followed by a-z, 0-9, -, _)
#[arg(long)]
slug: Option<String>,
#[arg(long)]
data_dir: Option<String>,
},
/// Expand a shortened code or custom slug to its destination URL (Feature 4)
Expand {
/// The short code or custom slug to expand
code: String,
#[arg(long)]
data_dir: Option<String>,
},
}
+41 -4
View File
@@ -6,6 +6,46 @@ use std::path::PathBuf;
use tar::Archive;
use tracing::{error, info};
pub fn perform_restore(
file_path: &std::path::Path,
data_dir: &std::path::Path,
) -> Result<(), Box<dyn std::error::Error>> {
// 1. Open the archive
let f = File::open(file_path)?;
let tar_gz = GzDecoder::new(f);
let mut archive = Archive::new(tar_gz);
// 2. Validate that the archive contains the expected BZOD database files
let mut has_admin = false;
let mut has_content = false;
let mut has_analytics = false;
let mut has_system = false;
for entry_res in archive.entries()? {
let entry = entry_res?;
let path = entry.path()?;
let file_name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
match file_name {
"admin.db" => has_admin = true,
"content.db" => has_content = true,
"analytics.db" => has_analytics = true,
"system.db" => has_system = true,
_ => {}
}
}
if !has_admin || !has_content || !has_analytics || !has_system {
return Err("Archive is missing one or more required database files (admin.db, content.db, analytics.db, system.db)".into());
}
// 3. Unpack archive to data_dir
let f2 = File::open(file_path)?;
let tar_gz2 = GzDecoder::new(f2);
let mut archive2 = Archive::new(tar_gz2);
archive2.unpack(data_dir)?;
Ok(())
}
pub async fn run(
file: String,
data_dir: Option<String>,
@@ -40,10 +80,7 @@ pub async fn run(
}
info!("Restoring backup from: {:?}", file_path);
let f = File::open(&file_path)?;
let tar_gz = GzDecoder::new(f);
let mut archive = Archive::new(tar_gz);
archive.unpack(&config.data_dir)?;
perform_restore(&file_path, &config.data_dir)?;
info!("Database files successfully restored.");
Ok(())
+73
View File
@@ -0,0 +1,73 @@
use crate::config::Config;
use crate::db::Db;
use std::path::PathBuf;
pub async fn run(
target_url: String,
slug: Option<String>,
data_dir: Option<String>,
mut config: Config,
) -> Result<(), Box<dyn std::error::Error>> {
// 1. Basic URL validation
if reqwest::Url::parse(&target_url).is_err() {
return Err("Invalid destination URL format".into());
}
if let Some(d) = data_dir {
config.data_dir = PathBuf::from(d);
}
let db = Db::init(&config)?;
// 2. Validate/normalize slug/code
let code = match slug {
Some(s) => {
let normalized = s.trim().to_lowercase();
if !crate::utils::validation::validate_custom_slug(&normalized) {
return Err(
"Custom slug must start with ! followed by 1-24 characters of a-z, 0-9, -, _"
.into(),
);
}
normalized
}
None => crate::utils::random::generate_token(3),
};
// 3. Persist URL
let conn = db.content.lock().unwrap();
let res = crate::db::content::create_url_extended(
&conn,
&code,
&target_url,
None,
None,
&[],
None,
None,
None,
);
match res {
Ok(_) => {
let proto = if config.cookie_secure {
"https"
} else {
"http"
};
let base_url = config
.base_url
.clone()
.unwrap_or_else(|| format!("{}://localhost:{}", proto, config.port));
// Output only the shortened URL as requested
println!("{}/{}", base_url, code);
Ok(())
}
Err(rusqlite::Error::SqliteFailure(err, _))
if err.code == rusqlite::ErrorCode::ConstraintViolation =>
{
Err("Short code/slug already exists".into())
}
Err(e) => Err(e.into()),
}
}
+14
View File
@@ -51,6 +51,20 @@ pub async fn perform_backup(
std::fs::create_dir_all(&out_dir)?;
}
// Force checkpoint on all databases to flush WAL contents to the main DB files
if let Ok(conn) = db.admin.lock() {
let _ = conn.execute("PRAGMA wal_checkpoint(TRUNCATE);", []);
}
if let Ok(conn) = db.content.lock() {
let _ = conn.execute("PRAGMA wal_checkpoint(TRUNCATE);", []);
}
if let Ok(conn) = db.analytics.lock() {
let _ = conn.execute("PRAGMA wal_checkpoint(TRUNCATE);", []);
}
if let Ok(conn) = db.system.lock() {
let _ = conn.execute("PRAGMA wal_checkpoint(TRUNCATE);", []);
}
let date_str = Utc::now().format("%Y-%m-%d-%H%M%S").to_string();
let tar_name = format!("{}-bzod-backup.tar.gz", date_str);
let tar_path = out_dir.join(tar_name);
+10
View File
@@ -44,6 +44,16 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
Commands::Doctor { data_dir } => {
bzod::cli::doctor::run(data_dir, config).await?;
}
Commands::Shorten {
target_url,
slug,
data_dir,
} => {
bzod::cli::shorten::run(target_url, slug, data_dir, config).await?;
}
Commands::Expand { code, data_dir } => {
bzod::cli::expand::run(code, data_dir, config).await?;
}
}
Ok(())
+1
View File
@@ -3,6 +3,7 @@ pub mod network;
pub mod random;
pub mod system;
pub mod time;
pub mod validation;
pub use hashing::sha256_hash;
pub use network::get_client_ip;
+19
View File
@@ -0,0 +1,19 @@
pub fn validate_custom_slug(slug: &str) -> bool {
if !slug.starts_with('!') {
return false;
}
let rest = &slug[1..];
if rest.is_empty() || rest.len() > 24 {
return false;
}
rest.chars()
.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-' || c == '_')
}
pub fn validate_redirect_code(code: &str) -> bool {
(code.len() == 6 && code.chars().all(|c| c.is_ascii_hexdigit())) || validate_custom_slug(code)
}
pub fn validate_page_code(code: &str) -> bool {
(code.len() == 4 && code.chars().all(|c| c.is_ascii_hexdigit())) || validate_custom_slug(code)
}
+248 -11
View File
@@ -399,6 +399,7 @@ pub async fn urls_get(
pub struct CreateUrlForm {
pub destination: String,
pub code: String,
pub custom_slug: String,
pub title: String,
pub description: String,
pub tags: String,
@@ -406,6 +407,9 @@ pub struct CreateUrlForm {
pub expires_at: String,
pub password: String,
pub max_access_count: String,
pub utm_source: String,
pub utm_medium: String,
pub utm_campaign: String,
}
// POST /admin/urls/create
@@ -426,13 +430,48 @@ pub async fn urls_create(
}
let ip = get_client_ip(&headers, connect_info);
let mut code = form.code.trim().to_lowercase();
// Custom Slug takes priority if provided
let mut code = form.custom_slug.trim().to_lowercase();
if code.is_empty() {
code = generate_token(3);
} else {
if code.len() != 6 || !code.chars().all(|c| c.is_ascii_hexdigit()) {
return Redirect::to("/admin/urls?error=Custom code must be exactly 6 hex characters")
code = form.code.trim().to_lowercase();
if code.is_empty() {
code = generate_token(3);
} else {
if code.len() != 6 || !code.chars().all(|c| c.is_ascii_hexdigit()) {
return Redirect::to(
"/admin/urls?error=Custom code must be exactly 6 hex characters",
)
.into_response();
}
}
} else {
if !crate::utils::validation::validate_custom_slug(&code) {
return Redirect::to("/admin/urls?error=Custom slug must start with ! followed by 1-24 characters of a-z, 0-9, -, _")
.into_response();
}
}
let mut dest = form.destination.trim().to_string();
if let Ok(mut parsed) = reqwest::Url::parse(&dest) {
let mut has_utm = false;
{
let mut query = parsed.query_pairs_mut();
if !form.utm_source.trim().is_empty() {
query.append_pair("utm_source", form.utm_source.trim());
has_utm = true;
}
if !form.utm_medium.trim().is_empty() {
query.append_pair("utm_medium", form.utm_medium.trim());
has_utm = true;
}
if !form.utm_campaign.trim().is_empty() {
query.append_pair("utm_campaign", form.utm_campaign.trim());
has_utm = true;
}
}
if has_utm {
dest = parsed.to_string();
}
}
@@ -489,7 +528,7 @@ pub async fn urls_create(
crate::db::content::create_url_extended(
&conn,
&code,
&form.destination,
&dest,
title_opt,
desc_opt,
&tags_list,
@@ -519,7 +558,7 @@ pub async fn urls_create(
Err(rusqlite::Error::SqliteFailure(err, _))
if err.code == rusqlite::ErrorCode::ConstraintViolation =>
{
Redirect::to("/admin/urls?error=Short code already exists").into_response()
Redirect::to("/admin/urls?error=Short code/slug already exists").into_response()
}
Err(e) => Redirect::to(&format!("/admin/urls?error=Database error: {}", e)).into_response(),
}
@@ -608,6 +647,7 @@ pub struct CreatePageForm {
pub title: String,
pub slug: String,
pub code: String,
pub custom_slug: String,
pub state: String,
pub html_content: String,
pub csrf_token: String,
@@ -631,12 +671,24 @@ pub async fn pages_create(
}
let ip = get_client_ip(&headers, connect_info);
let mut code = form.code.trim().to_lowercase();
// Custom Slug takes priority if provided
let mut code = form.custom_slug.trim().to_lowercase();
if code.is_empty() {
code = generate_token(2);
code = form.code.trim().to_lowercase();
if code.is_empty() {
code = generate_token(2);
} else {
if code.len() != 4 || !code.chars().all(|c| c.is_ascii_hexdigit()) {
return Redirect::to(
"/admin/pages?error=Custom code must be exactly 4 hex characters",
)
.into_response();
}
}
} else {
if code.len() != 4 || !code.chars().all(|c| c.is_ascii_hexdigit()) {
return Redirect::to("/admin/pages?error=Custom code must be exactly 4 hex characters")
if !crate::utils::validation::validate_custom_slug(&code) {
return Redirect::to("/admin/pages?error=Custom slug must start with ! followed by 1-24 characters of a-z, 0-9, -, _")
.into_response();
}
}
@@ -1301,3 +1353,188 @@ pub async fn status_get(State(state): State<AppState>, jar: CookieJar) -> Respon
template.into_response()
}
// POST /admin/settings/restore
pub async fn restore_backup_post(
State(state): State<AppState>,
jar: CookieJar,
headers: HeaderMap,
connect_info: Option<ConnectInfo<SocketAddr>>,
mut multipart: axum::extract::Multipart,
) -> Response {
let (user, session_id) = match require_auth(&state, &jar).await {
Ok(u) => u,
Err(redir) => return redir.into_response(),
};
let ip = get_client_ip(&headers, connect_info);
let mut file_bytes = Vec::new();
let mut confirm_text = String::new();
let mut csrf_token = String::new();
while let Ok(Some(field)) = multipart.next_field().await {
let name = field.name().unwrap_or("").to_string();
if name == "backup_file" {
if let Ok(bytes) = field.bytes().await {
file_bytes = bytes.to_vec();
}
} else if name == "confirm_text" {
if let Ok(text) = field.text().await {
confirm_text = text.trim().to_string();
}
} else if name == "csrf_token" {
if let Ok(token) = field.text().await {
csrf_token = token.trim().to_string();
}
}
}
if !verify_csrf(&session_id, &csrf_token) {
return Redirect::to("/admin/settings?error=Invalid CSRF token").into_response();
}
if confirm_text != "RESTORE" {
return Redirect::to("/admin/settings?error=Confirmation text must be exactly 'RESTORE'")
.into_response();
}
if file_bytes.is_empty() {
return Redirect::to("/admin/settings?error=No backup file uploaded").into_response();
}
// Save uploaded archive to a temporary file
let temp_file_path =
std::env::temp_dir().join(format!("bzod_restore_{}.tar.gz", uuid::Uuid::new_v4()));
if let Err(e) = std::fs::write(&temp_file_path, &file_bytes) {
return Redirect::to(&format!(
"/admin/settings?error=Failed to write temp file: {}",
e
))
.into_response();
}
// Log RESTORE_INITIATED audit event before restore
{
let conn = state.admin_db.lock().unwrap();
let _ = write_audit_log(
&conn,
&state,
&user.username,
"RESTORE_INITIATED",
Some("system"),
Some("tarball"),
Some(&ip),
headers.get("user-agent").and_then(|h| h.to_str().ok()),
);
}
// Call the perform_restore engine inside closed connection blocks
let restore_res = {
// Temporarily suspend access to active SQLite connections
let mut admin_conn = state.admin_db.lock().unwrap();
let mut content_conn = state.content_db.lock().unwrap();
let mut analytics_conn = state.analytics_db.lock().unwrap();
let mut system_conn = state.system_db.lock().unwrap();
// 1. Close current connections by replacing them with dummy in-memory DBs
*admin_conn = match rusqlite::Connection::open_in_memory() {
Ok(c) => c,
Err(e) => {
return Redirect::to(&format!(
"/admin/settings?error=Failed to open temp in-memory DB: {}",
e
))
.into_response()
}
};
*content_conn = match rusqlite::Connection::open_in_memory() {
Ok(c) => c,
Err(e) => {
return Redirect::to(&format!(
"/admin/settings?error=Failed to open temp in-memory DB: {}",
e
))
.into_response()
}
};
*analytics_conn = match rusqlite::Connection::open_in_memory() {
Ok(c) => c,
Err(e) => {
return Redirect::to(&format!(
"/admin/settings?error=Failed to open temp in-memory DB: {}",
e
))
.into_response()
}
};
*system_conn = match rusqlite::Connection::open_in_memory() {
Ok(c) => c,
Err(e) => {
return Redirect::to(&format!(
"/admin/settings?error=Failed to open temp in-memory DB: {}",
e
))
.into_response()
}
};
// 2. Perform restore unpacking/validation
let res = crate::cli::restore::perform_restore(&temp_file_path, &state.config.data_dir);
// 3. Reinitialize database connections
let new_admin = rusqlite::Connection::open(state.config.data_dir.join("admin.db"));
let new_content = rusqlite::Connection::open(state.config.data_dir.join("content.db"));
let new_analytics = rusqlite::Connection::open(state.config.data_dir.join("analytics.db"));
let new_system = rusqlite::Connection::open(state.config.data_dir.join("system.db"));
match (new_admin, new_content, new_analytics, new_system) {
(Ok(adm), Ok(cnt), Ok(any), Ok(sys)) => {
let _ = crate::db::sqlite::enable_wal(&adm, "admin");
let _ = crate::db::sqlite::enable_wal(&cnt, "content");
let _ = crate::db::sqlite::enable_wal(&any, "analytics");
let _ = crate::db::sqlite::enable_wal(&sys, "system");
let _ = crate::db::sqlite::enable_foreign_keys(&adm, "admin");
let _ = crate::db::sqlite::enable_foreign_keys(&cnt, "content");
let _ = crate::db::sqlite::enable_foreign_keys(&any, "analytics");
let _ = crate::db::sqlite::enable_foreign_keys(&sys, "system");
*admin_conn = adm;
*content_conn = cnt;
*analytics_conn = any;
*system_conn = sys;
}
_ => {
return Redirect::to("/admin/settings?error=Failed to reopen restored databases")
.into_response();
}
}
res
};
let _ = std::fs::remove_file(&temp_file_path);
match restore_res {
Ok(_) => {
// Write database restore success log to newly restored admin db
{
let conn = state.admin_db.lock().unwrap();
let _ = write_audit_log(
&conn,
&state,
&user.username,
"DATABASE_RESTORE",
Some("system"),
Some("tarball"),
Some(&ip),
headers.get("user-agent").and_then(|h| h.to_str().ok()),
);
}
Redirect::to("/admin/login").into_response()
}
Err(e) => {
Redirect::to(&format!("/admin/settings?error=Restore failed: {}", e)).into_response()
}
}
}
+38 -5
View File
@@ -33,6 +33,9 @@ pub struct CreateUrlRequest {
pub expires_at: Option<String>,
pub password: Option<String>,
pub max_access_count: Option<i64>,
pub utm_source: Option<String>,
pub utm_medium: Option<String>,
pub utm_campaign: Option<String>,
}
#[derive(Deserialize)]
@@ -84,17 +87,47 @@ pub async fn api_create_url(
if code.is_empty() {
code = generate_token(3); // 6 hex
} else {
if code.len() != 6 || !code.chars().all(|c| c.is_ascii_hexdigit()) {
if !crate::utils::validation::validate_redirect_code(&code) {
return (
StatusCode::BAD_REQUEST,
Json(ApiError {
error: "Short code must be 6 hex characters".to_string(),
error: "Short code must be 6 hex characters or a custom slug starting with !"
.to_string(),
}),
)
.into_response();
}
}
let mut dest = payload.destination.trim().to_string();
if let Ok(mut parsed) = reqwest::Url::parse(&dest) {
let mut has_utm = false;
{
let mut query = parsed.query_pairs_mut();
if let Some(ref src) = payload.utm_source {
if !src.trim().is_empty() {
query.append_pair("utm_source", src.trim());
has_utm = true;
}
}
if let Some(ref med) = payload.utm_medium {
if !med.trim().is_empty() {
query.append_pair("utm_medium", med.trim());
has_utm = true;
}
}
if let Some(ref camp) = payload.utm_campaign {
if !camp.trim().is_empty() {
query.append_pair("utm_campaign", camp.trim());
has_utm = true;
}
}
}
if has_utm {
dest = parsed.to_string();
}
}
let password_hash = if let Some(ref pwd) = payload.password {
if pwd.is_empty() {
None
@@ -121,7 +154,7 @@ pub async fn api_create_url(
match crate::db::content::create_url_extended(
&conn,
&code,
&payload.destination,
&dest,
payload.title.as_deref(),
payload.description.as_deref(),
&tags,
@@ -390,11 +423,11 @@ pub async fn api_create_page(
if code.is_empty() {
code = generate_token(2); // 4 hex
} else {
if code.len() != 4 || !code.chars().all(|c| c.is_ascii_hexdigit()) {
if !crate::utils::validation::validate_page_code(&code) {
return (
StatusCode::BAD_REQUEST,
Json(ApiError {
error: "Short code must be 4 hex characters".to_string(),
error: "Short code must be 4 hex characters or start with ! followed by 1-24 characters of a-z, 0-9, -, _".to_string(),
}),
)
.into_response();
+1 -1
View File
@@ -21,7 +21,7 @@ pub async fn resolve_page(
headers: HeaderMap,
connect_info: Option<ConnectInfo<SocketAddr>>,
) -> Response {
if code.len() != 4 || !code.chars().all(|c| c.is_ascii_hexdigit()) {
if !crate::utils::validation::validate_page_code(&code) {
return (StatusCode::NOT_FOUND, "Not Found").into_response();
}
+1 -1
View File
@@ -72,7 +72,7 @@ pub async fn qr_handler(
let code = parts[0];
let ext = parts[1].to_lowercase();
if code.len() != 6 || !code.chars().all(|c| c.is_ascii_hexdigit()) {
if !crate::utils::validation::validate_redirect_code(code) {
return (StatusCode::NOT_FOUND, "Not Found").into_response();
}
+2 -2
View File
@@ -24,8 +24,8 @@ pub async fn resolve_redirect(
headers: HeaderMap,
connect_info: Option<ConnectInfo<SocketAddr>>,
) -> Response {
// Basic validation of code (must be 6 hex characters)
if code.len() != 6 || !code.chars().all(|c| c.is_ascii_hexdigit()) {
// Basic validation of code (must be 6 hex characters or a valid custom slug)
if !crate::utils::validation::validate_redirect_code(&code) {
return (StatusCode::NOT_FOUND, "Not Found").into_response();
}
+1
View File
@@ -48,6 +48,7 @@ pub fn create_router(state: AppState) -> Router {
)
.route("/admin/settings/compact", post(admin::compact_db_post))
.route("/admin/settings/backup", get(admin::download_backup))
.route("/admin/settings/restore", post(admin::restore_backup_post))
.route("/admin/settings/bulk-qr", post(admin::bulk_qr_export_post))
.route(
"/admin/settings/api-keys/create",