A very performant and light (2mb in memory) link shortener and tracker. Written in Rust and React and uses Postgres/SQLite.
nekomimi.pet
1use crate::auth::AuthenticatedUser;
2use crate::{
3 error::AppError,
4 models::{
5 AuthResponse, Claims, ClickStats, CreateLink, Link, LoginRequest, RegisterRequest,
6 SourceStats, User, UserResponse,
7 },
8 AppState,
9};
10use actix_web::{web, HttpRequest, HttpResponse, Responder};
11use argon2::{
12 password_hash::{rand_core::OsRng, SaltString},
13 PasswordVerifier,
14};
15use argon2::{Argon2, PasswordHash, PasswordHasher};
16use jsonwebtoken::{encode, EncodingKey, Header};
17use lazy_static::lazy_static;
18use regex::Regex;
19
20lazy_static! {
21 static ref VALID_CODE_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_-]{1,32}$").unwrap();
22}
23
24pub async fn create_short_url(
25 state: web::Data<AppState>,
26 user: AuthenticatedUser,
27 payload: web::Json<CreateLink>,
28) -> Result<impl Responder, AppError> {
29 tracing::debug!("Creating short URL with user_id: {}", user.user_id);
30
31 validate_url(&payload.url)?;
32
33 let short_code = if let Some(ref custom_code) = payload.custom_code {
34 validate_custom_code(custom_code)?;
35
36 tracing::debug!("Checking if custom code {} exists", custom_code);
37 // Check if code is already taken
38 if let Some(_) = sqlx::query_as::<_, Link>("SELECT * FROM links WHERE short_code = $1")
39 .bind(custom_code)
40 .fetch_optional(&state.db)
41 .await?
42 {
43 return Err(AppError::InvalidInput(
44 "Custom code already taken".to_string(),
45 ));
46 }
47
48 custom_code.clone()
49 } else {
50 generate_short_code()
51 };
52
53 // Start transaction
54 let mut tx = state.db.begin().await?;
55
56 tracing::debug!("Inserting new link with short_code: {}", short_code);
57 let link = sqlx::query_as::<_, Link>(
58 "INSERT INTO links (original_url, short_code, user_id) VALUES ($1, $2, $3) RETURNING *",
59 )
60 .bind(&payload.url)
61 .bind(&short_code)
62 .bind(user.user_id)
63 .fetch_one(&mut *tx)
64 .await?;
65
66 if let Some(ref source) = payload.source {
67 tracing::debug!("Adding click source: {}", source);
68 sqlx::query("INSERT INTO clicks (link_id, source) VALUES ($1, $2)")
69 .bind(link.id)
70 .bind(source)
71 .execute(&mut *tx)
72 .await?;
73 }
74
75 tx.commit().await?;
76 Ok(HttpResponse::Created().json(link))
77}
78
79fn validate_custom_code(code: &str) -> Result<(), AppError> {
80 if !VALID_CODE_REGEX.is_match(code) {
81 return Err(AppError::InvalidInput(
82 "Custom code must be 1-32 characters long and contain only letters, numbers, underscores, and hyphens".to_string()
83 ));
84 }
85
86 // Add reserved words check
87 let reserved_words = ["api", "health", "admin", "static", "assets"];
88 if reserved_words.contains(&code.to_lowercase().as_str()) {
89 return Err(AppError::InvalidInput(
90 "This code is reserved and cannot be used".to_string(),
91 ));
92 }
93
94 Ok(())
95}
96
97fn validate_url(url: &String) -> Result<(), AppError> {
98 if url.is_empty() {
99 return Err(AppError::InvalidInput("URL cannot be empty".to_string()));
100 }
101 if !url.starts_with("http://") && !url.starts_with("https://") {
102 return Err(AppError::InvalidInput(
103 "URL must start with http:// or https://".to_string(),
104 ));
105 }
106 Ok(())
107}
108
109pub async fn redirect_to_url(
110 state: web::Data<AppState>,
111 path: web::Path<String>,
112 req: HttpRequest,
113) -> Result<impl Responder, AppError> {
114 let short_code = path.into_inner();
115
116 // Extract query source if present
117 let query_source = req
118 .uri()
119 .query()
120 .and_then(|q| web::Query::<std::collections::HashMap<String, String>>::from_query(q).ok())
121 .and_then(|params| params.get("source").cloned());
122
123 let mut tx = state.db.begin().await?;
124
125 let link = sqlx::query_as::<_, Link>(
126 "UPDATE links SET clicks = clicks + 1 WHERE short_code = $1 RETURNING *",
127 )
128 .bind(&short_code)
129 .fetch_optional(&mut *tx)
130 .await?;
131
132 match link {
133 Some(link) => {
134 // Record click with both user agent and query source
135 let user_agent = req
136 .headers()
137 .get("user-agent")
138 .and_then(|h| h.to_str().ok())
139 .unwrap_or("unknown")
140 .to_string();
141
142 sqlx::query("INSERT INTO clicks (link_id, source, query_source) VALUES ($1, $2, $3)")
143 .bind(link.id)
144 .bind(user_agent)
145 .bind(query_source)
146 .execute(&mut *tx)
147 .await?;
148
149 tx.commit().await?;
150
151 Ok(HttpResponse::TemporaryRedirect()
152 .append_header(("Location", link.original_url))
153 .finish())
154 }
155 None => Err(AppError::NotFound),
156 }
157}
158
159pub async fn get_all_links(
160 state: web::Data<AppState>,
161 user: AuthenticatedUser,
162) -> Result<impl Responder, AppError> {
163 let links = sqlx::query_as::<_, Link>(
164 "SELECT * FROM links WHERE user_id = $1 ORDER BY created_at DESC",
165 )
166 .bind(user.user_id)
167 .fetch_all(&state.db)
168 .await?;
169
170 Ok(HttpResponse::Ok().json(links))
171}
172
173pub async fn health_check(state: web::Data<AppState>) -> impl Responder {
174 match sqlx::query("SELECT 1").execute(&state.db).await {
175 Ok(_) => HttpResponse::Ok().json("Healthy"),
176 Err(_) => HttpResponse::ServiceUnavailable().json("Database unavailable"),
177 }
178}
179
180fn generate_short_code() -> String {
181 use base62::encode;
182 use uuid::Uuid;
183
184 let uuid = Uuid::new_v4();
185 encode(uuid.as_u128() as u64).chars().take(8).collect()
186}
187
188pub async fn register(
189 state: web::Data<AppState>,
190 payload: web::Json<RegisterRequest>,
191) -> Result<impl Responder, AppError> {
192 // Check if any users exist
193 let user_count = sqlx::query!("SELECT COUNT(*) as count FROM users")
194 .fetch_one(&state.db)
195 .await?
196 .count
197 .unwrap_or(0);
198
199 // If users exist, registration is closed - no exceptions
200 if user_count > 0 {
201 return Err(AppError::Auth("Registration is closed".to_string()));
202 }
203
204 // Verify admin token for first user
205 match (&state.admin_token, &payload.admin_token) {
206 (Some(stored_token), Some(provided_token)) if stored_token == provided_token => {
207 // Token matches, proceed with registration
208 }
209 _ => return Err(AppError::Auth("Invalid admin setup token".to_string())),
210 }
211
212 // Check if email already exists
213 let exists = sqlx::query!("SELECT id FROM users WHERE email = $1", payload.email)
214 .fetch_optional(&state.db)
215 .await?;
216
217 if exists.is_some() {
218 return Err(AppError::Auth("Email already registered".to_string()));
219 }
220
221 let salt = SaltString::generate(&mut OsRng);
222 let argon2 = Argon2::default();
223 let password_hash = argon2
224 .hash_password(payload.password.as_bytes(), &salt)
225 .map_err(|e| AppError::Auth(e.to_string()))?
226 .to_string();
227
228 let user = sqlx::query_as!(
229 User,
230 "INSERT INTO users (email, password_hash) VALUES ($1, $2) RETURNING *",
231 payload.email,
232 password_hash
233 )
234 .fetch_one(&state.db)
235 .await?;
236
237 let claims = Claims::new(user.id);
238 let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
239 let token = encode(
240 &Header::default(),
241 &claims,
242 &EncodingKey::from_secret(secret.as_bytes()),
243 )
244 .map_err(|e| AppError::Auth(e.to_string()))?;
245
246 Ok(HttpResponse::Ok().json(AuthResponse {
247 token,
248 user: UserResponse {
249 id: user.id,
250 email: user.email,
251 },
252 }))
253}
254
255pub async fn login(
256 state: web::Data<AppState>,
257 payload: web::Json<LoginRequest>,
258) -> Result<impl Responder, AppError> {
259 let user = sqlx::query_as!(User, "SELECT * FROM users WHERE email = $1", payload.email)
260 .fetch_optional(&state.db)
261 .await?
262 .ok_or_else(|| AppError::Auth("Invalid credentials".to_string()))?;
263
264 let argon2 = Argon2::default();
265 let parsed_hash =
266 PasswordHash::new(&user.password_hash).map_err(|e| AppError::Auth(e.to_string()))?;
267
268 if argon2
269 .verify_password(payload.password.as_bytes(), &parsed_hash)
270 .is_err()
271 {
272 return Err(AppError::Auth("Invalid credentials".to_string()));
273 }
274
275 let claims = Claims::new(user.id);
276 let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
277 let token = encode(
278 &Header::default(),
279 &claims,
280 &EncodingKey::from_secret(secret.as_bytes()),
281 )
282 .map_err(|e| AppError::Auth(e.to_string()))?;
283
284 Ok(HttpResponse::Ok().json(AuthResponse {
285 token,
286 user: UserResponse {
287 id: user.id,
288 email: user.email,
289 },
290 }))
291}
292
293pub async fn delete_link(
294 state: web::Data<AppState>,
295 user: AuthenticatedUser,
296 path: web::Path<i32>,
297) -> Result<impl Responder, AppError> {
298 let link_id = path.into_inner();
299
300 // Start transaction
301 let mut tx = state.db.begin().await?;
302
303 // Verify the link belongs to the user
304 let link = sqlx::query!(
305 "SELECT id FROM links WHERE id = $1 AND user_id = $2",
306 link_id,
307 user.user_id
308 )
309 .fetch_optional(&mut *tx)
310 .await?;
311
312 if link.is_none() {
313 return Err(AppError::NotFound);
314 }
315
316 // Delete associated clicks first due to foreign key constraint
317 sqlx::query!("DELETE FROM clicks WHERE link_id = $1", link_id)
318 .execute(&mut *tx)
319 .await?;
320
321 // Delete the link
322 sqlx::query!("DELETE FROM links WHERE id = $1", link_id)
323 .execute(&mut *tx)
324 .await?;
325
326 tx.commit().await?;
327
328 Ok(HttpResponse::NoContent().finish())
329}
330
331pub async fn get_link_clicks(
332 state: web::Data<AppState>,
333 user: AuthenticatedUser,
334 path: web::Path<i32>,
335) -> Result<impl Responder, AppError> {
336 let link_id = path.into_inner();
337
338 // Verify the link belongs to the user
339 let link = sqlx::query!(
340 "SELECT id FROM links WHERE id = $1 AND user_id = $2",
341 link_id,
342 user.user_id
343 )
344 .fetch_optional(&state.db)
345 .await?;
346
347 if link.is_none() {
348 return Err(AppError::NotFound);
349 }
350
351 let clicks = sqlx::query_as!(
352 ClickStats,
353 r#"
354 SELECT
355 DATE(created_at)::date as "date!",
356 COUNT(*)::bigint as "clicks!"
357 FROM clicks
358 WHERE link_id = $1
359 GROUP BY DATE(created_at)
360 ORDER BY DATE(created_at) ASC -- Changed from DESC to ASC
361 LIMIT 30
362 "#,
363 link_id
364 )
365 .fetch_all(&state.db)
366 .await?;
367
368 Ok(HttpResponse::Ok().json(clicks))
369}
370
371pub async fn get_link_sources(
372 state: web::Data<AppState>,
373 user: AuthenticatedUser,
374 path: web::Path<i32>,
375) -> Result<impl Responder, AppError> {
376 let link_id = path.into_inner();
377
378 // Verify the link belongs to the user
379 let link = sqlx::query!(
380 "SELECT id FROM links WHERE id = $1 AND user_id = $2",
381 link_id,
382 user.user_id
383 )
384 .fetch_optional(&state.db)
385 .await?;
386
387 if link.is_none() {
388 return Err(AppError::NotFound);
389 }
390
391 let sources = sqlx::query_as!(
392 SourceStats,
393 r#"
394 SELECT
395 query_source as "source!",
396 COUNT(*)::bigint as "count!"
397 FROM clicks
398 WHERE link_id = $1
399 AND query_source IS NOT NULL
400 AND query_source != ''
401 GROUP BY query_source
402 ORDER BY COUNT(*) DESC
403 LIMIT 10
404 "#,
405 link_id
406 )
407 .fetch_all(&state.db)
408 .await?;
409
410 Ok(HttpResponse::Ok().json(sources))
411}