use crate::services::qr::{generate_qr_png, generate_qr_svg}; use crate::state::AppState; use crate::utils::get_client_ip; use axum::{ extract::{ConnectInfo, Path, State}, http::{HeaderMap, StatusCode}, response::{IntoResponse, Response}, }; use std::net::SocketAddr; use serde_json::json; // GET /api/qr/:file (e.g. /api/qr/abcdef.png or /api/qr/abcdef.svg or JSON stats /api/qr/abcdef) pub async fn qr_handler( State(state): State, headers: HeaderMap, connect_info: Option>, Path(file): Path, ) -> Response { let parts: Vec<&str> = file.split('.').collect(); if parts.len() != 2 { // No extension: this is a JSON stats request! let auth_header = headers.get("Authorization").and_then(|h| h.to_str().ok()); let authenticated = if let Some(auth) = auth_header { let admin_conn = state.admin_db.lock().unwrap(); let users_conn = state.users_db.lock().unwrap(); matches!( crate::auth::session::authenticate_api_key(&admin_conn, &users_conn, auth), Ok(Some(_user)) ) } else { false }; if !authenticated { return (StatusCode::UNAUTHORIZED, "Unauthorized").into_response(); } // We need to look up owner_user_id, target_id, and status from global_slugs let (owner_user_id, target_id, slug_status) = { let system_conn = state.system_db.lock().unwrap(); let mut stmt = match system_conn.prepare( "SELECT owner_user_id, target_id, status FROM global_slugs WHERE slug = ?1;", ) { Ok(s) => s, Err(_) => { return (StatusCode::INTERNAL_SERVER_ERROR, "Database error").into_response() } }; use rusqlite::OptionalExtension; match stmt .query_row(rusqlite::params![&file], |row| { Ok(( row.get::<_, i64>(0)?, row.get::<_, String>(1)?, row.get::<_, String>(2)?, )) }) .optional() { Ok(Some((uid, tid, status))) => (uid, tid, status), Ok(None) => return (StatusCode::NOT_FOUND, "URL not found").into_response(), Err(_) => { return (StatusCode::INTERNAL_SERVER_ERROR, "Database error").into_response() } } }; if slug_status == "disabled" { return (StatusCode::GONE, "This content has been disabled").into_response(); } else if slug_status != "active" { return (StatusCode::NOT_FOUND, "URL not found").into_response(); } let user_dbs = match state.get_user_dbs(owner_user_id) { Ok(dbs) => dbs, Err(_) => return (StatusCode::INTERNAL_SERVER_ERROR, "Database error").into_response(), }; let qr_scans = { let conn = user_dbs.analytics.lock().unwrap(); crate::db::qr::get_qr_scan_count(&conn, &target_id).unwrap_or(0) }; let direct_clicks = { let conn = user_dbs.analytics.lock().unwrap(); conn.query_row( "SELECT COUNT(*) FROM visits WHERE target_id = ?1;", rusqlite::params![target_id], |row| row.get(0), ) .unwrap_or(0) }; return axum::response::Json(json!({ "direct_clicks": direct_clicks, "qr_scans": qr_scans })) .into_response(); } let code = parts[0]; let ext = parts[1].to_lowercase(); if !crate::utils::validation::validate_redirect_code(code) && !crate::utils::validation::validate_page_code(code) { return (StatusCode::NOT_FOUND, "Not Found").into_response(); } // We need to look up owner_user_id, target_type, target_id, and status from global_slugs let (owner_user_id, target_type, target_id, slug_status) = { let system_conn = state.system_db.lock().unwrap(); let mut stmt = match system_conn .prepare("SELECT owner_user_id, target_type, target_id, status FROM global_slugs WHERE slug = ?1;") { Ok(s) => s, Err(_) => return (StatusCode::INTERNAL_SERVER_ERROR, "Database error").into_response(), }; use rusqlite::OptionalExtension; match stmt .query_row(rusqlite::params![code], |row| { Ok(( row.get::<_, i64>(0)?, row.get::<_, String>(1)?, row.get::<_, String>(2)?, row.get::<_, String>(3)?, )) }) .optional() { Ok(Some(info)) => info, Ok(None) => return (StatusCode::NOT_FOUND, "Not Found").into_response(), Err(_) => return (StatusCode::INTERNAL_SERVER_ERROR, "Database error").into_response(), } }; if slug_status == "disabled" { return (StatusCode::GONE, "This content has been disabled").into_response(); } else if slug_status != "active" { return (StatusCode::NOT_FOUND, "Not Found").into_response(); } // Construct public base URL let proto = if state.config.cookie_secure { "https" } else { "http" }; let host_header = headers .get("host") .and_then(|h| h.to_str().ok()) .unwrap_or("localhost:8654"); let base_url = state .config .base_url .clone() .unwrap_or_else(|| format!("{}://{}", proto, host_header)); let full_url = if target_type == "page" { format!("{}/p/{}", base_url.trim_end_matches('/'), code) } else { format!("{}/{}", base_url.trim_end_matches('/'), code) }; // Generate QR code based on format let (body, content_type) = if ext == "svg" { match generate_qr_svg(&full_url) { Ok(svg) => (svg.into_bytes(), "image/svg+xml"), Err(e) => { return ( StatusCode::INTERNAL_SERVER_ERROR, format!("QR generation error: {}", e), ) .into_response() } } } else if ext == "png" { match generate_qr_png(&full_url, 256) { Ok(png) => (png, "image/png"), Err(e) => { return ( StatusCode::INTERNAL_SERVER_ERROR, format!("QR generation error: {}", e), ) .into_response() } } } else { return ( StatusCode::BAD_REQUEST, "Unsupported format. Use .png or .svg", ) .into_response(); }; // Log the QR access event in a try-catch style let _ = { let ip = get_client_ip(&headers, connect_info); let user_agent = headers .get("user-agent") .and_then(|h| h.to_str().ok()) .map(|s| s.to_string()); if let Ok(user_dbs) = state.get_user_dbs(owner_user_id) { if let Ok(analytics_conn) = user_dbs.analytics.lock() { let _ = crate::db::qr::log_qr_access( &analytics_conn, &target_id, Some(ip.as_str()), user_agent.as_deref(), ); } } }; Response::builder() .header("content-type", content_type) .header("cache-control", "public, max-age=86400") // cache for 1 day .body(axum::body::Body::from(body)) .unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response()) }