feat(kms): reserve and expose KV2 wrap-budget consumption (#6019)

2026-08-19 02:56:18 +00:00 · 2026-08-13 08:05:00 +08:00
parent 4a41325d1a
commit ca06c7ec2c
10 changed files with 538 additions and 15 deletions
@@ -657,6 +657,7 @@ impl KmsBackend for AwsKmsBackend {
                created_by: None,
                rotation_due: false,
                rotation_due_reason: None,
+                wrap_budget_reserved: None,
            });
        }

@@ -271,6 +271,7 @@ impl StaticKmsBackend {
            created_by: None,
            rotation_due: false,
            rotation_due_reason: None,
+            wrap_budget_reserved: None,
        })
    }

@@ -58,6 +58,12 @@ pub struct VaultKmsClient {
    /// triggered — shutdown drops the whole client — but kept as the single
    /// hook a future lifecycle owner can cancel through.
    cancel: CancellationToken,
+    /// Per-key in-process remainder of the persisted wrap-budget reservation
+    /// (see [`VaultKmsClient::consume_wrap_budget`]). Grows with the master
+    /// keys this node wraps under and is never pruned; that set is small by
+    /// construction. The outer lock is only ever held to look up or insert the
+    /// per-key entry, never across an await.
+    wrap_budgets: std::sync::Mutex<HashMap<String, Arc<tokio::sync::Mutex<WrapBudget>>>>,
 }

 /// Key data stored in Vault
@@ -112,6 +118,51 @@ struct VaultKeyData {
    /// deserializing.
    #[serde(default)]
    baseline_version: Option<u32>,
+    /// Wrap operations reserved against this key's *current* master key
+    /// material, in blocks of [`WRAP_BUDGET_BLOCK`].
+    ///
+    /// AES-256-GCM caps one key at 2^32 encryptions under random 96-bit nonces
+    /// (NIST SP 800-38D), and this backend wraps every DEK locally with the
+    /// current material, so this approximates how much of that bound the
+    /// cluster has consumed. Nodes reserve whole blocks up front and count
+    /// individual wraps in process memory only, so the persisted value can run
+    /// ahead of the wraps actually performed but — on builds that know the
+    /// field — never behind: a crash discards unused in-memory budget, never a
+    /// counted wrap. Two documented ways the value can still understate: wraps
+    /// performed while a reservation write kept failing (logged at warn, and
+    /// re-covered by the next reservation that lands), and an old build
+    /// rewriting this record on any lifecycle write, which drops the field it
+    /// does not know and regresses the count to zero.
+    ///
+    /// Reset to 0 by [`VaultKmsClient::rotate_key`]'s pointer-switch commit:
+    /// the GCM bound is per key material, and rotation installs fresh material.
+    #[serde(default)]
+    wrap_budget_reserved: u64,
+}
+
+/// Wrap operations reserved from the key record per reservation write.
+///
+/// Large enough that the once-per-block CAS write disappears against a million
+/// data-path wraps, small enough that the crash-time overestimate (at most one
+/// discarded block per node) stays negligible against the 2^32 bound.
+const WRAP_BUDGET_BLOCK: u64 = 1_000_000;
+
+/// In-process remainder of one key's persisted wrap-budget reservation.
+#[derive(Debug, Default)]
+struct WrapBudget {
+    /// Master key version the grant was taken against, only ever moved
+    /// forward. A *newer* version about to wrap means the key rotated: fresh
+    /// material has a fresh nonce budget and a zeroed persisted counter, so
+    /// the stale grant (and any stale debt — the old material never wraps
+    /// again) is discarded. An *older* one is a wrap whose snapshot lost a
+    /// race with a rotation and is simply counted against the current grant.
+    version: u32,
+    /// Wraps still covered by the last block grant.
+    available: u64,
+    /// Budget granted in memory while reservation writes were failing — wraps
+    /// the persisted counter does not cover yet. Added onto the next
+    /// successful reservation so the persisted count catches back up.
+    unpersisted: u64,
 }

 impl UnknownFieldSummary {
@@ -469,9 +520,78 @@ impl VaultKmsClient {
            dek_crypto: AesDekCrypto::new(),
            retry: RetryPolicy::for_backend(kms_config, "vault-kv2", &config.address, config.namespace.as_deref(), "operations"),
            cancel: CancellationToken::new(),
+            wrap_budgets: std::sync::Mutex::new(HashMap::new()),
        })
    }

+    /// Count one DEK wrap against the key's persisted wrap budget.
+    ///
+    /// `key_version` is the master key version whose material is about to
+    /// wrap, from the same record snapshot the wrap itself uses. Budget is
+    /// taken from an in-process block; only when the block is exhausted (or
+    /// the key rotated under it) is a new block of [`WRAP_BUDGET_BLOCK`]
+    /// reserved by a check-and-set update of the key record — never a write
+    /// per wrap, so the data path pays one extra Vault round trip per million
+    /// wraps, not per object.
+    ///
+    /// Reserve-then-consume on purpose: the reservation lands before the wrap
+    /// it covers, so a crash can only ever discard reserved-but-unused budget
+    /// — the persisted count overestimates, never undercounts. The counter is
+    /// advisory observability, not a quota: a reservation that cannot be
+    /// persisted is logged and the wrap proceeds on an in-memory grant carried
+    /// as `unpersisted` debt, which the next successful reservation adds on
+    /// top of its own block. That fail-open grant is also what bounds the
+    /// warn to at most one per block of wraps. Infallible by design — no
+    /// Vault hiccup here may fail a PUT.
+    ///
+    /// Concurrent wraps of the same key briefly queue on the per-key lock
+    /// while the once-per-block reservation is in flight instead of each
+    /// issuing their own.
+    async fn consume_wrap_budget(&self, key_id: &str, key_version: u32) {
+        let budget = {
+            let mut budgets = self.wrap_budgets.lock().expect("wrap budget map lock poisoned");
+            match budgets.get(key_id) {
+                // Fast path spares the per-wrap key allocation `entry` needs.
+                Some(budget) => Arc::clone(budget),
+                None => Arc::clone(budgets.entry(key_id.to_string()).or_default()),
+            }
+        };
+        let mut budget = budget.lock().await;
+        if key_version > budget.version {
+            *budget = WrapBudget {
+                version: key_version,
+                ..WrapBudget::default()
+            };
+        }
+        // `key_version < budget.version` is a wrap whose record snapshot lost a
+        // race with a rotation. It is counted against the current grant rather
+        // than resetting to the old version: the newer grant is persisted on
+        // the post-rotation record, so the count stays an overestimate, and a
+        // burst of in-flight stale wraps cannot ping-pong the version tag into
+        // one reservation write each.
+        if budget.available == 0 {
+            let requested = WRAP_BUDGET_BLOCK.saturating_add(budget.unpersisted);
+            let reserved = self
+                .update_key_data_with_cas(key_id, |key_data| {
+                    key_data.wrap_budget_reserved = key_data.wrap_budget_reserved.saturating_add(requested);
+                    Ok(CasMutation::Write(()))
+                })
+                .await;
+            match reserved {
+                Ok(_) => {
+                    budget.available = WRAP_BUDGET_BLOCK;
+                    budget.unpersisted = 0;
+                }
+                Err(error) => {
+                    budget.available = WRAP_BUDGET_BLOCK;
+                    budget.unpersisted = requested;
+                    warn!(key_id, requested, %error, "Vault KMS wrap budget reservation failed; wraps continue uncounted");
+                }
+            }
+        }
+        budget.available -= 1;
+    }
+
    /// Snapshot the authenticated Vault client for a single request.
    ///
    /// Every Vault call takes its own snapshot so a credential rotation
@@ -1004,6 +1124,7 @@ impl VaultKmsClient {
            decode_stored_key_material(&request.master_key_id, &key_data.encrypted_key_material).inspect_err(|error| {
                warn!(key_id = %request.master_key_id, %error, "Vault KMS key material failed validation");
            })?;
+        self.consume_wrap_budget(&request.master_key_id, key_data.version).await;
        let (encrypted_key, nonce) = self.dek_crypto.encrypt(&key_material, &plaintext_key).await?;

        // Create data key envelope with master key version for rotation support
@@ -1037,6 +1158,7 @@ impl VaultKmsClient {
        ensure_key_status_permits(&request.key_id, &key_data.status, StateGatedOperation::Encrypt)?;
        let key_material = decode_stored_key_material(&request.key_id, &key_data.encrypted_key_material)
            .inspect_err(|error| warn!(key_id = %request.key_id, %error, "Vault KMS key material failed validation"))?;
+        self.consume_wrap_budget(&request.key_id, key_data.version).await;
        let (encrypted_key, nonce) = self.dek_crypto.encrypt(&key_material, &request.plaintext).await?;

        // Wrap the ciphertext in the same authenticated envelope that
@@ -1207,6 +1329,11 @@ impl VaultKmsClient {
            });
        }

+        // The re-wrap below encrypts with the current material under a fresh
+        // random nonce — one wrap off the budget, counted before any plaintext
+        // exists so the accounting never extends the plaintext's lifetime.
+        self.consume_wrap_budget(&envelope.master_key_id, current_version).await;
+
        let source_version =
            resolve_envelope_master_key_version(envelope.master_key_version, key_data.baseline_version, current_version);
        // Both materials are resolved before anything is unwrapped, so no
@@ -1372,6 +1499,7 @@ impl VaultKmsClient {
            rotated_at: None,
            encrypted_key_material: encrypted_material,
            baseline_version: None,
+            wrap_budget_reserved: 0,
        };

        // Create-only write: the not-found pre-check above is only advisory —
@@ -1420,6 +1548,7 @@ impl VaultKmsClient {
            created_by: None,
            rotation_due: false,
            rotation_due_reason: None,
+            wrap_budget_reserved: Some(key_data.wrap_budget_reserved),
        })
    }

@@ -1670,6 +1799,10 @@ impl VaultKmsClient {
        key_data.version = new_version;
        key_data.encrypted_key_material = new_material;
        key_data.rotated_at = Some(Zoned::now());
+        // Fresh material, fresh AES-GCM nonce budget: the wrap ceiling is per
+        // key material version, so the counter restarts with the same commit
+        // that makes the new material current.
+        key_data.wrap_budget_reserved = 0;
        self.cas_store_key_data(key_id, &key_data, cas).await?;

        info!(key_id, version = new_version, "Vault KMS master key rotated");
@@ -2176,6 +2309,7 @@ mod tests {
            rotated_at: None,
            encrypted_key_material: general_purpose::STANDARD.encode([0x42u8; 32]),
            baseline_version: None,
+            wrap_budget_reserved: 0,
        }
    }

@@ -2639,6 +2773,7 @@ mod tests {
            baseline_version: Some(1),
            deletion_date: None,
            rotated_at: None,
+            wrap_budget_reserved: 0,
        };

        let mut value = serde_json::to_value(&key_data).expect("serialize key data");
@@ -3041,6 +3176,7 @@ mod tests {
            rotated_at: None,
            encrypted_key_material: "material".to_string(),
            baseline_version: None,
+            wrap_budget_reserved: 0,
        };

        let mut value = serde_json::to_value(&key_data).expect("serialize");
@@ -3073,9 +3209,13 @@ mod tests {

    #[tokio::test]
    async fn wired_kv2_encrypt_round_trips_through_decrypt() {
-        // One key-record read for the encrypt, one for the decrypt.
+        // One key-record read plus the first wrap's budget reservation for the
+        // encrypt, one read for the decrypt.
        let (_vault, client) = scripted_client(vec![
            ScriptedResponse::ok(kv2_read_data(&healthy_key_data())),
+            ScriptedResponse::ok(kv2_metadata_read_data(1)),
+            ScriptedResponse::ok(kv2_read_data(&healthy_key_data())),
+            ScriptedResponse::ok(kv2_write_ack()),
            ScriptedResponse::ok(kv2_read_data(&healthy_key_data())),
        ])
        .await;
@@ -4491,9 +4631,17 @@ mod tests {
        }
    }

-    /// Encrypt against a scripted Vault serving `state`.
+    /// Encrypt against a scripted Vault serving `state`. The fresh client's
+    /// first wrap also reserves its budget block, so that exchange is scripted
+    /// alongside the key-record read.
    async fn encrypt_scripted(state: &KeyState, plaintext: &[u8]) -> EncryptResponse {
-        let (_vault, client) = scripted_client(vec![ScriptedResponse::ok(kv2_read_data(&state.key_data))]).await;
+        let (_vault, client) = scripted_client(vec![
+            ScriptedResponse::ok(kv2_read_data(&state.key_data)),
+            ScriptedResponse::ok(kv2_metadata_read_data(1)),
+            ScriptedResponse::ok(kv2_read_data(&state.key_data)),
+            ScriptedResponse::ok(kv2_write_ack()),
+        ])
+        .await;
        client
            .encrypt(
                &EncryptRequest {
@@ -4654,12 +4802,19 @@ mod tests {
        }
    }

-    /// Rewrap against a scripted Vault serving `state`, scripting the key record
-    /// plus every version record the state holds, so the implementation — not
-    /// the harness — decides which of them it needs. Returns the response
-    /// together with the requests the rewrap made.
+    /// Rewrap against a scripted Vault serving `state`, scripting the key
+    /// record, the fresh client's first-wrap budget reservation, and every
+    /// version record the state holds, so the implementation — not the harness
+    /// — decides which of them it needs (a no-op rewrap consumes neither the
+    /// reservation nor a version record). Returns the response together with
+    /// the requests the rewrap made.
    async fn rewrap_scripted(state: &KeyState, ciphertext: &[u8]) -> (RewrapDataKeyResponse, Vec<String>) {
-        let mut responses = vec![ScriptedResponse::ok(kv2_read_data(&state.key_data))];
+        let mut responses = vec![
+            ScriptedResponse::ok(kv2_read_data(&state.key_data)),
+            ScriptedResponse::ok(kv2_metadata_read_data(1)),
+            ScriptedResponse::ok(kv2_read_data(&state.key_data)),
+            ScriptedResponse::ok(kv2_write_ack()),
+        ];
        responses.extend(
            state
                .version_records
@@ -4727,6 +4882,11 @@ mod tests {
            requests,
            vec![
                "GET /v1/secret/data/rustfs/kms/keys/wired-key".to_string(),
+                // The fresh client's first wrap reserves its budget block...
+                "GET /v1/secret/metadata/rustfs/kms/keys/wired-key".to_string(),
+                "GET /v1/secret/data/rustfs/kms/keys/wired-key?version=1".to_string(),
+                "POST /v1/secret/data/rustfs/kms/keys/wired-key".to_string(),
+                // ...and the unwrap must resolve the frozen version-1 material.
                "GET /v1/secret/data/rustfs/kms/keys/wired-key/versions/1".to_string(),
            ],
            "the unwrap must resolve the frozen version-1 material: {requests:?}"
@@ -4817,7 +4977,10 @@ mod tests {
        assert!(response.rewrapped);
        assert_eq!(response.source_key_version, Some(1));
        assert_eq!(response.destination_key_version, Some(1));
-        assert_eq!(requests.len(), 1, "a never-rotated key has no version record to read: {requests:?}");
+        assert!(
+            !requests.iter().any(|line| line.contains("/versions/")),
+            "a never-rotated key has no version record to read: {requests:?}"
+        );
        let stamped: DataKeyEnvelope = serde_json::from_slice(&response.ciphertext).expect("stamped envelope must parse");
        assert_eq!(stamped.master_key_version, Some(1));
        let (plaintext, _) = decrypt_scripted(&state_v1, &response.ciphertext).await;
@@ -4831,7 +4994,7 @@ mod tests {
        assert_eq!(rotated_response.source_key_version, Some(1), "the baseline is what wrapped it");
        assert_eq!(rotated_response.destination_key_version, Some(2));
        assert!(
-            rotated_requests[1].ends_with("/versions/1"),
+            rotated_requests.iter().any(|line| line.ends_with("/versions/1")),
            "the unwrap must resolve the baseline material: {rotated_requests:?}"
        );
        let (rotated_plaintext, _) = decrypt_scripted(&state_v2, &rotated_response.ciphertext).await;
@@ -4855,7 +5018,16 @@ mod tests {
    #[tokio::test]
    async fn wired_kv2_rewrap_rejects_a_tampered_encryption_context() {
        let context = HashMap::from([("bucket".to_string(), "photos/cat.jpg".to_string())]);
-        let (vault, client) = scripted_client(vec![ScriptedResponse::ok(kv2_read_data(&healthy_key_data()))]).await;
+        // The scripted exchange covers the encrypt (key read plus the first
+        // wrap's budget reservation) and nothing else: the refused calls below
+        // must not add a single request.
+        let (vault, client) = scripted_client(vec![
+            ScriptedResponse::ok(kv2_read_data(&healthy_key_data())),
+            ScriptedResponse::ok(kv2_metadata_read_data(1)),
+            ScriptedResponse::ok(kv2_read_data(&healthy_key_data())),
+            ScriptedResponse::ok(kv2_write_ack()),
+        ])
+        .await;

        let encrypted = client
            .encrypt(
@@ -4890,9 +5062,254 @@ mod tests {

        assert_eq!(
            vault.requests().len(),
-            1,
-            "the context guard must run before any Vault read: {:?}",
+            4,
+            "the context guard must run before any Vault read — every request must belong to the encrypt: {:?}",
            vault.requests()
        );
    }
+
+    /// The wrap counter is block-reserved, never written per wrap: N wraps
+    /// with N < [`WRAP_BUDGET_BLOCK`] perform exactly one check-and-set
+    /// reservation write, and the persisted value is the full block — an
+    /// overestimate of the wraps actually performed, which is the direction a
+    /// crash must leave it in (the unused in-memory remainder dies with the
+    /// process, counted wraps never do). A revert to per-wrap persistence
+    /// fails the write count; a revert to not persisting at all fails the
+    /// stored value.
+    #[tokio::test]
+    async fn wired_generate_data_key_reserves_wrap_budget_in_blocks() {
+        let (vault, client) = scripted_kv2_client(&healthy_key_data()).await;
+        let request = integration_generate_request("wired-key");
+
+        const WRAPS: u64 = 3;
+        for _ in 0..WRAPS {
+            client
+                .generate_data_key(&request, None)
+                .await
+                .expect("wraps within the reserved block must succeed");
+        }
+
+        let requests = vault.requests();
+        assert_eq!(
+            requests.iter().filter(|line| line.starts_with("POST ")).count(),
+            1,
+            "{WRAPS} wraps inside one block must reserve exactly once: {requests:?}"
+        );
+
+        // "Crash": drop the client and its in-memory remainder, then read what
+        // Vault durably holds.
+        drop(client);
+        let snapshot = vault.kv2_snapshot().expect("stateful KV2 snapshot");
+        let reserved = snapshot.current_data["wrap_budget_reserved"]
+            .as_u64()
+            .expect("the key record must carry the reserved wrap budget");
+        assert_eq!(reserved, WRAP_BUDGET_BLOCK, "the reservation persists the whole block up front");
+        assert!(reserved >= WRAPS, "the persisted count must never understate the wraps performed");
+    }
+
+    /// Two nodes reserving against the same record must accumulate, not
+    /// clobber: the loser of the check-and-set race re-reads the record the
+    /// winner committed and adds its block on top, ending at two blocks. A
+    /// blind write here would silently erase the peer's reservation and
+    /// undercount its million wraps.
+    #[tokio::test]
+    async fn wired_wrap_budget_reservation_adds_on_top_after_losing_a_cas_race() {
+        let mut peer_reserved = healthy_key_data();
+        peer_reserved.wrap_budget_reserved = WRAP_BUDGET_BLOCK;
+
+        let (vault, client) = scripted_client(vec![
+            // The encrypt reads the key record...
+            ScriptedResponse::ok(kv2_read_data(&healthy_key_data())),
+            // ...and its first wrap reserves: attempt 1 observes no
+            // reservation yet...
+            ScriptedResponse::ok(kv2_metadata_read_data(1)),
+            ScriptedResponse::ok(kv2_read_data(&healthy_key_data())),
+            // ...but the peer's reservation committed in between, so the
+            // check-and-set write loses.
+            ScriptedResponse::error(400, CAS_CONFLICT_MESSAGE),
+            // Attempt 2 re-reads the record the peer committed and lands.
+            ScriptedResponse::ok(kv2_metadata_read_data(2)),
+            ScriptedResponse::ok(kv2_read_data(&peer_reserved)),
+            ScriptedResponse::ok(kv2_write_ack()),
+        ])
+        .await;
+
+        client
+            .encrypt(
+                &EncryptRequest {
+                    key_id: "wired-key".to_string(),
+                    plaintext: b"counted-once".to_vec(),
+                    encryption_context: HashMap::new(),
+                    grant_tokens: Vec::new(),
+                },
+                None,
+            )
+            .await
+            .expect("losing the reservation race must not fail the wrap");
+
+        let bodies = vault.request_bodies();
+        let lost = parse_write_body(&bodies[3]);
+        assert_eq!(lost["options"]["cas"], serde_json::json!(1), "{lost}");
+        assert_eq!(lost["data"]["wrap_budget_reserved"], serde_json::json!(WRAP_BUDGET_BLOCK), "{lost}");
+        let committed = parse_write_body(&bodies[6]);
+        assert_eq!(committed["options"]["cas"], serde_json::json!(2), "{committed}");
+        assert_eq!(
+            committed["data"]["wrap_budget_reserved"],
+            serde_json::json!(2 * WRAP_BUDGET_BLOCK),
+            "the retry must add its block on top of the peer's, not overwrite it: {committed}"
+        );
+    }
+
+    /// The counter is advisory observability, not a quota: a reservation that
+    /// cannot be persisted is warned about and the wrap proceeds. Turning the
+    /// scripted 5xx into a failed `generate_data_key` — a Vault hiccup failing
+    /// a PUT — is exactly the regression this test pins down.
+    #[tokio::test]
+    async fn wired_wrap_proceeds_and_warns_when_budget_reservation_fails() {
+        let logs = crate::test_support::CapturedLogs::default();
+        let subscriber = tracing_subscriber::fmt()
+            .with_ansi(false)
+            .with_max_level(tracing::Level::WARN)
+            .with_writer(logs.clone())
+            .finish();
+        let _guard = tracing::subscriber::set_default(subscriber);
+
+        let (vault, client) = scripted_client(vec![
+            // generate_data_key: the state-gate read and the wrap snapshot.
+            ScriptedResponse::ok(kv2_read_data(&healthy_key_data())),
+            ScriptedResponse::ok(kv2_read_data(&healthy_key_data())),
+            // The reservation's versioned read succeeds...
+            ScriptedResponse::ok(kv2_metadata_read_data(1)),
+            ScriptedResponse::ok(kv2_read_data(&healthy_key_data())),
+            // ...and its write fails.
+            ScriptedResponse::error(503, "sealed"),
+        ])
+        .await;
+
+        let data_key = client
+            .generate_data_key(&integration_generate_request("wired-key"), None)
+            .await
+            .expect("a failed budget reservation must never fail the wrap");
+        assert!(data_key.plaintext.is_some());
+        assert!(!data_key.ciphertext.is_empty());
+
+        let requests = vault.requests();
+        assert_eq!(requests.len(), 5, "{requests:?}");
+        assert_eq!(
+            requests.iter().filter(|line| line.starts_with("POST ")).count(),
+            1,
+            "the failed non-idempotent reservation write must not be replayed: {requests:?}"
+        );
+
+        let output = logs.output();
+        assert!(output.contains("WARN"), "the failure must be visible to operators: {output}");
+        assert!(
+            output.contains("Vault KMS wrap budget reservation failed"),
+            "the warn must name the degraded counter: {output}"
+        );
+        assert!(
+            !output.contains(&healthy_key_data().encrypted_key_material),
+            "the warn must not echo stored key material: {output}"
+        );
+    }
+
+    /// The AES-GCM wrap bound is per key material version, so the rotation's
+    /// pointer-switch commit — and only that commit — resets the persisted
+    /// counter. The baseline-pin write before it must still carry the
+    /// pre-rotation count (the old material is still current there), and the
+    /// immutable version records never carry the field at all.
+    #[tokio::test]
+    async fn wired_rotate_resets_wrap_budget_with_the_pointer_switch() {
+        let mut key_data = healthy_key_data();
+        key_data.wrap_budget_reserved = 123_456;
+        let (vault, client) = scripted_kv2_client(&key_data).await;
+
+        client.rotate_key("wired-key", None).await.expect("rotation must commit");
+
+        let snapshot = vault.kv2_snapshot().expect("stateful KV2 snapshot");
+        assert_eq!(snapshot.current_data["version"], serde_json::json!(2));
+        assert_eq!(
+            snapshot.current_data["wrap_budget_reserved"],
+            serde_json::json!(0),
+            "fresh material must start with a fresh nonce budget"
+        );
+        assert!(
+            snapshot
+                .version_records
+                .get(&1)
+                .expect("the first rotation must freeze the version-1 record")
+                .get("wrap_budget_reserved")
+                .is_none(),
+            "version records carry material, never the wrap counter"
+        );
+
+        // First rotation writes: freeze v1, pin the baseline, create v2,
+        // switch the pointer. Only the last one resets the counter.
+        let bodies = vault.request_bodies();
+        let baseline_pin = parse_write_body(&bodies[4]);
+        assert_eq!(
+            baseline_pin["data"]["wrap_budget_reserved"],
+            serde_json::json!(123_456),
+            "the pre-switch write still describes the old material: {baseline_pin}"
+        );
+        let switch = parse_write_body(&bodies[6]);
+        assert_eq!(switch["data"]["version"], serde_json::json!(2), "{switch}");
+        assert_eq!(switch["data"]["wrap_budget_reserved"], serde_json::json!(0), "{switch}");
+    }
+
+    /// The in-memory block is tied to the material version it was reserved
+    /// against: a wrap after a rotation must not spend the stale grant (whose
+    /// persisted counter the rotation just reset) but reserve a fresh block on
+    /// the new version's record.
+    #[tokio::test]
+    async fn wired_wrap_after_rotation_reserves_a_fresh_block() {
+        let (vault, client) = scripted_kv2_client(&healthy_key_data()).await;
+        let request = integration_generate_request("wired-key");
+
+        client.generate_data_key(&request, None).await.expect("wrap under v1");
+        client.rotate_key("wired-key", None).await.expect("rotate to v2");
+        client.generate_data_key(&request, None).await.expect("wrap under v2");
+
+        let snapshot = vault.kv2_snapshot().expect("stateful KV2 snapshot");
+        assert_eq!(snapshot.current_data["version"], serde_json::json!(2));
+        assert_eq!(
+            snapshot.current_data["wrap_budget_reserved"],
+            serde_json::json!(WRAP_BUDGET_BLOCK),
+            "a wrap under fresh material must reserve anew instead of spending the stale grant"
+        );
+    }
+
+    /// `describe_key` reports the persisted counter — the deletion worker's
+    /// census reads it from exactly this surface to publish the aggregate
+    /// wrap gauge.
+    #[tokio::test]
+    async fn wired_describe_key_reports_wrap_budget_for_the_census() {
+        let mut key_data = healthy_key_data();
+        key_data.wrap_budget_reserved = 42;
+        let (_vault, client) = scripted_client(vec![ScriptedResponse::ok(kv2_read_data(&key_data))]).await;
+
+        let described = client.describe_key("wired-key", None).await.expect("describe must succeed");
+        assert_eq!(described.wrap_budget_reserved, Some(42));
+    }
+
+    /// The persisted KV2 record round-trips its wrap counter, and a record
+    /// without the field — written by an older build, or rewritten by one,
+    /// which is the documented mixed-version regression — reads back as zero.
+    #[test]
+    fn vault_key_data_wrap_budget_round_trips_and_defaults_to_zero() {
+        let mut key_data = healthy_key_data();
+        key_data.wrap_budget_reserved = 42;
+
+        let mut value = serde_json::to_value(&key_data).expect("serialize");
+        let restored: VaultKeyData = serde_json::from_value(value.clone()).expect("round trip");
+        assert_eq!(restored.wrap_budget_reserved, 42);
+
+        value
+            .as_object_mut()
+            .expect("record must be a JSON object")
+            .remove("wrap_budget_reserved")
+            .expect("current records must carry the field");
+        let legacy: VaultKeyData = serde_json::from_value(value).expect("legacy record must deserialize");
+        assert_eq!(legacy.wrap_budget_reserved, 0);
+    }
 }
@@ -874,6 +874,7 @@ impl VaultTransitKmsClient {
            created_by: metadata.created_by,
            rotation_due: false,
            rotation_due_reason: None,
+            wrap_budget_reserved: None,
        })
    }