···
1
-
use actix_web::{web, HttpResponse, Responder, HttpRequest};
2
-
use jsonwebtoken::{encode, Header, EncodingKey};use crate::{error::AppError, models::{AuthResponse, Claims, CreateLink, Link, LoginRequest, RegisterRequest, User, UserResponse}, AppState};
4
-
use argon2::{password_hash::{rand_core::OsRng, SaltString}, PasswordVerifier};
5
-
use lazy_static::lazy_static;
6
-
use argon2::{Argon2, PasswordHash, PasswordHasher};
use crate::auth::AuthenticatedUser;
5
+
AuthResponse, Claims, CreateLink, Link, LoginRequest, RegisterRequest, User, UserResponse,
9
+
use actix_web::{web, HttpRequest, HttpResponse, Responder};
11
+
password_hash::{rand_core::OsRng, SaltString},
14
+
use argon2::{Argon2, PasswordHash, PasswordHasher};
15
+
use jsonwebtoken::{encode, EncodingKey, Header};
16
+
use lazy_static::lazy_static;
static ref VALID_CODE_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_-]{1,32}$").unwrap();
···
payload: web::Json<CreateLink>,
) -> Result<impl Responder, AppError> {
tracing::debug!("Creating short URL with user_id: {}", user.user_id);
validate_url(&payload.url)?;
let short_code = if let Some(ref custom_code) = payload.custom_code {
validate_custom_code(custom_code)?;
tracing::debug!("Checking if custom code {} exists", custom_code);
// Check if code is already taken
27
-
if let Some(_) = sqlx::query_as::<_, Link>(
28
-
"SELECT * FROM links WHERE short_code = $1"
31
-
.fetch_optional(&state.db)
37
+
if let Some(_) = sqlx::query_as::<_, Link>("SELECT * FROM links WHERE short_code = $1")
39
+
.fetch_optional(&state.db)
return Err(AppError::InvalidInput(
34
-
"Custom code already taken".to_string()
43
+
"Custom code already taken".to_string(),
let mut tx = state.db.begin().await?;
tracing::debug!("Inserting new link with short_code: {}", short_code);
let link = sqlx::query_as::<_, Link>(
48
-
"INSERT INTO links (original_url, short_code, user_id) VALUES ($1, $2, $3) RETURNING *"
57
+
"INSERT INTO links (original_url, short_code, user_id) VALUES ($1, $2, $3) RETURNING *",
if let Some(ref source) = payload.source {
tracing::debug!("Adding click source: {}", source);
59
-
"INSERT INTO clicks (link_id, source) VALUES ($1, $2)"
67
+
sqlx::query("INSERT INTO clicks (link_id, source) VALUES ($1, $2)")
Ok(HttpResponse::Created().json(link))
···
"Custom code must be 1-32 characters long and contain only letters, numbers, underscores, and hyphens".to_string()
// Add reserved words check
let reserved_words = ["api", "health", "admin", "static", "assets"];
if reserved_words.contains(&code.to_lowercase().as_str()) {
return Err(AppError::InvalidInput(
82
-
"This code is reserved and cannot be used".to_string()
89
+
"This code is reserved and cannot be used".to_string(),
···
return Err(AppError::InvalidInput("URL cannot be empty".to_string()));
if !url.starts_with("http://") && !url.starts_with("https://") {
94
-
return Err(AppError::InvalidInput("URL must start with http:// or https://".to_string()));
101
+
return Err(AppError::InvalidInput(
102
+
"URL must start with http:// or https://".to_string(),
···
) -> Result<impl Responder, AppError> {
let short_code = path.into_inner();
// Extract query source if present
107
-
let query_source = req.uri()
116
+
let query_source = req
.and_then(|q| web::Query::<std::collections::HashMap<String, String>>::from_query(q).ok())
.and_then(|params| params.get("source").cloned());
···
let mut tx = state.db.begin().await?;
let link = sqlx::query_as::<_, Link>(
115
-
"UPDATE links SET clicks = clicks + 1 WHERE short_code = $1 RETURNING *"
125
+
"UPDATE links SET clicks = clicks + 1 WHERE short_code = $1 RETURNING *",
.fetch_optional(&mut *tx)
···
// Record click with both user agent and query source
124
-
let user_agent = req.headers()
134
+
let user_agent = req
.and_then(|h| h.to_str().ok())
131
-
"INSERT INTO clicks (link_id, source, query_source) VALUES ($1, $2, $3)"
135
-
.bind(query_source)
141
+
sqlx::query("INSERT INTO clicks (link_id, source, query_source) VALUES ($1, $2, $3)")
144
+
.bind(query_source)
Ok(HttpResponse::TemporaryRedirect()
.append_header(("Location", link.original_url))
None => Err(AppError::NotFound),
···
) -> Result<impl Responder, AppError> {
let links = sqlx::query_as::<_, Link>(
154
-
"SELECT * FROM links WHERE user_id = $1 ORDER BY created_at DESC"
163
+
"SELECT * FROM links WHERE user_id = $1 ORDER BY created_at DESC",
···
Ok(HttpResponse::Ok().json(links))
163
-
pub async fn health_check(
164
-
state: web::Data<AppState>,
165
-
) -> impl Responder {
172
+
pub async fn health_check(state: web::Data<AppState>) -> impl Responder {
match sqlx::query("SELECT 1").execute(&state.db).await {
Ok(_) => HttpResponse::Ok().json("Healthy"),
Err(_) => HttpResponse::ServiceUnavailable().json("Database unavailable"),
···
fn generate_short_code() -> String {
let uuid = Uuid::new_v4();
encode(uuid.as_u128() as u64).chars().take(8).collect()
···
state: web::Data<AppState>,
payload: web::Json<RegisterRequest>,
) -> Result<impl Responder, AppError> {
184
-
let exists = sqlx::query!(
185
-
"SELECT id FROM users WHERE email = $1",
188
-
.fetch_optional(&state.db)
191
+
let exists = sqlx::query!("SELECT id FROM users WHERE email = $1", payload.email)
192
+
.fetch_optional(&state.db)
return Err(AppError::Auth("Email already registered".to_string()));
···
let salt = SaltString::generate(&mut OsRng);
let argon2 = Argon2::default();
197
-
let password_hash = argon2.hash_password(payload.password.as_bytes(), &salt)
201
+
let password_hash = argon2
202
+
.hash_password(payload.password.as_bytes(), &salt)
.map_err(|e| AppError::Auth(e.to_string()))?
···
215
-
&EncodingKey::from_secret(secret.as_bytes())
216
-
).map_err(|e| AppError::Auth(e.to_string()))?;
220
+
&EncodingKey::from_secret(secret.as_bytes()),
222
+
.map_err(|e| AppError::Auth(e.to_string()))?;
Ok(HttpResponse::Ok().json(AuthResponse {
···
state: web::Data<AppState>,
payload: web::Json<LoginRequest>,
) -> Result<impl Responder, AppError> {
231
-
let user = sqlx::query_as!(
233
-
"SELECT * FROM users WHERE email = $1",
236
-
.fetch_optional(&state.db)
238
-
.ok_or_else(|| AppError::Auth("Invalid credentials".to_string()))?;
237
+
let user = sqlx::query_as!(User, "SELECT * FROM users WHERE email = $1", payload.email)
238
+
.fetch_optional(&state.db)
240
+
.ok_or_else(|| AppError::Auth("Invalid credentials".to_string()))?;
let argon2 = Argon2::default();
241
-
let parsed_hash = PasswordHash::new(&user.password_hash)
242
-
.map_err(|e| AppError::Auth(e.to_string()))?;
244
+
PasswordHash::new(&user.password_hash).map_err(|e| AppError::Auth(e.to_string()))?;
244
-
if argon2.verify_password(payload.password.as_bytes(), &parsed_hash).is_err() {
247
+
.verify_password(payload.password.as_bytes(), &parsed_hash)
return Err(AppError::Auth("Invalid credentials".to_string()));
···
253
-
&EncodingKey::from_secret(secret.as_bytes())
254
-
).map_err(|e| AppError::Auth(e.to_string()))?;
258
+
&EncodingKey::from_secret(secret.as_bytes()),
260
+
.map_err(|e| AppError::Auth(e.to_string()))?;
Ok(HttpResponse::Ok().json(AuthResponse {
···
271
+
pub 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();
278
+
// Start transaction
279
+
let mut tx = state.db.begin().await?;
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",
287
+
.fetch_optional(&mut *tx)
290
+
if link.is_none() {
291
+
return Err(AppError::NotFound);
294
+
// Delete associated clicks first due to foreign key constraint
295
+
sqlx::query!("DELETE FROM clicks WHERE link_id = $1", link_id)
300
+
sqlx::query!("DELETE FROM links WHERE id = $1", link_id)
304
+
tx.commit().await?;
306
+
Ok(HttpResponse::NoContent().finish())
nekomimi.pet