diff --git a/src/request/keyspace.rs b/src/request/keyspace.rs index 79f7225d..366e995c 100644 --- a/src/request/keyspace.rs +++ b/src/request/keyspace.rs @@ -160,23 +160,57 @@ impl TruncateKeyspace for Vec { } } +/// Is this key field UNSET, as opposed to carrying the empty logical key? +/// +/// An encoded key always carries its 4-byte keyspace prefix, so on the wire an empty +/// byte string means "not set" — never "the empty logical key", which encodes to the +/// bare prefix. Guarding on this keeps the codecs off `pretruncate_bytes`' length +/// assertion for fields a shared-lock wrapper leaves unset, and mirrors client-go's +/// `codecV2.DecodeKey`, which returns early on a zero-length key. +fn is_unset_key(key: &[u8]) -> bool { + key.is_empty() +} + impl TruncateKeyspace for Vec { fn truncate_keyspace(mut self, keyspace: Keyspace) -> Self { if !matches!(keyspace, Keyspace::Enable { .. }) { return self; } for lock in &mut self { - take_mut::take(&mut lock.key, |key| { - Key::from(key).truncate_keyspace(keyspace).into() - }); - take_mut::take(&mut lock.primary_lock, |primary| { - Key::from(primary).truncate_keyspace(keyspace).into() - }); + // Convert this lock's OWN key fields, then recurse into any shared-lock + // members — the shape of client-go's `codecV2.decodeLockInfo`, which + // decodes Key/PrimaryLock/Secondaries and *then* walks SharedLockInfos, + // with no wrapper special case. + // + // Per TiKV's writer (`SharedLocks::into_lock_info`, txn_types/src/lock.rs) + // a wrapper sets `lock_type`, `shared_lock_infos` and `key` — the key it + // locks — and leaves `primary_lock`/`lock_version` at their defaults. Each + // member is built from that SAME raw key plus its own primary and version. + // So the wrapper's key must be converted like any other (skipping it would + // hand `scan_locks` a physical key beside decoded member keys), while its + // unset fields must be left alone — hence no wrapper special case, just the + // per-field guard below. + if !is_unset_key(&lock.key) { + take_mut::take(&mut lock.key, |key| { + Key::from(key).truncate_keyspace(keyspace).into() + }); + } + if !is_unset_key(&lock.primary_lock) { + take_mut::take(&mut lock.primary_lock, |primary| { + Key::from(primary).truncate_keyspace(keyspace).into() + }); + } for secondary in lock.secondaries.iter_mut() { + if is_unset_key(secondary) { + continue; + } take_mut::take(secondary, |secondary| { Key::from(secondary).truncate_keyspace(keyspace).into() }); } + take_mut::take(&mut lock.shared_lock_infos, |members| { + members.truncate_keyspace(keyspace) + }); } self } @@ -188,6 +222,17 @@ impl EncodeKeyspace for Vec { return self; } for lock in &mut self { + // Deliberately NOT symmetric with the TruncateKeyspace impl above. There, + // empty bytes can only mean "unset", because an encoded key always carries + // its 4-byte prefix. Here the input is a LOGICAL key, and the empty logical + // key is valid in API v2 — it encodes to the bare prefix. Skipping empties + // would strand a lock on the empty key with no prefix, and `scan_locks` -> + // `resolve_locks` (transaction/client.rs) round-trips exactly that way, so + // its region lookup would then use an empty physical key. + // + // Shared-lock wrappers do not need the unset-field guard here: resolution + // refuses them (`reject_shared_locks`) before anything acts on the encoded + // result. take_mut::take(&mut lock.key, |key| { Key::from(key).encode_keyspace(keyspace, key_mode).into() }); @@ -203,6 +248,9 @@ impl EncodeKeyspace for Vec { .into() }); } + take_mut::take(&mut lock.shared_lock_infos, |members| { + members.encode_keyspace(keyspace, key_mode) + }); } self } @@ -477,4 +525,101 @@ mod tests { let locks = vec![lock]; assert_eq!(locks.clone().truncate_keyspace(keyspace), locks); } + + /// A shared-lock wrapper carries the key it locks, and its members are built from + /// that same raw key (TiKV's `SharedLocks::into_lock_info`). Both must be converted: + /// skipping the wrapper would hand `scan_locks` a physical key beside decoded member + /// keys. The wrapper's *other* fields are a different matter — see + /// [`unset_lock_key_fields_survive_truncation`]. + #[test] + fn shared_lock_wrapper_keys_are_converted_alongside_their_members() { + use crate::proto::kvrpcpb::{LockInfo, Op}; + let keyspace = Keyspace::Enable { keyspace_id: 0 }; + + let wrapper = LockInfo { + key: vec![b'x', 0, 0, 0, b'k'], + lock_type: Op::SharedLock as i32, + shared_lock_infos: vec![LockInfo { + key: vec![b'x', 0, 0, 0, b'm'], + primary_lock: vec![b'x', 0, 0, 0, b'p'], + lock_version: 8, + ..Default::default() + }], + ..Default::default() + }; + + let out = vec![wrapper].truncate_keyspace(keyspace); + assert_eq!( + out[0].key, + vec![b'k'], + "the wrapper's own key must be decoded, not left physical" + ); + assert_eq!(out[0].shared_lock_infos[0].key, vec![b'm']); + assert_eq!(out[0].shared_lock_infos[0].primary_lock, vec![b'p']); + + let back = out.encode_keyspace(keyspace, KeyMode::Txn); + assert_eq!(back[0].key, vec![b'x', 0, 0, 0, b'k']); + assert_eq!(back[0].shared_lock_infos[0].key, vec![b'x', 0, 0, 0, b'm']); + } + + /// The empty logical key is VALID in API v2: it encodes to the bare keyspace + /// prefix. `scan_locks` -> `resolve_locks` round-trips locks through + /// truncate-then-encode, so a lock on the empty key must regain its prefix — + /// otherwise resolution would look up the region for an empty physical key. + /// This is why the encode side does not share the truncate side's unset guard. + #[test] + fn a_lock_on_the_empty_logical_key_round_trips() { + use crate::proto::kvrpcpb::LockInfo; + let keyspace = Keyspace::Enable { keyspace_id: 0 }; + + let physical = vec![LockInfo { + key: vec![b'x', 0, 0, 0], + primary_lock: vec![b'x', 0, 0, 0], + ..Default::default() + }]; + let logical = physical.clone().truncate_keyspace(keyspace); + assert!(logical[0].key.is_empty(), "the empty logical key"); + + let back = logical.encode_keyspace(keyspace, KeyMode::Txn); + assert_eq!( + back, physical, + "a lock on the empty logical key must regain its keyspace prefix" + ); + } + + /// TiKV's `SharedLocks::into_lock_info` leaves a wrapper's `primary_lock` at its + /// default, so the codec meets genuinely unset fields in practice — they must be + /// skipped, not run into `pretruncate_bytes`' length assertion. + #[test] + fn unset_lock_key_fields_survive_truncation() { + use crate::proto::kvrpcpb::{LockInfo, Op}; + let keyspace = Keyspace::Enable { keyspace_id: 0 }; + + // The writer's actual shape: key and members set, primary_lock left default. + let realistic = LockInfo { + key: vec![b'x', 0, 0, 0, b'k'], + lock_type: Op::SharedLock as i32, + shared_lock_infos: vec![LockInfo { + key: vec![b'x', 0, 0, 0, b'k'], + primary_lock: vec![b'x', 0, 0, 0, b'p'], + ..Default::default() + }], + ..Default::default() + }; + let out = vec![realistic].truncate_keyspace(keyspace); + assert_eq!(out[0].key, vec![b'k']); + assert!( + out[0].primary_lock.is_empty(), + "the wrapper's unset primary must be skipped, not truncated" + ); + assert_eq!(out[0].shared_lock_infos[0].primary_lock, vec![b'p']); + + // Defensively, a wrapper with no key at all must not panic either. + let keyless = LockInfo { + lock_type: Op::SharedLock as i32, + ..Default::default() + }; + let out = vec![keyless].truncate_keyspace(keyspace); + assert!(out[0].key.is_empty()); + } } diff --git a/src/request/plan.rs b/src/request/plan.rs index b7fac56f..6ef78f93 100644 --- a/src/request/plan.rs +++ b/src/request/plan.rs @@ -841,6 +841,10 @@ where has_more_batch = false; } + // BEFORE any filter: a shared-lock wrapper's fields (including + // `use_async_commit`) must not be read — filtering on them would silently + // drop the real member locks. Refuse instead; see `reject_shared_locks`. + crate::transaction::reject_shared_locks(&locks)?; if self.options.async_commit_only { locks = locks .into_iter() diff --git a/src/transaction/lock.rs b/src/transaction/lock.rs index 4a8adb63..98bb3683 100644 --- a/src/transaction/lock.rs +++ b/src/transaction/lock.rs @@ -43,6 +43,35 @@ pub(crate) fn format_key_for_log(key: &[u8]) -> String { format!("len={}, prefix={}", key.len(), HexRepr(&key[..prefix_len])) } +/// Refuse to resolve SHARED locks — loudly, before any of them can be mis-handled. +/// +/// The contract (`kvrpcpb.LockInfo.shared_lock_infos`) is explicit: a shared lock's +/// real holders live ONLY in `shared_lock_infos` — "DO NOT read from the wrapper +/// LockInfo", whose own `key`/`lock_version` are unset. This client does not implement +/// shared-lock resolution yet, and every partial handling is worse than none: +/// resolving the wrapper checks transaction 0; filtering on wrapper fields silently +/// drops the members; and the pessimistic-lock special cases in this resolver do not +/// know `SharedPessimisticLock`. Until support lands, an explicit error is the only +/// answer that cannot roll back a live transaction or skip a dead one. +/// +/// Servers that predate shared locks never produce them, so this is a no-op there. +pub(crate) fn reject_shared_locks(locks: &[kvrpcpb::LockInfo]) -> Result<()> { + let shared = |l: &kvrpcpb::LockInfo| { + !l.shared_lock_infos.is_empty() + || l.lock_type == kvrpcpb::Op::SharedLock as i32 + || l.lock_type == kvrpcpb::Op::SharedPessimisticLock as i32 + }; + if locks.iter().any(shared) { + return Err(Error::StringError( + "shared locks (SharedLock/SharedPessimisticLock) are not supported by this \ + client yet; refusing to resolve them — resolving the wrapper would target \ + the wrong transaction" + .to_owned(), + )); + } + Ok(()) +} + /// _Resolves_ the given locks. Returns locks still live. When there is no live locks, all the given locks are resolved. /// /// If a key has a lock, the latest status of the key is unknown. We need to "resolve" the lock, @@ -56,6 +85,7 @@ pub async fn resolve_locks( keyspace: Keyspace, ) -> Result /* live_locks */> { debug!("resolving locks"); + reject_shared_locks(&locks)?; let ts = pd_client.clone().get_timestamp().await?; let caller_start_ts = timestamp.version(); let current_ts = ts.version(); @@ -300,6 +330,9 @@ impl LockResolver { pd_client: Arc, // TODO: make pd_client a member of LockResolver keyspace: Keyspace, ) -> Result<()> { + // Defense in depth: CleanupLocks::execute refuses these before its filters, + // but this entry point is public within the crate. + reject_shared_locks(&locks)?; if locks.is_empty() { return Ok(()); } @@ -619,6 +652,35 @@ mod tests { use crate::mock::MockPdClient; use crate::proto::errorpb; + #[test] + fn shared_locks_are_refused_never_misresolved() { + let plain = kvrpcpb::LockInfo { + key: b"k1".to_vec(), + lock_version: 7, + ..Default::default() + }; + assert!(reject_shared_locks(std::slice::from_ref(&plain)).is_ok()); + + // A wrapper: key/lock_version deliberately unset per the contract — resolving + // it would check transaction 0. Must be refused, not resolved or filtered. + let wrapper = kvrpcpb::LockInfo { + shared_lock_infos: vec![kvrpcpb::LockInfo { + key: b"k2".to_vec(), + lock_version: 8, + ..Default::default() + }], + ..Default::default() + }; + assert!(reject_shared_locks(&[plain.clone(), wrapper]).is_err()); + + // Also refused when only the op marks it shared (empty member list). + let by_op = kvrpcpb::LockInfo { + lock_type: kvrpcpb::Op::SharedPessimisticLock as i32, + ..Default::default() + }; + assert!(reject_shared_locks(&[by_op]).is_err()); + } + #[rstest::rstest] #[case(Keyspace::Disable)] #[case(Keyspace::Enable { keyspace_id: 0 })] diff --git a/src/transaction/mod.rs b/src/transaction/mod.rs index 1f2cc181..573992b0 100644 --- a/src/transaction/mod.rs +++ b/src/transaction/mod.rs @@ -28,6 +28,7 @@ mod client; mod lock; pub mod lowering; mod requests; +pub(crate) use lock::reject_shared_locks; pub use lock::LockResolver; pub use lock::ResolveLocksContext; pub use lock::ResolveLocksOptions;