Denial Of Service in Axum with Cockroachdb

When building a web service with the Axum web framework in Rust and using CockroachDB as the backend, certain request patterns and database interactions can create denial-of-service (DoS) conditions. DoS in this context refers to scenarios where legitimate requests cannot be served in a timely manner due to resource exhaustion or unbounded contention, not necessarily a network-layer flood.

One common pattern is long-running or unbounded SQL queries executed on a single database connection or without proper concurrency controls. CockroachDB is a distributed SQL database that provides strong consistency and horizontal scalability, but it still requires careful query design. In Axum, if a handler executes a query that performs a full table scan or joins large datasets without indexes, the query may consume significant database-side resources (CPU, I/O, and transaction contention). This can block other transactions and increase latency for all requests, effectively creating a DoS condition for the service.

A second dimension is connection management. Axum applications often use a connection pool (for example via bb8 or sqlx) to talk to CockroachDB. If the pool size is too small relative to the incoming request rate, requests will wait for an available connection, causing thread starvation and increased response times. Because CockroachDB may take longer to execute heavy queries, connections are held for extended periods, exacerbating pool exhaustion. This is especially impactful when combined with synchronous blocking calls in async handlers, which can degrade the runtime performance of an Axum service.

A third dimension involves retries and idempotency. CockroachDB recommends idempotent retries for transient errors, but if Axum handlers are not designed to be idempotent—especially for mutating operations—automatic retries from the client or the database driver can amplify load. Combined with a lack of request deduplication or rate limiting at the application layer, this can lead to repeated execution of costly operations, increasing load on the database and contributing to service unavailability for other users.

These DoS risks are detectable by middleBrick’s unauthenticated scans, which include checks for Rate Limiting, Input Validation, and Unsafe Consumption among twelve parallel security checks. By submitting your API endpoint to middleBrick, you can identify whether long-running queries, missing indexes, or weak concurrency configurations are present without requiring credentials or agent installation.

To reduce DoS risk in an Axum service using CockroachDB, focus on query efficiency, connection pool sizing, and idempotent handler design. Below are concrete, syntactically correct examples that demonstrate these principles.

1. Use indexed queries and limit result sets

Ensure your SQL includes appropriate WHERE clauses and that the relevant columns are indexed in CockroachDB. This avoids full table scans and reduces query latency.

-- CockroachDB SQL: create an index to support efficient filtering
CREATE INDEX IF NOT EXISTS idx_users_email ON users(email);

-- Example SQL query (used from Axum): fetch a single user by indexed column
SELECT id, email, name FROM users WHERE email = $1 LIMIT 1;

2. Configure a sensible connection pool in Axum

Size the pool based on expected concurrency and keep queries short. Here is an example using sqlx with a PostgreSQL-compatible CockroachDB connection string.

use sqlx::postgres::PgPoolOptions;
use sqlx::PgPool;

#[tokio::main]
async fn main() -> Result<(), sqlx::Error> {
    let database_url = std::env::var("DATABASE_URL")
        .expect("DATABASE_URL must be set");
    // Configure pool size and timeouts to avoid resource exhaustion
    let pool = PgPool::connect_lazy_with(
        sqlx::postgres::PgConnectOptions::from_str(&database_url)?
            .max_connections(20)
            .connect_timeout(std::time::Duration::from_secs(2)),
    );
    // Use pool in Axum routes...
    Ok(())
}

3. Make handlers idempotent and avoid long transactions

Keep transactions short and design mutating endpoints to be safely retried. Use optimistic concurrency control where appropriate.

use axum::{
    routing::post,
    Router,
};
use sqlx::PgPool;
use std::net::SocketAddr;

async fn create_order_handler(
    pool: PgPool,
    axum::Json(payload): axum::Json<CreateOrderPayload>
) -> Result<axum::Json<OrderResponse>, (axum::http::StatusCode, String)> {
    let mut tx = pool.begin().await.map_err(|e| {
        (axum::http::StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
    })?;

    // Idempotency key check to avoid duplicate processing
    let exists: (bool,) = sqlx::query_as(
        "SELECT idempotency_key FROM idempotency WHERE idempotency_key = $1"
    )
    .bind(&payload.idempotency_key)
    .fetch_one(&mut *tx)
    .await
    .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    if exists {
        tx.rollback().await.unwrap_or(());
        return Err((axum::http::StatusCode::CONFLICT, "Duplicate request".into()));
    }

    // Short, indexed mutation
    sqlx::query(
        "INSERT INTO orders (user_id, amount, idempotency_key) VALUES ($1, $2, $3)"
    )
    .bind(payload.user_id)
    .bind(payload.amount)
    .bind(&payload.idempotency_key)
    .execute(&mut *tx)
    .await
    .map_err(|e| (axum::http::StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    tx.commit().await.map_err(|e| {
        (axum::http::StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
    })?;

    Ok(axum::Json(OrderResponse { order_id: 42 }))
}

4. Apply rate limiting and timeouts at the application layer

Use middleware to enforce request-rate ceilings and set query timeouts to prevent unbounded execution.

use axum::middleware::from_fn_with_state;

async fn rate_limit_middleware(
    // simplified example; use a crate like `governor` or `axum-extra` in production
    State(_state): State<AppState>,
    request: Request,
    next: Next,
) -> Result<Response, (StatusCode, String)> {
    // Implement token-bucket or fixed-window logic here
    Ok(next.run(request).await)
}

let app = Router::new()
    .route("/orders", post(create_order_handler))
    .layer(from_fn_with_state((), rate_limit_middleware));