feat: enforce API key authentication on all data routes (#45)

Author: ajainvivek

Cargo.lock

@@ -28,6 +28,7 @@ tower = "0.4"
 tower-http = { version = "0.5", features = ["cors", "trace", "limit"] }
 
 # Async
+async-trait = "0.1"
 tokio.workspace = true
 tokio-stream = "0.1"
 futures.workspace = true

crates/reasondb-server/Cargo.toml

@@ -5,13 +5,19 @@
 //! - API key authentication via `X-API-Key: <key>` header
 //! - Optional authentication (for public endpoints)
 
+use async_trait::async_trait;
 use axum::{
+    extract::{FromRequestParts, Request, State},
     http::{header, request::Parts, StatusCode},
+    middleware::Next,
     response::{IntoResponse, Response},
     Json,
 };
 use reasondb_core::{ApiKey, KeyPrefix, Permission, Permissions};
 use serde::Serialize;
+use std::sync::Arc;
+
+use crate::state::AppState;
 
 /// Authenticated API key (extracted from request)
 #[derive(Debug, Clone)]
@@ -174,3 +180,106 @@ pub fn extract_api_key(parts: &Parts) -> Option<String> {
 
     None
 }
+
+/// Axum middleware that enforces API key authentication on all routes when
+/// `REASONDB_AUTH_ENABLED=true`.
+///
+/// Public routes (`/health`, `/metrics`, `/swagger-ui`, `/api-docs`) bypass
+/// auth so monitoring and documentation remain accessible without a key.
+pub async fn auth_middleware<
+    R: reasondb_core::llm::ReasoningEngine + Clone + Send + Sync + 'static,
+>(
+    State(state): State<Arc<AppState<R>>>,
+    request: Request,
+    next: Next,
+) -> Response {
+    // Auth disabled — pass through
+    if !state.config.auth.enabled {
+        return next.run(request).await;
+    }
+
+    // Public paths that never require auth
+    let path = request.uri().path().to_owned();
+    if path == "/health"
+        || path == "/metrics"
+        || path.starts_with("/swagger-ui")
+        || path.starts_with("/api-docs")
+    {
+        return next.run(request).await;
+    }
+
+    // Extract API key from headers before consuming the request
+    let raw_key = {
+        let headers = request.headers();
+        // Try Authorization: Bearer <key>
+        let from_auth = headers
+            .get(header::AUTHORIZATION)
+            .and_then(|v| v.to_str().ok())
+            .and_then(|v| v.strip_prefix("Bearer ").map(|s| s.trim().to_string()));
+        // Fall back to X-API-Key
+        let from_x_key = headers
+            .get("X-API-Key")
+            .and_then(|v| v.to_str().ok())
+            .map(|s| s.trim().to_string());
+        from_auth.or(from_x_key)
+    };
+
+    let raw_key = match raw_key {
+        Some(k) => k,
+        None => return AuthError::MissingKey.into_response(),
+    };
+
+    // Validate against master key
+    if let Some(ref master_key) = state.config.auth.master_key {
+        if raw_key == *master_key {
+            return next.run(request).await;
+        }
+    }
+
+    // Validate against stored API keys
+    match state.api_key_store.authenticate(&raw_key) {
+        Ok(Some(key)) if key.is_active => next.run(request).await,
+        Ok(Some(_)) => AuthError::RevokedKey.into_response(),
+        Ok(None) => AuthError::InvalidKey.into_response(),
+        Err(e) => AuthError::Internal(e.to_string()).into_response(),
+    }
+}
+
+/// `FromRequestParts` impl so handlers can optionally extract an
+/// `AuthenticatedKey` without the middleware (used by auth management routes).
+#[async_trait]
+impl<R> FromRequestParts<Arc<AppState<R>>> for AuthenticatedKey
+where
+    R: reasondb_core::llm::ReasoningEngine + Clone + Send + Sync + 'static,
+{
+    type Rejection = AuthError;
+
+    async fn from_request_parts(
+        parts: &mut Parts,
+        state: &Arc<AppState<R>>,
+    ) -> Result<Self, Self::Rejection> {
+        if !state.config.auth.enabled {
+            return Ok(AuthenticatedKey::anonymous());
+        }
+
+        let raw_key = extract_api_key(parts).ok_or(AuthError::MissingKey)?;
+
+        if let Some(ref master_key) = state.config.auth.master_key {
+            if raw_key == *master_key {
+                return Ok(AuthenticatedKey::master());
+            }
+        }
+
+        let key = state
+            .api_key_store
+            .authenticate(&raw_key)
+            .map_err(|e| AuthError::Internal(e.to_string()))?
+            .ok_or(AuthError::InvalidKey)?;
+
+        if !key.is_active {
+            return Err(AuthError::RevokedKey);
+        }
+
+        Ok(AuthenticatedKey { key })
+    }
+}

crates/reasondb-server/src/auth.rs

@@ -61,6 +61,7 @@ pub mod replication;
 pub mod routes;
 pub mod state;
 
+pub use auth::auth_middleware;
 pub use error::{ApiError, ApiResult, ErrorResponse};
 pub use metrics::{init_metrics, metrics_handler, metrics_middleware};
 #[cfg(feature = "telemetry")]
@@ -108,6 +109,12 @@ pub fn create_server<R: ReasoningEngine + Clone + Send + Sync + 'static>(
     // Add Prometheus metrics endpoint
     app = app.route("/metrics", get(metrics::metrics_handler));
 
+    // Add API key authentication middleware (enforces auth when REASONDB_AUTH_ENABLED=true)
+    app = app.layer(axum::middleware::from_fn_with_state(
+        state.clone(),
+        auth_middleware,
+    ));
+
     // Add rate limiting middleware
     if state.config.rate_limit.enabled {
         info!(

crates/reasondb-server/src/lib.rs

@@ -5,6 +5,10 @@ terraform {
       source  = "hashicorp/aws"
       version = "~> 5.0"
     }
+    null = {
+      source  = "hashicorp/null"
+      version = "~> 3.0"
+    }
   }
 }
 
@@ -103,6 +107,8 @@ resource "aws_instance" "reasondb" {
     llm_model      = var.llm_model
     llm_base_url   = var.llm_base_url
     reasondb_image = var.reasondb_image
+    auth_enabled   = var.auth_enabled
+    master_key     = var.master_key
   })
 
   tags = { Name = "${var.name_prefix}-instance" }
@@ -144,3 +150,68 @@ resource "aws_eip" "reasondb" {
 
   depends_on = [aws_instance.reasondb]
 }
+
+# ---------------------------------------------------------------------------
+# Container updater — re-runs whenever image or config variables change.
+#
+# user_data only executes once at first boot, so this null_resource handles
+# zero-downtime container updates on subsequent `terraform apply` runs by
+# SSH-ing into the instance and restarting the Docker container in-place.
+# ---------------------------------------------------------------------------
+
+resource "null_resource" "container_update" {
+  # Re-trigger whenever any of these values change
+  triggers = {
+    image        = var.reasondb_image
+    llm_provider = var.llm_provider
+    llm_model    = var.llm_model
+    llm_base_url = var.llm_base_url
+    auth_enabled = tostring(var.auth_enabled)
+    # Use a hash of the key so the value doesn't appear in state
+    master_key_hash = sha256(var.master_key)
+    llm_key_hash    = sha256(var.llm_api_key)
+    # Re-run if the instance is replaced
+    instance_id = aws_instance.reasondb.id
+  }
+
+  connection {
+    type        = "ssh"
+    user        = "ubuntu"
+    private_key = file(var.ssh_private_key_path)
+    host        = aws_eip.reasondb.public_ip
+  }
+
+  provisioner "remote-exec" {
+    inline = [
+      "echo '=== Pulling ${var.reasondb_image} ==='",
+      "sudo docker pull ${var.reasondb_image}",
+      "echo '=== Stopping existing container ==='",
+      "sudo docker stop reasondb 2>/dev/null || true",
+      "sudo docker rm   reasondb 2>/dev/null || true",
+      "echo '=== Starting updated container ==='",
+      "sudo docker run -d \\",
+      "  --name reasondb \\",
+      "  --restart unless-stopped \\",
+      "  -p 4444:4444 \\",
+      "  -v /data:/data \\",
+      "  -e REASONDB_HOST=0.0.0.0 \\",
+      "  -e REASONDB_PORT=4444 \\",
+      "  -e REASONDB_PATH=/data/reasondb.redb \\",
+      "  -e REASONDB_LLM_PROVIDER=${var.llm_provider} \\",
+      "  -e REASONDB_LLM_API_KEY=${var.llm_api_key} \\",
+      "  -e REASONDB_MODEL=${var.llm_model} \\",
+      "  -e REASONDB_LLM_BASE_URL=${var.llm_base_url} \\",
+      "  -e REASONDB_RATE_LIMIT_RPM=300 \\",
+      "  -e REASONDB_RATE_LIMIT_RPH=5000 \\",
+      "  -e REASONDB_RATE_LIMIT_BURST=30 \\",
+      "  -e REASONDB_WORKER_COUNT=4 \\",
+      "  -e REASONDB_AUTH_ENABLED=${var.auth_enabled} \\",
+      "  -e REASONDB_MASTER_KEY=${var.master_key} \\",
+      "  ${var.reasondb_image}",
+      "echo '=== Waiting for health check ==='",
+      "for i in $(seq 1 24); do curl -sf http://localhost:4444/health && echo '' && break || sleep 5; done",
+    ]
+  }
+
+  depends_on = [aws_eip.reasondb, aws_volume_attachment.data]
+}

deploy/aws/terraform/main.tf

@@ -11,8 +11,9 @@ instance_type = "t3.medium"
 # EBS data volume size in GB
 volume_size_gb = 20
 
-# Paste the contents of your ~/.ssh/id_rsa.pub (or equivalent)
-ssh_public_key = "ssh-rsa AAAA... you@example.com"
+# Paste the contents of your ~/.ssh/id_ed25519.pub (or equivalent)
+ssh_public_key       = "ssh-ed25519 AAAA... you@example.com"
+ssh_private_key_path = "~/.ssh/id_ed25519"
 
 # Restrict access to your IP for security, or use 0.0.0.0/0 for open access
 allowed_cidr = "0.0.0.0/0"
@@ -28,5 +29,11 @@ llm_api_key  = "sk-..."
 # Optional: custom base URL (e.g. for Ollama or a proxy)
 # llm_base_url = "http://localhost:11434/v1"
 
-# Docker image (use a specific version tag in production, e.g. "ajainvivek/reasondb:0.1.0")
-reasondb_image = "ajainvivek/reasondb:latest"
+# Docker image (use a specific version tag in production, e.g. "brainfishai/reasondb:0.5.7")
+reasondb_image = "brainfishai/reasondb:latest"
+
+# ── Authentication ─────────────────────────────────────────────────────────────
+# Enable API key auth (recommended for any non-local deployment)
+auth_enabled = true
+# Master admin key — generate with: python3 -c "import secrets; print('rdb_master_' + secrets.token_hex(24))"
+master_key = "rdb_master_..."

deploy/aws/terraform/terraform.tfvars.example

@@ -68,4 +68,6 @@ docker run -d \
   -e REASONDB_RATE_LIMIT_RPH="5000" \
   -e REASONDB_RATE_LIMIT_BURST="30" \
   -e REASONDB_WORKER_COUNT="4" \
+  -e REASONDB_AUTH_ENABLED="${auth_enabled}" \
+  -e REASONDB_MASTER_KEY="${master_key}" \
   ${reasondb_image}

deploy/aws/terraform/user_data.sh.tpl

@@ -23,10 +23,16 @@ variable "volume_size_gb" {
 }
 
 variable "ssh_public_key" {
-  description = "SSH public key material (contents of ~/.ssh/id_rsa.pub or similar)"
+  description = "SSH public key material (contents of ~/.ssh/id_ed25519.pub or similar)"
   type        = string
 }
 
+variable "ssh_private_key_path" {
+  description = "Path to the SSH private key file used by the container updater provisioner (e.g. ~/.ssh/id_ed25519)"
+  type        = string
+  default     = "~/.ssh/id_ed25519"
+}
+
 variable "allowed_cidr" {
   description = "CIDR block allowed to reach port 4444 and 22. Use your IP (e.g. 1.2.3.4/32) or 0.0.0.0/0 for open access."
   type        = string
@@ -68,3 +74,16 @@ variable "name_prefix" {
   type        = string
   default     = "reasondb-testing"
 }
+
+variable "auth_enabled" {
+  description = "Enable API key authentication (recommended for production)"
+  type        = bool
+  default     = false
+}
+
+variable "master_key" {
+  description = "Master admin key for the ReasonDB instance (required when auth_enabled = true)"
+  type        = string
+  sensitive   = true
+  default     = ""
+}