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, DatabasePool, Link, LoginRequest,
6 RegisterRequest, 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;
19use sqlx::{Postgres, Sqlite};
20
21lazy_static! {
22 static ref VALID_CODE_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_-]{1,32}$").unwrap();
23}
24
25pub async fn create_short_url(
26 state: web::Data<AppState>,
27 user: AuthenticatedUser,
28 payload: web::Json<CreateLink>,
29) -> Result<impl Responder, AppError> {
30 tracing::debug!("Creating short URL with user_id: {}", user.user_id);
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 // Check if code exists using match on pool type
37 let exists = match &state.db {
38 DatabasePool::Postgres(pool) => {
39 sqlx::query_as::<_, Link>("SELECT * FROM links WHERE short_code = $1")
40 .bind(custom_code)
41 .fetch_optional(pool)
42 .await?
43 }
44 DatabasePool::Sqlite(pool) => {
45 sqlx::query_as::<_, Link>("SELECT * FROM links WHERE short_code = ?1")
46 .bind(custom_code)
47 .fetch_optional(pool)
48 .await?
49 }
50 };
51
52 if exists.is_some() {
53 return Err(AppError::InvalidInput(
54 "Custom code already taken".to_string(),
55 ));
56 }
57 custom_code.clone()
58 } else {
59 generate_short_code()
60 };
61
62 // Start transaction based on pool type
63 let result = match &state.db {
64 DatabasePool::Postgres(pool) => {
65 let mut tx = pool.begin().await?;
66
67 let link = sqlx::query_as::<_, Link>(
68 "INSERT INTO links (original_url, short_code, user_id) VALUES ($1, $2, $3) RETURNING *"
69 )
70 .bind(&payload.url)
71 .bind(&short_code)
72 .bind(user.user_id)
73 .fetch_one(&mut *tx)
74 .await?;
75
76 if let Some(ref source) = payload.source {
77 sqlx::query("INSERT INTO clicks (link_id, source) VALUES ($1, $2)")
78 .bind(link.id)
79 .bind(source)
80 .execute(&mut *tx)
81 .await?;
82 }
83
84 tx.commit().await?;
85 link
86 }
87 DatabasePool::Sqlite(pool) => {
88 let mut tx = pool.begin().await?;
89
90 let link = sqlx::query_as::<_, Link>(
91 "INSERT INTO links (original_url, short_code, user_id) VALUES (?1, ?2, ?3) RETURNING *"
92 )
93 .bind(&payload.url)
94 .bind(&short_code)
95 .bind(user.user_id)
96 .fetch_one(&mut *tx)
97 .await?;
98
99 if let Some(ref source) = payload.source {
100 sqlx::query("INSERT INTO clicks (link_id, source) VALUES (?1, ?2)")
101 .bind(link.id)
102 .bind(source)
103 .execute(&mut *tx)
104 .await?;
105 }
106
107 tx.commit().await?;
108 link
109 }
110 };
111
112 Ok(HttpResponse::Created().json(result))
113}
114
115fn validate_custom_code(code: &str) -> Result<(), AppError> {
116 if !VALID_CODE_REGEX.is_match(code) {
117 return Err(AppError::InvalidInput(
118 "Custom code must be 1-32 characters long and contain only letters, numbers, underscores, and hyphens".to_string()
119 ));
120 }
121
122 // Add reserved words check
123 let reserved_words = ["api", "health", "admin", "static", "assets"];
124 if reserved_words.contains(&code.to_lowercase().as_str()) {
125 return Err(AppError::InvalidInput(
126 "This code is reserved and cannot be used".to_string(),
127 ));
128 }
129
130 Ok(())
131}
132
133fn validate_url(url: &String) -> Result<(), AppError> {
134 if url.is_empty() {
135 return Err(AppError::InvalidInput("URL cannot be empty".to_string()));
136 }
137 if !url.starts_with("http://") && !url.starts_with("https://") {
138 return Err(AppError::InvalidInput(
139 "URL must start with http:// or https://".to_string(),
140 ));
141 }
142 Ok(())
143}
144
145pub async fn redirect_to_url(
146 state: web::Data<AppState>,
147 path: web::Path<String>,
148 req: HttpRequest,
149) -> Result<impl Responder, AppError> {
150 let short_code = path.into_inner();
151
152 // Extract query source if present
153 let query_source = req
154 .uri()
155 .query()
156 .and_then(|q| web::Query::<std::collections::HashMap<String, String>>::from_query(q).ok())
157 .and_then(|params| params.get("source").cloned());
158
159 let link = match &state.db {
160 DatabasePool::Postgres(pool) => {
161 let mut tx = pool.begin().await?;
162 let link = sqlx::query_as::<_, Link>(
163 "UPDATE links SET clicks = clicks + 1 WHERE short_code = $1 RETURNING *",
164 )
165 .bind(&short_code)
166 .fetch_optional(&mut *tx)
167 .await?;
168 tx.commit().await?;
169 link
170 }
171 DatabasePool::Sqlite(pool) => {
172 let mut tx = pool.begin().await?;
173 let link = sqlx::query_as::<_, Link>(
174 "UPDATE links SET clicks = clicks + 1 WHERE short_code = ?1 RETURNING *",
175 )
176 .bind(&short_code)
177 .fetch_optional(&mut *tx)
178 .await?;
179 tx.commit().await?;
180 link
181 }
182 };
183
184 match link {
185 Some(link) => {
186 // Handle click recording based on database type
187 match &state.db {
188 DatabasePool::Postgres(pool) => {
189 let mut tx = pool.begin().await?;
190 let user_agent = req
191 .headers()
192 .get("user-agent")
193 .and_then(|h| h.to_str().ok())
194 .unwrap_or("unknown")
195 .to_string();
196
197 sqlx::query(
198 "INSERT INTO clicks (link_id, source, query_source) VALUES ($1, $2, $3)",
199 )
200 .bind(link.id)
201 .bind(user_agent)
202 .bind(query_source)
203 .execute(&mut *tx)
204 .await?;
205
206 tx.commit().await?;
207 }
208 DatabasePool::Sqlite(pool) => {
209 let mut tx = pool.begin().await?;
210 let user_agent = req
211 .headers()
212 .get("user-agent")
213 .and_then(|h| h.to_str().ok())
214 .unwrap_or("unknown")
215 .to_string();
216
217 sqlx::query(
218 "INSERT INTO clicks (link_id, source, query_source) VALUES (?1, ?2, ?3)",
219 )
220 .bind(link.id)
221 .bind(user_agent)
222 .bind(query_source)
223 .execute(&mut *tx)
224 .await?;
225
226 tx.commit().await?;
227 }
228 };
229
230 Ok(HttpResponse::TemporaryRedirect()
231 .append_header(("Location", link.original_url))
232 .finish())
233 }
234 None => Err(AppError::NotFound),
235 }
236}
237
238pub async fn get_all_links(
239 state: web::Data<AppState>,
240 user: AuthenticatedUser,
241) -> Result<impl Responder, AppError> {
242 let links = match &state.db {
243 DatabasePool::Postgres(pool) => {
244 sqlx::query_as::<_, Link>(
245 "SELECT * FROM links WHERE user_id = $1 ORDER BY created_at DESC",
246 )
247 .bind(user.user_id)
248 .fetch_all(pool)
249 .await?
250 }
251 DatabasePool::Sqlite(pool) => {
252 sqlx::query_as::<_, Link>(
253 "SELECT * FROM links WHERE user_id = ?1 ORDER BY created_at DESC",
254 )
255 .bind(user.user_id)
256 .fetch_all(pool)
257 .await?
258 }
259 };
260
261 Ok(HttpResponse::Ok().json(links))
262}
263
264pub async fn health_check(state: web::Data<AppState>) -> impl Responder {
265 let is_healthy = match &state.db {
266 DatabasePool::Postgres(pool) => sqlx::query("SELECT 1").execute(pool).await.is_ok(),
267 DatabasePool::Sqlite(pool) => sqlx::query("SELECT 1").execute(pool).await.is_ok(),
268 };
269
270 if is_healthy {
271 HttpResponse::Ok().json("Healthy")
272 } else {
273 HttpResponse::ServiceUnavailable().json("Database unavailable")
274 }
275}
276
277fn generate_short_code() -> String {
278 use base62::encode;
279 use uuid::Uuid;
280
281 let uuid = Uuid::new_v4();
282 encode(uuid.as_u128() as u64).chars().take(32).collect()
283}
284
285pub async fn register(
286 state: web::Data<AppState>,
287 payload: web::Json<RegisterRequest>,
288) -> Result<impl Responder, AppError> {
289 // Check if any users exist
290 let user_count = match &state.db {
291 DatabasePool::Postgres(pool) => {
292 let mut tx = pool.begin().await?;
293 let count = sqlx::query_as::<Postgres, (i64,)>("SELECT COUNT(*)::bigint FROM users")
294 .fetch_one(&mut *tx)
295 .await?
296 .0;
297 tx.commit().await?;
298 count
299 }
300 DatabasePool::Sqlite(pool) => {
301 let mut tx = pool.begin().await?;
302 let count = sqlx::query_as::<Sqlite, (i64,)>("SELECT COUNT(*) FROM users")
303 .fetch_one(&mut *tx)
304 .await?
305 .0;
306 tx.commit().await?;
307 count
308 }
309 };
310
311 // If users exist, registration is closed - no exceptions
312 if user_count > 0 {
313 return Err(AppError::Auth("Registration is closed".to_string()));
314 }
315
316 // Verify admin token for first user
317 match (&state.admin_token, &payload.admin_token) {
318 (Some(stored_token), Some(provided_token)) if stored_token == provided_token => {
319 // Token matches, proceed with registration
320 }
321 _ => return Err(AppError::Auth("Invalid admin setup token".to_string())),
322 }
323
324 // Check if email already exists
325 let exists = match &state.db {
326 DatabasePool::Postgres(pool) => {
327 let mut tx = pool.begin().await?;
328 let exists =
329 sqlx::query_as::<Postgres, (i32,)>("SELECT id FROM users WHERE email = $1")
330 .bind(&payload.email)
331 .fetch_optional(&mut *tx)
332 .await?;
333 tx.commit().await?;
334 exists
335 }
336 DatabasePool::Sqlite(pool) => {
337 let mut tx = pool.begin().await?;
338 let exists = sqlx::query_as::<Sqlite, (i32,)>("SELECT id FROM users WHERE email = ?")
339 .bind(&payload.email)
340 .fetch_optional(&mut *tx)
341 .await?;
342 tx.commit().await?;
343 exists
344 }
345 };
346
347 if exists.is_some() {
348 return Err(AppError::Auth("Email already registered".to_string()));
349 }
350
351 let salt = SaltString::generate(&mut OsRng);
352 let argon2 = Argon2::default();
353 let password_hash = argon2
354 .hash_password(payload.password.as_bytes(), &salt)
355 .map_err(|e| AppError::Auth(e.to_string()))?
356 .to_string();
357
358 // Insert new user
359 let user = match &state.db {
360 DatabasePool::Postgres(pool) => {
361 let mut tx = pool.begin().await?;
362 let user = sqlx::query_as::<Postgres, User>(
363 "INSERT INTO users (email, password_hash) VALUES ($1, $2) RETURNING *",
364 )
365 .bind(&payload.email)
366 .bind(&password_hash)
367 .fetch_one(&mut *tx)
368 .await?;
369 tx.commit().await?;
370 user
371 }
372 DatabasePool::Sqlite(pool) => {
373 let mut tx = pool.begin().await?;
374 let user = sqlx::query_as::<Sqlite, User>(
375 "INSERT INTO users (email, password_hash) VALUES (?, ?) RETURNING *",
376 )
377 .bind(&payload.email)
378 .bind(&password_hash)
379 .fetch_one(&mut *tx)
380 .await?;
381 tx.commit().await?;
382 user
383 }
384 };
385
386 let claims = Claims::new(user.id);
387 let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
388 let token = encode(
389 &Header::default(),
390 &claims,
391 &EncodingKey::from_secret(secret.as_bytes()),
392 )
393 .map_err(|e| AppError::Auth(e.to_string()))?;
394
395 Ok(HttpResponse::Ok().json(AuthResponse {
396 token,
397 user: UserResponse {
398 id: user.id,
399 email: user.email,
400 },
401 }))
402}
403
404pub async fn login(
405 state: web::Data<AppState>,
406 payload: web::Json<LoginRequest>,
407) -> Result<impl Responder, AppError> {
408 let user = match &state.db {
409 DatabasePool::Postgres(pool) => {
410 let mut tx = pool.begin().await?;
411 let user = sqlx::query_as::<Postgres, User>("SELECT * FROM users WHERE email = $1")
412 .bind(&payload.email)
413 .fetch_optional(&mut *tx)
414 .await?;
415 tx.commit().await?;
416 user
417 }
418 DatabasePool::Sqlite(pool) => {
419 let mut tx = pool.begin().await?;
420 let user = sqlx::query_as::<Sqlite, User>("SELECT * FROM users WHERE email = ?")
421 .bind(&payload.email)
422 .fetch_optional(&mut *tx)
423 .await?;
424 tx.commit().await?;
425 user
426 }
427 }
428 .ok_or_else(|| AppError::Auth("Invalid credentials".to_string()))?;
429
430 let argon2 = Argon2::default();
431 let parsed_hash =
432 PasswordHash::new(&user.password_hash).map_err(|e| AppError::Auth(e.to_string()))?;
433
434 if argon2
435 .verify_password(payload.password.as_bytes(), &parsed_hash)
436 .is_err()
437 {
438 return Err(AppError::Auth("Invalid credentials".to_string()));
439 }
440
441 let claims = Claims::new(user.id);
442 let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
443 let token = encode(
444 &Header::default(),
445 &claims,
446 &EncodingKey::from_secret(secret.as_bytes()),
447 )
448 .map_err(|e| AppError::Auth(e.to_string()))?;
449
450 Ok(HttpResponse::Ok().json(AuthResponse {
451 token,
452 user: UserResponse {
453 id: user.id,
454 email: user.email,
455 },
456 }))
457}
458
459pub async fn delete_link(
460 state: web::Data<AppState>,
461 user: AuthenticatedUser,
462 path: web::Path<i32>,
463) -> Result<impl Responder, AppError> {
464 let link_id = path.into_inner();
465
466 match &state.db {
467 DatabasePool::Postgres(pool) => {
468 let mut tx = pool.begin().await?;
469
470 // Verify the link belongs to the user
471 let link = sqlx::query_as::<Postgres, (i32,)>(
472 "SELECT id FROM links WHERE id = $1 AND user_id = $2",
473 )
474 .bind(link_id)
475 .bind(user.user_id)
476 .fetch_optional(&mut *tx)
477 .await?;
478
479 if link.is_none() {
480 return Err(AppError::NotFound);
481 }
482
483 // Delete associated clicks first due to foreign key constraint
484 sqlx::query("DELETE FROM clicks WHERE link_id = $1")
485 .bind(link_id)
486 .execute(&mut *tx)
487 .await?;
488
489 // Delete the link
490 sqlx::query("DELETE FROM links WHERE id = $1")
491 .bind(link_id)
492 .execute(&mut *tx)
493 .await?;
494
495 tx.commit().await?;
496 }
497 DatabasePool::Sqlite(pool) => {
498 let mut tx = pool.begin().await?;
499
500 // Verify the link belongs to the user
501 let link = sqlx::query_as::<Sqlite, (i32,)>(
502 "SELECT id FROM links WHERE id = ? AND user_id = ?",
503 )
504 .bind(link_id)
505 .bind(user.user_id)
506 .fetch_optional(&mut *tx)
507 .await?;
508
509 if link.is_none() {
510 return Err(AppError::NotFound);
511 }
512
513 // Delete associated clicks first due to foreign key constraint
514 sqlx::query("DELETE FROM clicks WHERE link_id = ?")
515 .bind(link_id)
516 .execute(&mut *tx)
517 .await?;
518
519 // Delete the link
520 sqlx::query("DELETE FROM links WHERE id = ?")
521 .bind(link_id)
522 .execute(&mut *tx)
523 .await?;
524
525 tx.commit().await?;
526 }
527 }
528
529 Ok(HttpResponse::NoContent().finish())
530}
531
532pub async fn get_link_clicks(
533 state: web::Data<AppState>,
534 user: AuthenticatedUser,
535 path: web::Path<i32>,
536) -> Result<impl Responder, AppError> {
537 let link_id = path.into_inner();
538
539 // Verify the link belongs to the user
540 let link = match &state.db {
541 DatabasePool::Postgres(pool) => {
542 let mut tx = pool.begin().await?;
543 let link = sqlx::query_as::<Postgres, (i32,)>(
544 "SELECT id FROM links WHERE id = $1 AND user_id = $2",
545 )
546 .bind(link_id)
547 .bind(user.user_id)
548 .fetch_optional(&mut *tx)
549 .await?;
550 tx.commit().await?;
551 link
552 }
553 DatabasePool::Sqlite(pool) => {
554 let mut tx = pool.begin().await?;
555 let link = sqlx::query_as::<Sqlite, (i32,)>(
556 "SELECT id FROM links WHERE id = ? AND user_id = ?",
557 )
558 .bind(link_id)
559 .bind(user.user_id)
560 .fetch_optional(&mut *tx)
561 .await?;
562 tx.commit().await?;
563 link
564 }
565 };
566
567 if link.is_none() {
568 return Err(AppError::NotFound);
569 }
570
571 let clicks = match &state.db {
572 DatabasePool::Postgres(pool) => {
573 sqlx::query_as::<Postgres, ClickStats>(
574 r#"
575 SELECT
576 DATE(created_at)::date as "date!",
577 COUNT(*)::bigint as "clicks!"
578 FROM clicks
579 WHERE link_id = $1
580 GROUP BY DATE(created_at)
581 ORDER BY DATE(created_at) ASC
582 LIMIT 30
583 "#,
584 )
585 .bind(link_id)
586 .fetch_all(pool)
587 .await?
588 }
589 DatabasePool::Sqlite(pool) => {
590 sqlx::query_as::<Sqlite, ClickStats>(
591 r#"
592 SELECT
593 DATE(created_at) as "date!",
594 COUNT(*) as "clicks!"
595 FROM clicks
596 WHERE link_id = ?
597 GROUP BY DATE(created_at)
598 ORDER BY DATE(created_at) ASC
599 LIMIT 30
600 "#,
601 )
602 .bind(link_id)
603 .fetch_all(pool)
604 .await?
605 }
606 };
607
608 Ok(HttpResponse::Ok().json(clicks))
609}
610
611pub async fn get_link_sources(
612 state: web::Data<AppState>,
613 user: AuthenticatedUser,
614 path: web::Path<i32>,
615) -> Result<impl Responder, AppError> {
616 let link_id = path.into_inner();
617
618 // Verify the link belongs to the user
619 let link = match &state.db {
620 DatabasePool::Postgres(pool) => {
621 let mut tx = pool.begin().await?;
622 let link = sqlx::query_as::<Postgres, (i32,)>(
623 "SELECT id FROM links WHERE id = $1 AND user_id = $2",
624 )
625 .bind(link_id)
626 .bind(user.user_id)
627 .fetch_optional(&mut *tx)
628 .await?;
629 tx.commit().await?;
630 link
631 }
632 DatabasePool::Sqlite(pool) => {
633 let mut tx = pool.begin().await?;
634 let link = sqlx::query_as::<Sqlite, (i32,)>(
635 "SELECT id FROM links WHERE id = ? AND user_id = ?",
636 )
637 .bind(link_id)
638 .bind(user.user_id)
639 .fetch_optional(&mut *tx)
640 .await?;
641 tx.commit().await?;
642 link
643 }
644 };
645
646 if link.is_none() {
647 return Err(AppError::NotFound);
648 }
649
650 let sources = match &state.db {
651 DatabasePool::Postgres(pool) => {
652 sqlx::query_as::<Postgres, SourceStats>(
653 r#"
654 SELECT
655 query_source as "source!",
656 COUNT(*)::bigint as "count!"
657 FROM clicks
658 WHERE link_id = $1
659 AND query_source IS NOT NULL
660 AND query_source != ''
661 GROUP BY query_source
662 ORDER BY COUNT(*) DESC
663 LIMIT 10
664 "#,
665 )
666 .bind(link_id)
667 .fetch_all(pool)
668 .await?
669 }
670 DatabasePool::Sqlite(pool) => {
671 sqlx::query_as::<Sqlite, SourceStats>(
672 r#"
673 SELECT
674 query_source as "source!",
675 COUNT(*) as "count!"
676 FROM clicks
677 WHERE link_id = ?
678 AND query_source IS NOT NULL
679 AND query_source != ''
680 GROUP BY query_source
681 ORDER BY COUNT(*) DESC
682 LIMIT 10
683 "#,
684 )
685 .bind(link_id)
686 .fetch_all(pool)
687 .await?
688 }
689 };
690
691 Ok(HttpResponse::Ok().json(sources))
692}