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 let exists = sqlx::query!("SELECT id FROM users WHERE email = $1", payload.email)
193 .fetch_optional(&state.db)
194 .await?;
195
196 if exists.is_some() {
197 return Err(AppError::Auth("Email already registered".to_string()));
198 }
199
200 let salt = SaltString::generate(&mut OsRng);
201 let argon2 = Argon2::default();
202 let password_hash = argon2
203 .hash_password(payload.password.as_bytes(), &salt)
204 .map_err(|e| AppError::Auth(e.to_string()))?
205 .to_string();
206
207 let user = sqlx::query_as!(
208 User,
209 "INSERT INTO users (email, password_hash) VALUES ($1, $2) RETURNING *",
210 payload.email,
211 password_hash
212 )
213 .fetch_one(&state.db)
214 .await?;
215
216 let claims = Claims::new(user.id);
217 let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
218 let token = encode(
219 &Header::default(),
220 &claims,
221 &EncodingKey::from_secret(secret.as_bytes()),
222 )
223 .map_err(|e| AppError::Auth(e.to_string()))?;
224
225 Ok(HttpResponse::Ok().json(AuthResponse {
226 token,
227 user: UserResponse {
228 id: user.id,
229 email: user.email,
230 },
231 }))
232}
233
234pub async fn login(
235 state: web::Data<AppState>,
236 payload: web::Json<LoginRequest>,
237) -> Result<impl Responder, AppError> {
238 let user = sqlx::query_as!(User, "SELECT * FROM users WHERE email = $1", payload.email)
239 .fetch_optional(&state.db)
240 .await?
241 .ok_or_else(|| AppError::Auth("Invalid credentials".to_string()))?;
242
243 let argon2 = Argon2::default();
244 let parsed_hash =
245 PasswordHash::new(&user.password_hash).map_err(|e| AppError::Auth(e.to_string()))?;
246
247 if argon2
248 .verify_password(payload.password.as_bytes(), &parsed_hash)
249 .is_err()
250 {
251 return Err(AppError::Auth("Invalid credentials".to_string()));
252 }
253
254 let claims = Claims::new(user.id);
255 let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
256 let token = encode(
257 &Header::default(),
258 &claims,
259 &EncodingKey::from_secret(secret.as_bytes()),
260 )
261 .map_err(|e| AppError::Auth(e.to_string()))?;
262
263 Ok(HttpResponse::Ok().json(AuthResponse {
264 token,
265 user: UserResponse {
266 id: user.id,
267 email: user.email,
268 },
269 }))
270}
271
272pub async fn delete_link(
273 state: web::Data<AppState>,
274 user: AuthenticatedUser,
275 path: web::Path<i32>,
276) -> Result<impl Responder, AppError> {
277 let link_id = path.into_inner();
278
279 // Start transaction
280 let mut tx = state.db.begin().await?;
281
282 // Verify the link belongs to the user
283 let link = sqlx::query!(
284 "SELECT id FROM links WHERE id = $1 AND user_id = $2",
285 link_id,
286 user.user_id
287 )
288 .fetch_optional(&mut *tx)
289 .await?;
290
291 if link.is_none() {
292 return Err(AppError::NotFound);
293 }
294
295 // Delete associated clicks first due to foreign key constraint
296 sqlx::query!("DELETE FROM clicks WHERE link_id = $1", link_id)
297 .execute(&mut *tx)
298 .await?;
299
300 // Delete the link
301 sqlx::query!("DELETE FROM links WHERE id = $1", link_id)
302 .execute(&mut *tx)
303 .await?;
304
305 tx.commit().await?;
306
307 Ok(HttpResponse::NoContent().finish())
308}
309
310pub async fn get_link_clicks(
311 state: web::Data<AppState>,
312 user: AuthenticatedUser,
313 path: web::Path<i32>,
314) -> Result<impl Responder, AppError> {
315 let link_id = path.into_inner();
316
317 // Verify the link belongs to the user
318 let link = sqlx::query!(
319 "SELECT id FROM links WHERE id = $1 AND user_id = $2",
320 link_id,
321 user.user_id
322 )
323 .fetch_optional(&state.db)
324 .await?;
325
326 if link.is_none() {
327 return Err(AppError::NotFound);
328 }
329
330 let clicks = sqlx::query_as!(
331 ClickStats,
332 r#"
333 SELECT
334 DATE(created_at)::date as "date!",
335 COUNT(*)::bigint as "clicks!"
336 FROM clicks
337 WHERE link_id = $1
338 GROUP BY DATE(created_at)
339 ORDER BY DATE(created_at) ASC -- Changed from DESC to ASC
340 LIMIT 30
341 "#,
342 link_id
343 )
344 .fetch_all(&state.db)
345 .await?;
346
347 Ok(HttpResponse::Ok().json(clicks))
348}
349
350pub async fn get_link_sources(
351 state: web::Data<AppState>,
352 user: AuthenticatedUser,
353 path: web::Path<i32>,
354) -> Result<impl Responder, AppError> {
355 let link_id = path.into_inner();
356
357 // Verify the link belongs to the user
358 let link = sqlx::query!(
359 "SELECT id FROM links WHERE id = $1 AND user_id = $2",
360 link_id,
361 user.user_id
362 )
363 .fetch_optional(&state.db)
364 .await?;
365
366 if link.is_none() {
367 return Err(AppError::NotFound);
368 }
369
370 let sources = sqlx::query_as!(
371 SourceStats,
372 r#"
373 SELECT
374 query_source as "source!",
375 COUNT(*)::bigint as "count!"
376 FROM clicks
377 WHERE link_id = $1
378 AND query_source IS NOT NULL
379 AND query_source != ''
380 GROUP BY query_source
381 ORDER BY COUNT(*) DESC
382 LIMIT 10
383 "#,
384 link_id
385 )
386 .fetch_all(&state.db)
387 .await?;
388
389 Ok(HttpResponse::Ok().json(sources))
390}