transaction: refuse to resolve shared locks instead of mis-handling them

src/request/keyspace.rs

@@ -160,23 +160,57 @@ impl TruncateKeyspace for Vec<KvPair> {
     }
 }
 
+/// 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<crate::proto::kvrpcpb::LockInfo> {
     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<crate::proto::kvrpcpb::LockInfo> {
             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<crate::proto::kvrpcpb::LockInfo> {
                         .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());
+    }
 }

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()

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<Vec<kvrpcpb::LockInfo> /* 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<impl PdClient>, // 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 })]

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;