A very performant and light (2mb in memory) link shortener and tracker. Written in Rust and React and uses Postgres/SQLite.
1use actix_cors::Cors;
2use actix_web::{web, App, HttpResponse, HttpServer};
3use anyhow::Result;
4use clap::Parser;
5use rust_embed::RustEmbed;
6use simplelink::check_and_generate_admin_token;
7use simplelink::models::DatabasePool;
8use simplelink::{create_db_pool, run_migrations};
9use simplelink::{handlers, AppState};
10use sqlx::{Postgres, Sqlite};
11use tracing::{error, info};
12
13#[derive(Parser, Debug)]
14#[command(author, version, about, long_about = None)]
15#[derive(RustEmbed)]
16#[folder = "static/"]
17struct Asset;
18
19async fn serve_static_file(path: &str) -> HttpResponse {
20 match Asset::get(path) {
21 Some(content) => {
22 let mime = mime_guess::from_path(path).first_or_octet_stream();
23 HttpResponse::Ok()
24 .content_type(mime.as_ref())
25 .body(content.data.into_owned())
26 }
27 None => HttpResponse::NotFound().body("404 Not Found"),
28 }
29}
30
31async fn create_initial_links(pool: &DatabasePool) -> Result<()> {
32 if let Ok(links) = std::env::var("INITIAL_LINKS") {
33 for link_entry in links.split(';') {
34 let parts: Vec<&str> = link_entry.split(',').collect();
35 if parts.len() >= 2 {
36 let url = parts[0];
37 let code = parts[1];
38
39 match pool {
40 DatabasePool::Postgres(pool) => {
41 sqlx::query(
42 "INSERT INTO links (original_url, short_code, user_id)
43 VALUES ($1, $2, $3)
44 ON CONFLICT (short_code)
45 DO UPDATE SET short_code = EXCLUDED.short_code
46 WHERE links.original_url = EXCLUDED.original_url",
47 )
48 .bind(url)
49 .bind(code)
50 .bind(1)
51 .execute(pool)
52 .await?;
53 }
54 DatabasePool::Sqlite(pool) => {
55 // First check if the exact combination exists
56 let exists = sqlx::query_scalar::<_, bool>(
57 "SELECT EXISTS(
58 SELECT 1 FROM links
59 WHERE original_url = ?1
60 AND short_code = ?2
61 )",
62 )
63 .bind(url)
64 .bind(code)
65 .fetch_one(pool)
66 .await?;
67
68 // Only insert if the exact combination doesn't exist
69 if !exists {
70 sqlx::query(
71 "INSERT INTO links (original_url, short_code, user_id)
72 VALUES (?1, ?2, ?3)",
73 )
74 .bind(url)
75 .bind(code)
76 .bind(1)
77 .execute(pool)
78 .await?;
79 info!("Created initial link: {} -> {} for user_id: 1", code, url);
80 } else {
81 info!("Skipped existing link: {} -> {} for user_id: 1", code, url);
82 }
83 }
84 }
85 }
86 }
87 }
88 Ok(())
89}
90
91async fn create_admin_user(pool: &DatabasePool, email: &str, password: &str) -> Result<()> {
92 use argon2::{
93 password_hash::{rand_core::OsRng, SaltString},
94 Argon2, PasswordHasher,
95 };
96
97 let salt = SaltString::generate(&mut OsRng);
98 let argon2 = Argon2::default();
99 let password_hash = argon2
100 .hash_password(password.as_bytes(), &salt)
101 .map_err(|e| anyhow::anyhow!("Password hashing error: {}", e))?
102 .to_string();
103
104 match pool {
105 DatabasePool::Postgres(pool) => {
106 sqlx::query(
107 "INSERT INTO users (email, password_hash)
108 VALUES ($1, $2)
109 ON CONFLICT (email) DO NOTHING",
110 )
111 .bind(email)
112 .bind(&password_hash)
113 .execute(pool)
114 .await?;
115 }
116 DatabasePool::Sqlite(pool) => {
117 sqlx::query(
118 "INSERT OR IGNORE INTO users (email, password_hash)
119 VALUES (?1, ?2)",
120 )
121 .bind(email)
122 .bind(&password_hash)
123 .execute(pool)
124 .await?;
125 }
126 }
127 info!("Created admin user: {}", email);
128 Ok(())
129}
130
131#[actix_web::main]
132async fn main() -> Result<()> {
133 // Load environment variables from .env file
134 dotenv::dotenv().ok();
135
136 // Initialize logging
137 tracing_subscriber::fmt::init();
138
139 // Create database connection pool
140 let pool = create_db_pool().await?;
141 run_migrations(&pool).await?;
142
143 // First check if admin credentials are provided in environment variables
144 let admin_credentials = match (
145 std::env::var("SIMPLELINK_USER"),
146 std::env::var("SIMPLELINK_PASS"),
147 ) {
148 (Ok(user), Ok(pass)) => Some((user, pass)),
149 _ => None,
150 };
151
152 if let Some((email, password)) = admin_credentials {
153 // Now check for existing users
154 let user_count = match &pool {
155 DatabasePool::Postgres(pool) => {
156 let mut tx = pool.begin().await?;
157 let count =
158 sqlx::query_as::<Postgres, (i64,)>("SELECT COUNT(*)::bigint FROM users")
159 .fetch_one(&mut *tx)
160 .await?
161 .0;
162 tx.commit().await?;
163 count
164 }
165 DatabasePool::Sqlite(pool) => {
166 let mut tx = pool.begin().await?;
167 let count = sqlx::query_as::<Sqlite, (i64,)>("SELECT COUNT(*) FROM users")
168 .fetch_one(&mut *tx)
169 .await?
170 .0;
171 tx.commit().await?;
172 count
173 }
174 };
175
176 if user_count == 0 {
177 info!("No users found, creating admin user: {}", email);
178 match create_admin_user(&pool, &email, &password).await {
179 Ok(_) => info!("Successfully created admin user"),
180 Err(e) => {
181 error!("Failed to create admin user: {}", e);
182 return Err(anyhow::anyhow!("Failed to create admin user: {}", e));
183 }
184 }
185 }
186 } else {
187 info!(
188 "No admin credentials provided in environment variables, skipping admin user creation"
189 );
190 }
191
192 // Create initial links from environment variables
193 create_initial_links(&pool).await?;
194
195 let admin_token = check_and_generate_admin_token(&pool).await?;
196
197 let state = AppState {
198 db: pool,
199 admin_token,
200 };
201
202 let host = std::env::var("SERVER_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
203 let port = std::env::var("SERVER_PORT").unwrap_or_else(|_| "8080".to_string());
204 info!("Starting server at http://{}:{}", host, port);
205
206 // Start HTTP server
207 HttpServer::new(move || {
208 let cors = Cors::default()
209 .allow_any_origin()
210 .allow_any_method()
211 .allow_any_header()
212 .max_age(3600);
213
214 App::new()
215 .wrap(cors)
216 .app_data(web::Data::new(state.clone()))
217 .service(
218 web::scope("/api")
219 .route("/shorten", web::post().to(handlers::create_short_url))
220 .route("/links", web::get().to(handlers::get_all_links))
221 .route("/links/{id}", web::delete().to(handlers::delete_link))
222 .route(
223 "/links/{id}/clicks",
224 web::get().to(handlers::get_link_clicks),
225 )
226 .route(
227 "/links/{id}/sources",
228 web::get().to(handlers::get_link_sources),
229 )
230 .route("/links/{id}", web::patch().to(handlers::edit_link))
231 .route("/auth/register", web::post().to(handlers::register))
232 .route("/auth/login", web::post().to(handlers::login))
233 .route(
234 "/auth/check-first-user",
235 web::get().to(handlers::check_first_user),
236 )
237 .route("/health", web::get().to(handlers::health_check)),
238 )
239 .service(web::resource("/{short_code}").route(web::get().to(handlers::redirect_to_url)))
240 .default_service(web::route().to(|req: actix_web::HttpRequest| async move {
241 let path = req.path().trim_start_matches('/');
242 let path = if path.is_empty() { "index.html" } else { path };
243 serve_static_file(path).await
244 }))
245 })
246 .workers(2)
247 .backlog(10_000)
248 .bind(format!("{}:{}", host, port))?
249 .run()
250 .await?;
251
252 Ok(())
253}