fix(kms): handle concurrent Vault KV2 baseline races

crates/kms/src/backends/scripted_vault.rs

@@ -21,6 +21,7 @@
 //! assertions. It intentionally implements just enough HTTP/1.1 for the
 //! `vaultrs` reqwest client: no keep-alive, no chunked bodies.
 
+use std::collections::BTreeMap;
 use std::sync::{Arc, Mutex};
 
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
@@ -67,6 +68,7 @@ pub(crate) struct ScriptedVault {
     /// Base address (`http://127.0.0.1:port`) to point a Vault client at.
     pub(crate) address: String,
     requests: Arc<Mutex<Vec<(String, String)>>>,
+    kv2_state: Option<Arc<Mutex<Kv2State>>>,
 }
 
 impl ScriptedVault {
@@ -83,13 +85,16 @@ impl ScriptedVault {
         tokio::spawn(async move {
             let mut responses = responses.into_iter();
             loop {
-                let Ok((mut stream, _)) = listener.accept().await else {
+                let Ok((stream, _)) = listener.accept().await else {
                     return;
                 };
-                let Some(request) = read_request(&mut stream).await else {
+                let Some((request_line, body, mut stream)) = read_request(stream).await else {
                     continue;
                 };
-                recorded.lock().expect("scripted vault request log poisoned").push(request);
+                recorded
+                    .lock()
+                    .expect("scripted vault request log poisoned")
+                    .push((request_line, body));
                 let response = responses
                     .next()
                     .unwrap_or_else(|| ScriptedResponse::error(599, "scripted vault: script exhausted"));
@@ -104,7 +109,59 @@ impl ScriptedVault {
             }
         });
 
-        Self { address, requests }
+        Self {
+            address,
+            requests,
+            kv2_state: None,
+        }
+    }
+
+    /// Bind a small stateful KV2 responder for concurrency tests.
+    ///
+    /// Unlike Self::serve, this responder evaluates CAS writes against a
+    /// shared in-memory record and handles connections concurrently. It models
+    /// only the KV2 data and metadata paths used by the rotation protocol; an
+    /// unknown request receives a 599 response so a test cannot silently
+    /// under-specify the Vault exchange.
+    pub(crate) async fn serve_kv2(key_path: &str, key_data: serde_json::Value) -> Self {
+        let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind scripted vault listener");
+        let address = format!("http://{}", listener.local_addr().expect("scripted vault local addr"));
+        let requests = Arc::new(Mutex::new(Vec::new()));
+        let state = Arc::new(Mutex::new(Kv2State::new(key_data)));
+        let recorded = Arc::clone(&requests);
+        let state_for_server = Arc::clone(&state);
+        let key_path = key_path.to_string();
+
+        tokio::spawn(async move {
+            loop {
+                let Ok((stream, _)) = listener.accept().await else {
+                    return;
+                };
+                let recorded = Arc::clone(&recorded);
+                let state = Arc::clone(&state_for_server);
+                let key_path = key_path.clone();
+                tokio::spawn(async move {
+                    let Some((request_line, body, stream)) = read_request(stream).await else {
+                        return;
+                    };
+                    recorded
+                        .lock()
+                        .expect("scripted vault request log poisoned")
+                        .push((request_line.clone(), body.clone()));
+                    let response = state
+                        .lock()
+                        .expect("scripted KV2 state poisoned")
+                        .respond(&key_path, &request_line, &body);
+                    write_response(stream, response).await;
+                });
+            }
+        });
+
+        Self {
+            address,
+            requests,
+            kv2_state: Some(state),
+        }
     }
 
     /// The `METHOD /path` lines of every request served so far, in order.
@@ -128,13 +185,191 @@ impl ScriptedVault {
             .map(|(_, body)| body.clone())
             .collect()
     }
+
+    /// Snapshot the in-memory KV2 state used by Self::serve_kv2.
+    pub(crate) fn kv2_snapshot(&self) -> Option<Kv2Snapshot> {
+        self.kv2_state.as_ref().map(|state| {
+            let state = state.lock().expect("scripted KV2 state poisoned");
+            Kv2Snapshot {
+                current_data: state.current_data.clone(),
+                current_secret_version: state.current_secret_version,
+                version_records: state.version_records.clone(),
+            }
+        })
+    }
+}
+
+/// State captured by the stateful KV2 responder for assertions in wiring tests.
+#[derive(Debug, Clone)]
+pub(crate) struct Kv2Snapshot {
+    pub(crate) current_data: serde_json::Value,
+    pub(crate) current_secret_version: u64,
+    pub(crate) version_records: BTreeMap<u32, serde_json::Value>,
+}
+
+#[derive(Debug)]
+struct Kv2State {
+    current_data: serde_json::Value,
+    current_secret_version: u64,
+    history: BTreeMap<u64, serde_json::Value>,
+    version_records: BTreeMap<u32, serde_json::Value>,
+}
+
+impl Kv2State {
+    fn new(current_data: serde_json::Value) -> Self {
+        Self {
+            history: BTreeMap::from([(1, current_data.clone())]),
+            current_data,
+            current_secret_version: 1,
+            version_records: BTreeMap::new(),
+        }
+    }
+
+    fn respond(&mut self, key_path: &str, request_line: &str, body: &str) -> ScriptedResponse {
+        let Some((method, path)) = request_line.split_once(' ') else {
+            return ScriptedResponse::error(599, "scripted KV2: malformed request line");
+        };
+        let data_path = format!("/v1/secret/data/{key_path}");
+        let metadata_path = format!("/v1/secret/metadata/{key_path}");
+        let version_data_prefix = format!("{data_path}/versions/");
+        let version_metadata_path = format!("{metadata_path}/versions");
+
+        if method == "GET" && path == metadata_path {
+            return ScriptedResponse::ok(metadata_data(self.current_secret_version));
+        }
+
+        if method == "LIST" && path == version_metadata_path {
+            let keys = self.version_records.keys().map(u32::to_string).collect::<Vec<_>>();
+            return ScriptedResponse::ok(serde_json::json!({ "keys": keys }));
+        }
+
+        if let Some(version) = path
+            .strip_prefix(&version_data_prefix)
+            .and_then(|value| value.parse::<u32>().ok())
+        {
+            return match method {
+                "GET" => self
+                    .version_records
+                    .get(&version)
+                    .cloned()
+                    .map(read_data)
+                    .unwrap_or_else(|| ScriptedResponse::error(404, "not found")),
+                "POST" => self.create_version_record(version, body),
+                _ => ScriptedResponse::error(599, "scripted KV2: unsupported version request"),
+            };
+        }
+
+        if method == "GET" && path.strip_prefix(&data_path).is_some() {
+            let version = path
+                .split_once("?version=")
+                .and_then(|(_, value)| value.parse::<u64>().ok())
+                .unwrap_or(self.current_secret_version);
+            return self
+                .history
+                .get(&version)
+                .cloned()
+                .map(read_data)
+                .unwrap_or_else(|| ScriptedResponse::error(404, "not found"));
+        }
+
+        if method == "POST" && path == data_path {
+            return self.write_current_record(body);
+        }
+
+        ScriptedResponse::error(599, "scripted KV2: unexpected request")
+    }
+
+    fn create_version_record(&mut self, version: u32, body: &str) -> ScriptedResponse {
+        let body: serde_json::Value = match serde_json::from_str(body) {
+            Ok(body) => body,
+            Err(_) => return ScriptedResponse::error(400, "invalid JSON"),
+        };
+        if body["options"]["cas"].as_u64() != Some(0) {
+            return ScriptedResponse::error(400, "version record requires create-only CAS");
+        }
+        if self.version_records.contains_key(&version) {
+            return ScriptedResponse::error(400, CAS_CONFLICT_MESSAGE);
+        }
+        let Some(data) = body.get("data").cloned() else {
+            return ScriptedResponse::error(400, "missing data");
+        };
+        self.version_records.insert(version, data);
+        write_ack(1)
+    }
+
+    fn write_current_record(&mut self, body: &str) -> ScriptedResponse {
+        let body: serde_json::Value = match serde_json::from_str(body) {
+            Ok(body) => body,
+            Err(_) => return ScriptedResponse::error(400, "invalid JSON"),
+        };
+        if body["options"]["cas"].as_u64() != Some(self.current_secret_version) {
+            return ScriptedResponse::error(400, CAS_CONFLICT_MESSAGE);
+        }
+        let Some(data) = body.get("data").cloned() else {
+            return ScriptedResponse::error(400, "missing data");
+        };
+        self.current_secret_version += 1;
+        self.current_data = data.clone();
+        self.history.insert(self.current_secret_version, data);
+        write_ack(self.current_secret_version)
+    }
+}
+
+const CAS_CONFLICT_MESSAGE: &str = "check-and-set parameter did not match the current version";
+
+fn metadata_data(current_version: u64) -> serde_json::Value {
+    serde_json::json!({
+        "cas_required": false,
+        "created_time": "2026-01-01T00:00:00Z",
+        "current_version": current_version,
+        "delete_version_after": "0s",
+        "max_versions": 0,
+        "oldest_version": 0,
+        "updated_time": "2026-01-01T00:00:00Z",
+        "custom_metadata": null,
+        "versions": {},
+    })
+}
+
+fn read_data(data: serde_json::Value) -> ScriptedResponse {
+    ScriptedResponse::ok(serde_json::json!({
+        "data": data,
+        "metadata": {
+            "created_time": "2026-01-01T00:00:00Z",
+            "deletion_time": "",
+            "custom_metadata": null,
+            "destroyed": false,
+            "version": 1,
+        },
+    }))
+}
+
+fn write_ack(version: u64) -> ScriptedResponse {
+    ScriptedResponse::ok(serde_json::json!({
+        "created_time": "2026-01-01T00:00:00Z",
+        "custom_metadata": null,
+        "deletion_time": "",
+        "destroyed": false,
+        "version": version,
+    }))
+}
+
+async fn write_response(mut stream: TcpStream, response: ScriptedResponse) {
+    if let ScriptedResponse::Http { status, body } = response {
+        let payload = format!(
+            "HTTP/1.1 {status} Scripted\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
+            body.len(),
+        );
+        let _ = stream.write_all(payload.as_bytes()).await;
+        let _ = stream.shutdown().await;
+    }
 }
 
 /// Read one HTTP/1.1 request (head plus content-length body) and return its
 /// `METHOD /path` line together with the body. Draining the body before
 /// responding keeps the client from seeing a connection reset while it is
 /// still writing.
-async fn read_request(stream: &mut TcpStream) -> Option<(String, String)> {
+async fn read_request(mut stream: TcpStream) -> Option<(String, String, TcpStream)> {
     let mut buffer = Vec::new();
     let mut chunk = [0u8; 4096];
     let head_end = loop {
@@ -178,5 +413,5 @@ async fn read_request(stream: &mut TcpStream) -> Option<(String, String)> {
     }
     body.truncate(content_length);
 
-    Some((format!("{method} {path}"), String::from_utf8_lossy(&body).into_owned()))
+    Some((format!("{method} {path}"), String::from_utf8_lossy(&body).into_owned(), stream))
 }

crates/kms/src/backends/vault.rs

@@ -696,8 +696,15 @@ impl VaultKmsClient {
     async fn ensure_version_history_consistent(&self, key_id: &str, key_data: &VaultKeyData) -> Result<()> {
         let recorded = self.recorded_version_numbers(key_id).await?;
 
+        // A first-rotation record set (the current v1 record and, at most, the
+        // next v2 record) can be in flight while another caller is between its
+        // create-only writes and the baseline CAS. The material check below
+        // validates the v1 record before adopting it; a later key with a
+        // missing persisted baseline still indicates a lost baseline and fails
+        // closed.
         if key_data.baseline_version.is_none()
             && let Some(oldest_recorded) = recorded.iter().min().copied()
+            && !(key_data.version == 1 && recorded.iter().all(|version| *version <= 2))
         {
             warn!(
                 key_id,
@@ -1895,6 +1902,19 @@ mod tests {
         (vault, client)
     }
 
+    async fn scripted_kv2_client(key_data: &VaultKeyData) -> (ScriptedVault, VaultKmsClient) {
+        let vault = ScriptedVault::serve_kv2(
+            "rustfs/kms/keys/wired-key",
+            serde_json::to_value(key_data).expect("serialize scripted KV2 key"),
+        )
+        .await;
+        let (vault_config, kms_config) = scripted_configs(&vault.address);
+        let client = VaultKmsClient::new(vault_config, &kms_config)
+            .await
+            .expect("scripted Vault client");
+        (vault, client)
+    }
+
     fn healthy_key_data() -> VaultKeyData {
         VaultKeyData {
             algorithm: "AES_256".to_string(),
@@ -3223,6 +3243,100 @@ mod tests {
         );
     }
 
+    /// Concurrent rotations use KV2 create-only records and a CAS pointer
+    /// switch. The stateful scripted Vault applies those preconditions to real
+    /// HTTP requests, so this test proves committed versions are unique and
+    /// contiguous instead of only checking that several calls returned.
+    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+    async fn wired_concurrent_kv2_rotations_commit_unique_monotonic_versions() {
+        use std::collections::HashSet;
+        use std::sync::Arc;
+
+        const ATTEMPTS: usize = 8;
+        let (vault, client) = scripted_kv2_client(&healthy_key_data()).await;
+        let client = Arc::new(client);
+        let barrier = Arc::new(tokio::sync::Barrier::new(ATTEMPTS));
+        let tasks: Vec<_> = (0..ATTEMPTS)
+            .map(|_| {
+                let client = Arc::clone(&client);
+                let barrier = Arc::clone(&barrier);
+                tokio::spawn(async move {
+                    barrier.wait().await;
+                    client.rotate_key("wired-key", None).await
+                })
+            })
+            .collect();
+
+        let mut committed_versions = Vec::new();
+        let mut errors = Vec::new();
+        for task in tasks {
+            match task.await.expect("join concurrent rotation task") {
+                Ok(result) => committed_versions.push(result.version),
+                Err(error) => errors.push(error),
+            }
+        }
+
+        assert!(
+            !committed_versions.is_empty(),
+            "at least one concurrent rotation must commit; errors: {errors:?}"
+        );
+        assert!(
+            errors.iter().all(
+                |error| matches!(error, KmsError::InvalidOperation { message } if message.contains("Concurrent modification"))
+            ),
+            "a lost CAS race must be the only expected failure: {errors:?}"
+        );
+
+        let mut sorted_versions = committed_versions.clone();
+        sorted_versions.sort_unstable();
+        let unique_versions: HashSet<_> = sorted_versions.iter().copied().collect();
+        assert_eq!(
+            unique_versions.len(),
+            sorted_versions.len(),
+            "concurrent rotations must never return a version twice: {committed_versions:?}"
+        );
+
+        let successful_rotations = u32::try_from(sorted_versions.len()).expect("test attempts fit u32");
+        let current_version = 1u32 + successful_rotations;
+        assert_eq!(
+            sorted_versions,
+            (2..=current_version).collect::<Vec<_>>(),
+            "committed versions must form one monotonic sequence: {sorted_versions:?}"
+        );
+
+        let snapshot = vault.kv2_snapshot().expect("stateful KV2 snapshot");
+        let persisted_version = snapshot.current_data["version"]
+            .as_u64()
+            .and_then(|version| u32::try_from(version).ok())
+            .expect("current KV2 record must carry a u32 key version");
+        assert_eq!(persisted_version, current_version);
+        assert!(
+            snapshot.current_secret_version >= u64::from(current_version),
+            "the KV2 secret version must advance with each committed pointer switch"
+        );
+        assert_eq!(
+            snapshot.version_records.keys().copied().collect::<Vec<_>>(),
+            (1..=current_version).collect::<Vec<_>>(),
+            "every committed KMS version must have exactly one immutable record"
+        );
+
+        let mut materials = HashSet::new();
+        for (version, record) in &snapshot.version_records {
+            let material = record["encrypted_key_material"]
+                .as_str()
+                .expect("version record must carry encrypted key material");
+            assert!(materials.insert(material), "version {version} reuses another version's material");
+        }
+        assert_eq!(
+            snapshot.current_data["encrypted_key_material"],
+            snapshot
+                .version_records
+                .get(&current_version)
+                .expect("current version record")["encrypted_key_material"],
+            "the top-level fast path must match the current immutable version record"
+        );
+    }
+
     /// The tags write-back after a create is a check-and-set read-modify-write
     /// that carries the key material over from the freshly read record.
     #[tokio::test]
@@ -3465,6 +3579,50 @@ mod tests {
         );
     }
 
+    /// A concurrent first rotation may leave both immutable records visible
+    /// before this caller observes the baseline CAS. That transient state must
+    /// continue through the create-only/adopt path rather than being reported
+    /// as a permanently lost baseline.
+    #[tokio::test]
+    async fn wired_rotate_recovers_concurrent_first_rotation_without_baseline() {
+        let key_data = healthy_key_data();
+        let record_v2 = VaultKeyVersionRecord {
+            version: 2,
+            encrypted_key_material: rotated_material(),
+            created_at: Zoned::now(),
+        };
+
+        let (vault, client) = scripted_client(vec![
+            ScriptedResponse::ok(kv2_metadata_read_data(1)),
+            ScriptedResponse::ok(kv2_read_data(&key_data)),
+            ScriptedResponse::ok(serde_json::json!({ "keys": ["1", "2"] })),
+            // Another rotation already created the baseline record.
+            ScriptedResponse::error(400, CAS_CONFLICT_MESSAGE),
+            ScriptedResponse::ok(kv2_read_version_record_data(&VaultKeyVersionRecord {
+                version: 1,
+                encrypted_key_material: key_data.encrypted_key_material.clone(),
+                created_at: key_data.created_at.clone(),
+            })),
+            // Persist the baseline pointer, then adopt the already-created v2.
+            ScriptedResponse::ok(kv2_write_ack()),
+            ScriptedResponse::error(400, CAS_CONFLICT_MESSAGE),
+            ScriptedResponse::ok(kv2_read_version_record_data(&record_v2)),
+            ScriptedResponse::ok(kv2_write_ack()),
+        ])
+        .await;
+
+        let rotated = client
+            .rotate_key("wired-key", None)
+            .await
+            .expect("an in-flight first rotation must be recoverable");
+        assert_eq!(rotated.version, 2);
+        assert_eq!(
+            parse_write_body(&vault.request_bodies()[8])["data"]["encrypted_key_material"],
+            serde_json::json!(record_v2.encrypted_key_material),
+            "the pointer must adopt the immutable v2 material"
+        );
+    }
+
     /// The mixed-version corruption path: a node older than versioned rotation
     /// performed a lifecycle write, which rewrites the whole key record and
     /// silently drops the `baseline_version` it does not know. Version records