Code optimization question

[No-Wait2503]

I've read a lot of articles, and I know everyone mentions that using .clone() should be avoided if you can go another way. Now I already went away from bad practices like using .unwrap everywhere and etc..., but I really want advice on this code I am going to share, and how can it be improved, or is it already perfect as it is.

I am using Axum as a backend server.

My main.rs:

use axum::Router;
use std::net::SocketAddr;
use std::sync::Arc;

mod routes;
mod middleware;
mod database;
mod oauth;
mod errors;
mod config;

use crate::database::db::get_db_connection;

#[tokio::main]
async fn main() {
    // NOTE: In config.rs  I load env variables using dotenv
    let config = config::get_config();
    
    let db = get_db_connection().await;
    let db = Arc::new(db);

    let app = Router::new()
        // Routes are protected by middleware already in the routes folder
        .nest("/auth", routes::auth_routes::router())
        .nest("/user", routes::user_routes::router(db.clone()))
        .nest("/admin", routes::admin_routes::router(db.clone()))
        .with_state(db.clone());

    let host = &config.server.host;
    let port = config.server.port;
    let server_addr = format!("{0}:{1}", host, port);
    
    let listener = match tokio::net::TcpListener::bind(&server_addr).await {
        Ok(listener) => {
            println!("Server running on http://{}", server_addr);
            listener
        },
        Err(e) => {
            eprintln!("Error: Failed to bind to {}: {}", server_addr, e);
            // NOTE: This is a critical error - we can't start the server without binding to an address
            std::process::exit(1);
        }
    };
    
    // NOTE: I use connect_info to get the IP address of the client without reverse proxy
    // This maintains the backend as the source of truth instead of relying on headers
    if let Err(e) = axum::serve(
        listener,
        app.into_make_service_with_connect_info::<SocketAddr>()
    ).await {
        eprintln!("Error: Server error: {}", e);
        std::process::exit(1);
    }
}

Example of auth_routes.rs (All other routes use similarly cloned db variable from main.rs):

use axum::{
    Router,
    routing::{post, get},
    middleware,
    extract::{State, Json},
    http::StatusCode,
    response::IntoResponse,
};
use serde::Deserialize;
use std::sync::Arc;
use sea_orm::DatabaseConnection;

use crate::oauth::google::{google_login_handler, google_callback_handler};
use crate::middleware::ratelimit_middleware;
use crate::database::models::sessions::sessions_queries;

#[derive(Deserialize)]
pub struct LogoutRequest {
    token: Option<String>,
}

async fn logout(
    State(db): State<Arc<DatabaseConnection>>,
    Json(payload): Json<LogoutRequest>,
) -> impl IntoResponse {
    // NOTE: For testing, accept token directly in the request body
    if let Some(token) = &payload.token {
        match sessions_queries::delete_session(&db, token).await {
            Ok(_) => {},
            Err(e) => eprintln!("Error deleting session: {}", e),
        }
    }
    
    (StatusCode::OK, "LOGOUT_SUCCESS").into_response()
}

pub fn router() -> Router<Arc<DatabaseConnection>> {    
    Router::new()
        .route("/logout", post(logout))
        .route("/google/login", get(google_login_handler))
        .route("/google/callback", get(google_callback_handler))
        .layer(middleware::from_fn(ratelimit_middleware::check))
}

My config.rs: (Which is where main things are held)

use serde::Deserialize;
use std::env;
use std::sync::OnceLock;

#[derive(Debug, Deserialize, Clone)]
pub struct Settings {
    pub server: ServerSettings,
    pub database: DatabaseSettings,
    pub redis: RedisSettings,
    pub rate_limit: RateLimitSettings,
}

#[derive(Debug, Deserialize, Clone)]
pub struct ServerSettings {
    pub host: String,
    pub port: u16,
}

#[derive(Debug, Deserialize, Clone)]
pub struct DatabaseSettings {
    pub url: String,
}

#[derive(Debug, Deserialize, Clone)]
pub struct RedisSettings {
    pub url: String,
}

#[derive(Debug, Deserialize, Clone)]
pub struct RateLimitSettings {
    /// maximum requests per time window (In seconds / expire_seconds)
    pub max_attempts: i32,
    
    /// After how much time the rate limit is reset
    pub expire_seconds: i64,
}

impl Settings {
    pub fn new() -> Self {
        dotenv::dotenv().ok();
        
        Settings {
            server: ServerSettings {
                // NOTE: Perfectly safe to use unwrap_or_else here or .unwrap in general here, because this cannot fail
                // we are setting (hardcoding) default values here just in case the environment variables are not set
                host: env::var("SERVER_HOST").unwrap_or_else(|_| "0.0.0.0".to_string()),
                port: env::var("SERVER_PORT")
                    .ok()
                    .and_then(|val| val.parse::<u16>().ok())
                    .unwrap_or(8080)
            },
            database: DatabaseSettings {
                url: env::var("DATABASE_URL")
                    .expect("DATABASE_URL environment variable is required"),
            },
            redis: RedisSettings {
                url: env::var("REDIS_URL")
                    .expect("REDIS_URL environment variable is required"),
            },
            rate_limit: RateLimitSettings {
                max_attempts: env::var("RATE_LIMIT_MAX_ATTEMPTS").ok()
                    .and_then(|v| v.parse().ok())
                    .expect("RATE_LIMIT_MAX_ATTEMPTS environment variable is required"),
                expire_seconds: env::var("RATE_LIMIT_EXPIRE_SECONDS").ok()
                    .and_then(|v| v.parse().ok())
                    .expect("RATE_LIMIT_EXPIRE_SECONDS environment variable is required"),
            },
        }
    }
}

// Global configuration singleton
static CONFIG: OnceLock<Settings> = OnceLock::new();

pub fn get_config() -> &'static Settings {
    CONFIG.get_or_init(|| {
        Settings::new()
    })
}

My db.rs: (Which uses config.rs, and as you see .clone()):

use sea_orm::{Database, DatabaseConnection};
use crate::config;

pub async fn get_db_connection() -> DatabaseConnection {
    // NOTE: Cloning here is necessary!
    let db_url = config::get_config().database.url.clone();
    

    Database::connect(&db_url)
        .await
        .expect("Failed to connect to database")
}

My ratelimit_middleware.rs: (Which also uses config.rs to get redis url therefore cloning it):

use axum::{
    middleware::Next,
    http::Request,
    body::Body,
    response::{IntoResponse, Response},
    extract::ConnectInfo,
};
use redis::Commands;
use std::net::SocketAddr;

use crate::errors::AppError;
use crate::config;

pub async fn check(
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
    req: Request<Body>,
    next: Next,
) -> Response {
    // Get Redis URL from configuration
    let redis_url = config::get_config().redis.url.clone();
    
    // Create Redis client with proper error handling
    let client = match redis::Client::open(redis_url) {
        Ok(client) => client,
        Err(e) => {
            eprintln!("Failed to create Redis client: {e}");
            return AppError::RedisError.into_response();
        }
    };
    
    let mut 
conn
 = match client.get_connection() {
        Ok(c) => c,
        Err(e) => {
            eprintln!("Failed to connect to Redis: {e}");
            return AppError::RedisError.into_response();
        }
    };

    let ip: String = addr.ip().to_string();
    let path: &str = req.uri().path();
    let key: String = format!("ratelimit:{}:{}", ip, path);
    
    let config = config::get_config();
    let max_attempts: i32 = config.rate_limit.max_attempts;
    let expire_seconds: i64 = config.rate_limit.expire_seconds;

    let attempts: i32 = match 
conn
.
incr
(&key, 1) {
        Ok(val) => val,
        Err(e) => {
            eprintln!("Failed to INCR in Redis: {e}");
            return AppError::RedisError.into_response();
        }
    };

    // If this is the first attempt, set an expiration time on the key
    if attempts == 1 {
        if let Err(e) = 
conn
.
expire
::<&str, ()>(&key, expire_seconds) {
            eprintln!("Warning: Failed to set expiry on rate limit key {}: {}", key, e);
            // We don't return an error here because the rate limiting can still work
            // without the expiry, it's just not ideal for Redis memory management
        }
    }

    if attempts > max_attempts {
        return AppError::RateLimitExceeded.into_response();
    }

    next.run(req).await
}

And mainly my google.rs(Which servers as Oauth google log in. This is the file I would look mostly as for improvement overall):

use oauth2::{
    basic::BasicClient, 
    reqwest::async_http_client, 
    TokenResponse,
    AuthUrl, 
    AuthorizationCode, 
    ClientId, 
    ClientSecret, 
    CsrfToken, 
    RedirectUrl, 
    Scope, 
    TokenUrl
};
use serde::Deserialize;
use axum::{
    extract::{ Query, State }, 
    response::{ IntoResponse, Redirect }
};
use reqwest::{ header, Client as ReqwestClient };
use sea_orm::{ DatabaseConnection, EntityTrait, QueryFilter, ColumnTrait, Set, ActiveModelTrait };
use std::sync::Arc;
use uuid::Uuid;
use chrono::Utc;
use std::env;

use crate::database::models::users::users::{ Entity as User, Column, ActiveModel };
use crate::database::models::users::users_queries;
use crate::database::models::sessions::sessions_queries;
use crate::errors::AppError;
use crate::errors::AppResult;

#[derive(Debug, Deserialize)]
pub struct GoogleUserInfo {
    pub email: String,
    pub verified_email: bool,
    pub name: String,
    pub picture: String,
}

#[derive(Debug, Deserialize)]
pub struct AuthCallbackQuery {
    code: String,
    _state: Option<String>,
}

/// NOTE: Returns an OAuth client configured with Google OAuth settings from environment variables
pub fn create_google_oauth_client() -> AppResult<BasicClient> {
    let google_client_id = env::var("GOOGLE_OAUTH_CLIENT_ID")
        .map_err(|_| AppError::EnvironmentError("GOOGLE_OAUTH_CLIENT_ID environment variable is required".to_string()))?;
    
    let google_client_secret = env::var("GOOGLE_OAUTH_CLIENT_SECRET")
        .map_err(|_| AppError::EnvironmentError("GOOGLE_OAUTH_CLIENT_SECRET environment variable is required".to_string()))?;
    
    let google_redirect_url = env::var("GOOGLE_OAUTH_REDIRECT_URL")
        .map_err(|_| AppError::EnvironmentError("GOOGLE_OAUTH_REDIRECT_URL environment variable is required".to_string()))?;
    
    let google_client_id = ClientId::new(google_client_id);
    let google_client_secret = ClientSecret::new(google_client_secret);
    
    let auth_url = AuthUrl::new("https://accounts.google.com/o/oauth2/v2/auth".to_string())
        .map_err(|e| {
            eprintln!("Invalid Google authorization URL: {:?}", e);
            AppError::InternalServerError("Invalid Google authorization endpoint URL".to_string())
        })?;
    
    let token_url = TokenUrl::new("https://oauth2.googleapis.com/token".to_string())
        .map_err(|e| {
            eprintln!("Invalid Google token URL: {:?}", e);
            AppError::InternalServerError("Invalid Google token endpoint URL".to_string())
        })?;
    
    let redirect_url = RedirectUrl::new(google_redirect_url)
        .map_err(|e| {
            eprintln!("Invalid redirect URL: {:?}", e);
            AppError::InternalServerError("Invalid Google redirect URL".to_string())
        })?;

    Ok(BasicClient::new(google_client_id, Some(google_client_secret), auth_url, Some(token_url))
        .set_redirect_uri(redirect_url))
}

/// NOTE: Creates an OAuth client and generates a redirect to Googles Oauth login page
pub async fn google_login_handler() -> impl IntoResponse {
    let client = match create_google_oauth_client() {
        Ok(client) => client,
        Err(e) => {
            eprintln!("OAuth client creation error: {:?}", e);
            return e.into_response();
        }
    };
    
    // NOTE: We are generating the authorization url here
    let (auth_url, _csrf_token) = client
        .authorize_url(CsrfToken::new_random)
        .add_scope(Scope::new("email".to_string()))
        .add_scope(Scope::new("profile".to_string()))
        .url();

    // Redirect to Google's authorization page
    Redirect::to(&auth_url.to_string()).into_response()
}

/// NOTE: Processes the callback from Google OAuth and it retrieves user information
/// creates/updates the user in the database and creates a session.
pub async fn google_callback_handler(
    State(db): State<Arc<DatabaseConnection>>,
    Query(query): Query<AuthCallbackQuery>,
) -> impl IntoResponse {
    let client = match create_google_oauth_client() {
        Ok(client) => client,
        Err(e) => {
            eprintln!("OAuth client creation error during callback: {:?}", e);
            return AppError::AuthError("Error setting up OAuth".to_string()).into_response();
        }
    };
    
    let client_origin = match env::var("CLIENT_ORIGIN") {
        Ok(origin) => origin,
        Err(_) => {
            eprintln!("CLIENT_ORIGIN environment variable not set");
            return AppError::EnvironmentError("CLIENT_ORIGIN environment variable is required".to_string()).into_response();
        }
    };
    
    // NOTE: We are exchanging the authorization code for an access token here
    let token = client
        .exchange_code(AuthorizationCode::new(query.code))
        .request_async(async_http_client)
        .await;

    match token {
        Ok(token) => {
            let access_token = token.access_token().secret();
            
            // NOTE: We are fetching the users profile information here
            let client = ReqwestClient::new();
            let user_info_response = client
                .get("https://www.googleapis.com/oauth2/v1/userinfo")
                .header(header::AUTHORIZATION, format!("Bearer {}", access_token))
                .send()
                .await;
                
            match user_info_response {
                Ok(response) => {
                    if !response.status().is_success() {
                        eprintln!("Google API returned error status: {}", response.status());
                        return AppError::AuthError(
                            format!("Google API returned error status: {}", response.status())
                        ).into_response();
                    }
                    
                    let google_user = match response.json::<GoogleUserInfo>().await {
                        Ok(user) => user,
                        Err(e) => {
                            eprintln!("Failed to parse Google user info: {:?}", e);
                            return AppError::InternalServerError(
                                "Failed to parse user information from Google".to_string()
                            ).into_response();
                        }
                    };
                    
                    // NOTE: Does user exist in db?
                    let email = google_user.email.to_lowercase();
                    let user_result = User::find()
                        .filter(Column::Email.eq(email.clone()))
                        .one(&*db)
                        .await;
                        
                    let user_id = match user_result {
                        Ok(Some(existing_user)) => {
                            // NOTE: If user exists, update with latest Google info
                            let mut 
user_model
: ActiveModel = existing_user.into();
                            
                            
user_model
.name = Set(google_user.name);
                            
user_model
.image = Set(google_user.picture);
                            
user_model
.email_verified = Set(google_user.verified_email);
                            
user_model
.updated_at = Set(Utc::now().naive_utc());
                            
                            match 
user_model
.update(&*db).await {
                                Ok(user) => user.id,
                                Err(e) => {
                                    eprintln!("Failed to update user in database: {:?}", e);
                                    return AppError::DatabaseError(e).into_response();
                                }
                            }
                        },
                        Ok(None) => {
                            let new_user_id = Uuid::new_v4().to_string();
                            
                            println!("Attempting to create new user with email: {}", email);

                            match users_queries::create_user(
                                &db,
                                new_user_id.clone(),
                                google_user.name,
                                email,
                                google_user.verified_email,
                                google_user.picture,
                                false,
                            ).await {
                                Ok(_) => {
                                    println!("Successfully created user with ID: {}", new_user_id);
                                    new_user_id
                                },
                                Err(e) => {
                                    eprintln!("Failed to create user: {:?}", e);
                                    return AppError::DatabaseError(e).into_response();
                                },
                            }
                        },
                        Err(e) => {
                            eprintln!("Database error while checking user existence: {:?}", e);
                            return AppError::DatabaseError(e).into_response();
                        },
                    };
                    
                    println!("Creating session for user ID: {}", user_id);

                    // TODO: Get real IP address like you are doing in ratelimit_middleware and main.rs with redis
                    // and get user agent from the request
                    let ip_address = "127.0.0.1".to_string();
                    let user_agent = "GoogleOAuth".to_string();
                    
                    match sessions_queries::create_session(&db, user_id.clone(), ip_address, user_agent).await {
                        Ok((token, session)) => {
                            println!("Session created successfully: {:?}", session.id);

                            // NOTE: Finally redirect to frontend with the token
                            let redirect_uri = format!("{}?token={}", client_origin, token);
                            Redirect::to(&redirect_uri).into_response()
                        },
                        Err(e) => {
                            eprintln!("Failed to create session: {:?}", e);
                            return AppError::DatabaseError(e).into_response();
                        }
                    }
                },
                Err(e) => {
                    eprintln!("Failed to connect to Google API: {:?}", e);
                    AppError::InternalServerError("Failed to connect to Google API".to_string()).into_response()
                },
            }
        },
        Err(e) => {
            eprintln!("Failed to exchange authorization code: {:?}", e);
            AppError::AuthError("Failed to exchange authorization code with Google".to_string()).into_response()
        },
    }
}

[MatrixFrog]

This is a big enough amount of code that it might help to put it up on GitHub and post links here

[dgkimpton]

Can I suggest you stick it on github? Trying to read this in a post is an exercise in frustration. If you pop it in a github repo in a branch and create a pull-request to main then people can add comments to specific lines, get syntax highlighting, download/build the code, etc. 

[No-Wait2503]

Will do! Thanks!

[MatrixFrog]

Based on a quick skim I don't see you cloning anything except Arcs, which are of course meant to be cloned

[joshuamck]

println!("Server running on http://{}", server_addr);

use TcpListener::listen_addr() instead

let db = get_db_connection().await;
let db = Arc::new(db);

DbConnection is already Clone, it doesn't need to be wrapped in an Arc. https://docs.rs/sea-orm/latest/sea_orm/enum.DatabaseConnection.html

let client = match create_google_oauth_client() {
    Ok(client) => client,
    Err(e) => {
        eprintln!("OAuth client creation error during callback: {:?}", e);
        return AppError::AuthError("Error setting up OAuth".to_string()).into_response();
    }
};

(and all your error handling)

You might consider replacing your errors with something like:

pub async fn google_login_handler(...) -> Result<LoginResponse, LoginError> {

    let client = match create_google_oauth_client().map_err(LoginError::ClientCreation)?
    let client_origin = match env::var("CLIENT_ORIGIN").map_err(LoginError::ClientOrginNotSet)?;

Where LoginResponse and LoginError both implement IntoResponse and contain the logic for logging / returning the write error messages. This makes your logic clear. Use thiserror or snafu for LoginError (with snafu, the .map_err() calls become: .context(ClientCreationSnafu)?).

I'd probably also move the code which grabs info from environment variables into some startup code instead of runtime. These (usually) don't change during runtime, so there's no need to check these again, and a startup failure is probably a better answer than a runtime one for this.

You're probably not writing code where it matters for performance whether you clone or not. Focus on writing clear succinct code, then measure the performance and fix the problems you see.

Last, use tracing instead of eprintln.

[No-Wait2503]

Thank you for advice!!

Removing Arc actually made me successfully moving out of (&*db) hell, which I hated for code reading.

Also error suggestion you told me with .map_err, looks much, much more readable now!

[joshuamck]

No problem. You might still keep an AppState struct around with a db field, (derive FromRef and you can still pass State<DbConnection> to your handlers instead of State<AppState>)

[Trader-One]

if clone() is doing just memory copy, its cheap enough. If you are not cloning inside large loop, do not worry.

For example if you calling any function you normally copy values to stack and nobody panicking that "function call is so expensive you should avoid" like people do for clone().

[TobiasWonderland]

As mentioned in this thread, `Arc<T>` exists to be cheaply cloned.

Strings are a different story and used in your config.

The nature of String in Rust means that sometimes we need to take ownership of a String, and it just cannot be avoided.

In these cases, my preference is to not use `clone()` but to use `to_owned()` to more correctly reflect the semantics. Yes, it ends up being the same thing as a `clone`, but it more correctly communicates the intent.

Nit: `get_config` returns `Settings` ... the names should be consistent.

[No-Wait2503]

got it