allow creation of user and links via args and env

Author: WaveringAna

src/auth.rs

@@ -1,7 +1,7 @@
+use crate::{error::AppError, models::Claims};
 use actix_web::{dev::Payload, FromRequest, HttpRequest};
 use jsonwebtoken::{decode, DecodingKey, Validation};
 use std::future::{ready, Ready};
-use crate::{error::AppError, models::Claims};
 
 pub struct AuthenticatedUser {
     pub user_id: i32,
@@ -12,19 +12,20 @@ impl FromRequest for AuthenticatedUser {
     type Future = Ready<Result<Self, Self::Error>>;
 
     fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
-        let auth_header = req.headers()
+        let auth_header = req
+            .headers()
             .get("Authorization")
             .and_then(|h| h.to_str().ok());
 
         if let Some(auth_header) = auth_header {
             if auth_header.starts_with("Bearer ") {
                 let token = &auth_header[7..];
-                let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
-                
+                let secret =
+                    std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
                 match decode::<Claims>(
                     token,
                     &DecodingKey::from_secret(secret.as_bytes()),
-                    &Validation::default()
+                    &Validation::default(),
                 ) {
                     Ok(token_data) => {
                         return ready(Ok(AuthenticatedUser {
@@ -35,7 +36,7 @@ impl FromRequest for AuthenticatedUser {
                 }
             }
         }
-        
         ready(Err(AppError::Unauthorized))
     }
-}
+}
+

src/handlers.rs

@@ -131,7 +131,7 @@ fn validate_custom_code(code: &str) -> Result<(), AppError> {
     Ok(())
 }
 
-fn validate_url(url: &String) -> Result<(), AppError> {
+fn validate_url(url: &str) -> Result<(), AppError> {
     if url.is_empty() {
         return Err(AppError::InvalidInput("URL cannot be empty".to_string()));
     }

src/main.rs

@@ -1,8 +1,10 @@
 use actix_cors::Cors;
 use actix_web::{web, App, HttpResponse, HttpServer};
 use anyhow::Result;
+use clap::Parser;
 use rust_embed::RustEmbed;
 use simplelink::check_and_generate_admin_token;
+use simplelink::models::DatabasePool;
 use simplelink::{create_db_pool, run_migrations};
 use simplelink::{handlers, AppState};
 use sqlx::{Postgres, Sqlite};
@@ -26,6 +28,106 @@ async fn serve_static_file(path: &str) -> HttpResponse {
     }
 }
 
+async fn create_initial_links(pool: &DatabasePool) -> Result<()> {
+    if let Ok(links) = std::env::var("INITIAL_LINKS") {
+        for link_entry in links.split(';') {
+            let parts: Vec<&str> = link_entry.split(',').collect();
+            if parts.len() >= 2 {
+                let url = parts[0];
+                let code = parts[1];
+
+                match pool {
+                    DatabasePool::Postgres(pool) => {
+                        sqlx::query(
+                            "INSERT INTO links (original_url, short_code, user_id)  
+                             VALUES ($1, $2, $3) 
+                             ON CONFLICT (short_code) 
+                             DO UPDATE SET short_code = EXCLUDED.short_code
+                             WHERE links.original_url = EXCLUDED.original_url",
+                        )
+                        .bind(url)
+                        .bind(code)
+                        .bind(1)
+                        .execute(pool)
+                        .await?;
+                    }
+                    DatabasePool::Sqlite(pool) => {
+                        // First check if the exact combination exists
+                        let exists = sqlx::query_scalar::<_, bool>(
+                            "SELECT EXISTS(
+                                SELECT 1 FROM links 
+                                WHERE original_url = ?1 
+                                AND short_code = ?2
+                            )",
+                        )
+                        .bind(url)
+                        .bind(code)
+                        .fetch_one(pool)
+                        .await?;
+
+                        // Only insert if the exact combination doesn't exist
+                        if !exists {
+                            sqlx::query(
+                                "INSERT INTO links (original_url, short_code, user_id) 
+                                 VALUES (?1, ?2, ?3)",
+                            )
+                            .bind(url)
+                            .bind(code)
+                            .bind(1)
+                            .execute(pool)
+                            .await?;
+                            info!("Created initial link: {} -> {} for user_id: 1", code, url);
+                        } else {
+                            info!("Skipped existing link: {} -> {} for user_id: 1", code, url);
+                        }
+                    }
+                }
+            }
+        }
+    }
+    Ok(())
+}
+
+async fn create_admin_user(pool: &DatabasePool, email: &str, password: &str) -> Result<()> {
+    use argon2::{
+        password_hash::{rand_core::OsRng, SaltString},
+        Argon2, PasswordHasher,
+    };
+
+    let salt = SaltString::generate(&mut OsRng);
+    let argon2 = Argon2::default();
+    let password_hash = argon2
+        .hash_password(password.as_bytes(), &salt)
+        .map_err(|e| anyhow::anyhow!("Password hashing error: {}", e))?
+        .to_string();
+
+    match pool {
+        DatabasePool::Postgres(pool) => {
+            sqlx::query(
+                "INSERT INTO users (email, password_hash) 
+                 VALUES ($1, $2) 
+                 ON CONFLICT (email) DO NOTHING",
+            )
+            .bind(email)
+            .bind(&password_hash)
+            .execute(pool)
+            .await?;
+        }
+        DatabasePool::Sqlite(pool) => {
+            sqlx::query(
+                "INSERT OR IGNORE INTO users (email, password_hash) 
+                 VALUES (?1, ?2)",
+            )
+            .bind(email)
+            .bind(&password_hash)
+            .execute(pool)
+            .await?;
+        }
+    }
+    info!("Created admin user: {}", email);
+    Ok(())
+}
+
 #[actix_web::main]
 async fn main() -> Result<()> {
     // Load environment variables from .env file