feat(enterprise): add auth TLS and recovery hardening

2026-04-29 22:01:43 +00:00
parent 2f3031cfc7
commit ed2c02bcf9
27 changed files with 2369 additions and 55 deletions
@@ -0,0 +1,87 @@
+use bson::{doc, Binary, Bson, Document};
+
+use crate::context::{CommandContext, ConnectionState};
+use crate::error::{CommandError, CommandResult};
+
+pub async fn handle_sasl_start(
+    cmd: &Document,
+    db: &str,
+    ctx: &CommandContext,
+    connection: &mut ConnectionState,
+) -> CommandResult<Document> {
+    let mechanism = cmd
+        .get_str("mechanism")
+        .map_err(|_| CommandError::InvalidArgument("missing SASL mechanism".into()))?;
+    if mechanism != "SCRAM-SHA-256" {
+        return Err(CommandError::InvalidArgument(format!(
+            "unsupported SASL mechanism: {mechanism}"
+        )));
+    }
+
+    let payload = payload_bytes(cmd)?;
+    let result = ctx
+        .auth
+        .start_scram_sha256(db, &payload)
+        .map_err(map_auth_error)?;
+    let conversation_id = connection.next_conversation_id();
+    connection
+        .sasl_conversations
+        .insert(conversation_id, result.conversation);
+
+    Ok(doc! {
+        "conversationId": conversation_id,
+        "done": false,
+        "payload": Binary { subtype: bson::spec::BinarySubtype::Generic, bytes: result.payload },
+        "ok": 1.0,
+    })
+}
+
+pub async fn handle_sasl_continue(
+    cmd: &Document,
+    ctx: &CommandContext,
+    connection: &mut ConnectionState,
+) -> CommandResult<Document> {
+    let conversation_id = cmd
+        .get_i32("conversationId")
+        .map_err(|_| CommandError::InvalidArgument("missing SASL conversationId".into()))?;
+    let payload = payload_bytes(cmd)?;
+    let conversation = connection
+        .sasl_conversations
+        .remove(&conversation_id)
+        .ok_or_else(|| CommandError::InvalidArgument("unknown SASL conversation".into()))?;
+    let result = ctx
+        .auth
+        .continue_scram_sha256(conversation, &payload)
+        .map_err(map_auth_error)?;
+    connection.authenticate(result.user);
+
+    Ok(doc! {
+        "conversationId": conversation_id,
+        "done": true,
+        "payload": Binary { subtype: bson::spec::BinarySubtype::Generic, bytes: result.payload },
+        "ok": 1.0,
+    })
+}
+
+fn payload_bytes(cmd: &Document) -> CommandResult<Vec<u8>> {
+    match cmd.get("payload") {
+        Some(Bson::Binary(binary)) => Ok(binary.bytes.clone()),
+        Some(Bson::String(value)) => Ok(value.as_bytes().to_vec()),
+        _ => Err(CommandError::InvalidArgument("missing SASL payload".into())),
+    }
+}
+
+fn map_auth_error(error: rustdb_auth::AuthError) -> CommandError {
+    match error {
+        rustdb_auth::AuthError::InvalidPayload(message) => CommandError::InvalidArgument(message),
+        rustdb_auth::AuthError::UnsupportedMechanism(message) => CommandError::InvalidArgument(message),
+        rustdb_auth::AuthError::Disabled => CommandError::Unauthorized("authentication is disabled".into()),
+        rustdb_auth::AuthError::UnknownConversation => {
+            CommandError::InvalidArgument("unknown SASL conversation".into())
+        }
+        rustdb_auth::AuthError::AuthenticationFailed => CommandError::AuthenticationFailed,
+        rustdb_auth::AuthError::UserAlreadyExists(message) => CommandError::DuplicateKey(message),
+        rustdb_auth::AuthError::UserNotFound(message) => CommandError::NamespaceNotFound(message),
+        rustdb_auth::AuthError::Persistence(message) => CommandError::InternalError(message),
+    }
+}