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, CreateLink, Link, LoginRequest, RegisterRequest, User, UserResponse,
6 },
7 AppState,
8};
9use actix_web::{web, HttpRequest, HttpResponse, Responder};
10use argon2::{
11 password_hash::{rand_core::OsRng, SaltString},
12 PasswordVerifier,
13};
14use argon2::{Argon2, PasswordHash, PasswordHasher};
15use jsonwebtoken::{encode, EncodingKey, Header};
16use lazy_static::lazy_static;
17use regex::Regex;
18
19lazy_static! {
20 static ref VALID_CODE_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_-]{1,32}$").unwrap();
21}
22
23pub async fn create_short_url(
24 state: web::Data<AppState>,
25 user: AuthenticatedUser,
26 payload: web::Json<CreateLink>,
27) -> Result<impl Responder, AppError> {
28 tracing::debug!("Creating short URL with user_id: {}", user.user_id);
29
30 validate_url(&payload.url)?;
31
32 let short_code = if let Some(ref custom_code) = payload.custom_code {
33 validate_custom_code(custom_code)?;
34
35 tracing::debug!("Checking if custom code {} exists", custom_code);
36 // Check if code is already taken
37 if let Some(_) = sqlx::query_as::<_, Link>("SELECT * FROM links WHERE short_code = $1")
38 .bind(custom_code)
39 .fetch_optional(&state.db)
40 .await?
41 {
42 return Err(AppError::InvalidInput(
43 "Custom code already taken".to_string(),
44 ));
45 }
46
47 custom_code.clone()
48 } else {
49 generate_short_code()
50 };
51
52 // Start transaction
53 let mut tx = state.db.begin().await?;
54
55 tracing::debug!("Inserting new link with short_code: {}", short_code);
56 let link = sqlx::query_as::<_, Link>(
57 "INSERT INTO links (original_url, short_code, user_id) VALUES ($1, $2, $3) RETURNING *",
58 )
59 .bind(&payload.url)
60 .bind(&short_code)
61 .bind(user.user_id)
62 .fetch_one(&mut *tx)
63 .await?;
64
65 if let Some(ref source) = payload.source {
66 tracing::debug!("Adding click source: {}", source);
67 sqlx::query("INSERT INTO clicks (link_id, source) VALUES ($1, $2)")
68 .bind(link.id)
69 .bind(source)
70 .execute(&mut *tx)
71 .await?;
72 }
73
74 tx.commit().await?;
75 Ok(HttpResponse::Created().json(link))
76}
77
78fn validate_custom_code(code: &str) -> Result<(), AppError> {
79 if !VALID_CODE_REGEX.is_match(code) {
80 return Err(AppError::InvalidInput(
81 "Custom code must be 1-32 characters long and contain only letters, numbers, underscores, and hyphens".to_string()
82 ));
83 }
84
85 // Add reserved words check
86 let reserved_words = ["api", "health", "admin", "static", "assets"];
87 if reserved_words.contains(&code.to_lowercase().as_str()) {
88 return Err(AppError::InvalidInput(
89 "This code is reserved and cannot be used".to_string(),
90 ));
91 }
92
93 Ok(())
94}
95
96fn validate_url(url: &String) -> Result<(), AppError> {
97 if url.is_empty() {
98 return Err(AppError::InvalidInput("URL cannot be empty".to_string()));
99 }
100 if !url.starts_with("http://") && !url.starts_with("https://") {
101 return Err(AppError::InvalidInput(
102 "URL must start with http:// or https://".to_string(),
103 ));
104 }
105 Ok(())
106}
107
108pub async fn redirect_to_url(
109 state: web::Data<AppState>,
110 path: web::Path<String>,
111 req: HttpRequest,
112) -> Result<impl Responder, AppError> {
113 let short_code = path.into_inner();
114
115 // Extract query source if present
116 let query_source = req
117 .uri()
118 .query()
119 .and_then(|q| web::Query::<std::collections::HashMap<String, String>>::from_query(q).ok())
120 .and_then(|params| params.get("source").cloned());
121
122 let mut tx = state.db.begin().await?;
123
124 let link = sqlx::query_as::<_, Link>(
125 "UPDATE links SET clicks = clicks + 1 WHERE short_code = $1 RETURNING *",
126 )
127 .bind(&short_code)
128 .fetch_optional(&mut *tx)
129 .await?;
130
131 match link {
132 Some(link) => {
133 // Record click with both user agent and query source
134 let user_agent = req
135 .headers()
136 .get("user-agent")
137 .and_then(|h| h.to_str().ok())
138 .unwrap_or("unknown")
139 .to_string();
140
141 sqlx::query("INSERT INTO clicks (link_id, source, query_source) VALUES ($1, $2, $3)")
142 .bind(link.id)
143 .bind(user_agent)
144 .bind(query_source)
145 .execute(&mut *tx)
146 .await?;
147
148 tx.commit().await?;
149
150 Ok(HttpResponse::TemporaryRedirect()
151 .append_header(("Location", link.original_url))
152 .finish())
153 }
154 None => Err(AppError::NotFound),
155 }
156}
157
158pub async fn get_all_links(
159 state: web::Data<AppState>,
160 user: AuthenticatedUser,
161) -> Result<impl Responder, AppError> {
162 let links = sqlx::query_as::<_, Link>(
163 "SELECT * FROM links WHERE user_id = $1 ORDER BY created_at DESC",
164 )
165 .bind(user.user_id)
166 .fetch_all(&state.db)
167 .await?;
168
169 Ok(HttpResponse::Ok().json(links))
170}
171
172pub async fn health_check(state: web::Data<AppState>) -> impl Responder {
173 match sqlx::query("SELECT 1").execute(&state.db).await {
174 Ok(_) => HttpResponse::Ok().json("Healthy"),
175 Err(_) => HttpResponse::ServiceUnavailable().json("Database unavailable"),
176 }
177}
178
179fn generate_short_code() -> String {
180 use base62::encode;
181 use uuid::Uuid;
182
183 let uuid = Uuid::new_v4();
184 encode(uuid.as_u128() as u64).chars().take(8).collect()
185}
186
187pub async fn register(
188 state: web::Data<AppState>,
189 payload: web::Json<RegisterRequest>,
190) -> Result<impl Responder, AppError> {
191 let exists = sqlx::query!("SELECT id FROM users WHERE email = $1", payload.email)
192 .fetch_optional(&state.db)
193 .await?;
194
195 if exists.is_some() {
196 return Err(AppError::Auth("Email already registered".to_string()));
197 }
198
199 let salt = SaltString::generate(&mut OsRng);
200 let argon2 = Argon2::default();
201 let password_hash = argon2
202 .hash_password(payload.password.as_bytes(), &salt)
203 .map_err(|e| AppError::Auth(e.to_string()))?
204 .to_string();
205
206 let user = sqlx::query_as!(
207 User,
208 "INSERT INTO users (email, password_hash) VALUES ($1, $2) RETURNING *",
209 payload.email,
210 password_hash
211 )
212 .fetch_one(&state.db)
213 .await?;
214
215 let claims = Claims::new(user.id);
216 let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
217 let token = encode(
218 &Header::default(),
219 &claims,
220 &EncodingKey::from_secret(secret.as_bytes()),
221 )
222 .map_err(|e| AppError::Auth(e.to_string()))?;
223
224 Ok(HttpResponse::Ok().json(AuthResponse {
225 token,
226 user: UserResponse {
227 id: user.id,
228 email: user.email,
229 },
230 }))
231}
232
233pub async fn login(
234 state: web::Data<AppState>,
235 payload: web::Json<LoginRequest>,
236) -> Result<impl Responder, AppError> {
237 let user = sqlx::query_as!(User, "SELECT * FROM users WHERE email = $1", payload.email)
238 .fetch_optional(&state.db)
239 .await?
240 .ok_or_else(|| AppError::Auth("Invalid credentials".to_string()))?;
241
242 let argon2 = Argon2::default();
243 let parsed_hash =
244 PasswordHash::new(&user.password_hash).map_err(|e| AppError::Auth(e.to_string()))?;
245
246 if argon2
247 .verify_password(payload.password.as_bytes(), &parsed_hash)
248 .is_err()
249 {
250 return Err(AppError::Auth("Invalid credentials".to_string()));
251 }
252
253 let claims = Claims::new(user.id);
254 let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
255 let token = encode(
256 &Header::default(),
257 &claims,
258 &EncodingKey::from_secret(secret.as_bytes()),
259 )
260 .map_err(|e| AppError::Auth(e.to_string()))?;
261
262 Ok(HttpResponse::Ok().json(AuthResponse {
263 token,
264 user: UserResponse {
265 id: user.id,
266 email: user.email,
267 },
268 }))
269}
270
271pub async fn delete_link(
272 state: web::Data<AppState>,
273 user: AuthenticatedUser,
274 path: web::Path<i32>,
275) -> Result<impl Responder, AppError> {
276 let link_id = path.into_inner();
277
278 // Start transaction
279 let mut tx = state.db.begin().await?;
280
281 // Verify the link belongs to the user
282 let link = sqlx::query!(
283 "SELECT id FROM links WHERE id = $1 AND user_id = $2",
284 link_id,
285 user.user_id
286 )
287 .fetch_optional(&mut *tx)
288 .await?;
289
290 if link.is_none() {
291 return Err(AppError::NotFound);
292 }
293
294 // Delete associated clicks first due to foreign key constraint
295 sqlx::query!("DELETE FROM clicks WHERE link_id = $1", link_id)
296 .execute(&mut *tx)
297 .await?;
298
299 // Delete the link
300 sqlx::query!("DELETE FROM links WHERE id = $1", link_id)
301 .execute(&mut *tx)
302 .await?;
303
304 tx.commit().await?;
305
306 Ok(HttpResponse::NoContent().finish())
307}