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