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

The re-vendored kvproto exposes shared locks (Op::SharedLock /
Op::SharedPessimisticLock). Their contract is unusual: a shared lock's real
holders live ONLY in kvrpcpb.LockInfo.shared_lock_infos — the proto comment is
explicit that you must "DO NOT read from the wrapper LockInfo", whose own key
and lock_version are unset.

This client does not implement shared-lock resolution, and every partial
handling of one is worse than none: resolving the wrapper checks transaction 0;
filtering on the wrapper's fields (use_async_commit, lock_type) silently drops
the real members; and the pessimistic-lock special cases in the resolver do not
know SharedPessimisticLock. So resolution REFUSES them with an explicit error —
in resolve_locks, in LockResolver's crate-public entry point, and in
CleanupLocks::execute before any filter runs. Until real 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.

The API-v2 keyspace codecs are made shared-lock-aware separately, because they
run on scan results that never reach the resolver. They now take the shape of
client-go's codecV2.decodeLockInfo (internal/apicodec/codec_v2.go), which
decodes a LockInfo's own Key/PrimaryLock/Secondaries and then recurses into
SharedLockInfos — it has no wrapper special case. This client now does the
same: convert the lock's OWN key fields, then recurse into the members.

An earlier revision skipped a wrapper's own fields entirely, reading the proto's
"DO NOT read from the wrapper LockInfo" as covering them. Checking the writer
(TiKV's SharedLocks::into_lock_info in components/txn_types/src/lock.rs,
identical at v8.5.7 and on master) shows
that is only half right, and the half it gets wrong matters:

  info.set_shared_lock_infos(shared_locks.into());
  info.set_key(raw_key);

A wrapper sets lock_type, shared_lock_infos and key — the key it locks — and
leaves primary_lock and lock_version at their defaults; each member is built
from that SAME raw key plus its own primary and version. So the caveat scopes
the per-transaction fields, not the locked key. Skipping the wrapper wholesale
would have handed scan_locks a physical key beside decoded member keys, while
converting it wholesale would have hit the length assertion on the unset
primary. Both fields are exercised by tests.

The two directions guard differently, on purpose. Truncating,
empty bytes can only mean "unset" — an encoded key always carries its 4-byte
prefix — so unset fields are skipped rather than run into pretruncate_bytes'
length assertion. Encoding, the input is a LOGICAL key and the empty logical key
is valid in API v2, so nothing is skipped: scan_locks -> resolve_locks
round-trips locks through truncate-then-encode, and skipping empties there would
strand a lock on the empty key with no prefix and send resolution after an empty
physical key. Wrappers need no encode-side guard because resolution refuses them
first. That panic hazard predates the re-vendor: a wrapper arrives the same way
whichever proto vintage parses it.

Full shared-lock support is deliberately follow-up work.

Split out of the kvproto re-vendor at pingyu's suggestion (#550).

Tests: 4 new — the resolver refusal (a plain lock passes; a wrapper carrying
members and a lock marked shared only by its op are both refused); the codec
converting a wrapper's own key alongside its members; a wrapper's unset fields
surviving truncation (the other reading, so both are pinned); and a lock on the
empty logical key round-tripping through truncate-then-encode. 71 lib tests
green.

The coverage is unit-level by necessity: CI pins TIKV_VERSION v8.5.5, which
predates shared locks, so no integration test in this repo can produce one, and
the refusal path is unreachable against that server. v8.5.6 is the first release
carrying shared locks, so a CI pin bump within the same release family would
make these paths reachable — left to a separate change.

Signed-off-by: Eduard R. <eduard@ralphovi.net>

Author: eduralph

src/request/keyspace.rs

@@ -160,23 +160,63 @@ 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() {
+                // Unlike `key`/`primary_lock` above, this guard is not an unset-field
+                // case: an unset `secondaries` is an empty Vec, with no elements to
+                // visit. A present-but-empty ELEMENT is malformed input (an encoded
+                // key always carries its prefix); skipping it merely keeps a corrupt
+                // entry from panicking the codec on `pretruncate_bytes`' length
+                // assertion.
+                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 +228,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 +254,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 +531,120 @@ 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());
+    }
+
+    /// An empty element INSIDE `secondaries` is a different case from the unset
+    /// fields above: an unset `secondaries` is an empty Vec with no elements, so a
+    /// present-but-empty element can only be malformed input. It is tolerated —
+    /// skipped rather than run into `pretruncate_bytes`' length assertion — while
+    /// the well-formed elements beside it still convert.
+    #[test]
+    fn a_malformed_empty_secondary_is_tolerated_not_truncated() {
+        use crate::proto::kvrpcpb::LockInfo;
+        let keyspace = Keyspace::Enable { keyspace_id: 0 };
+
+        let malformed = LockInfo {
+            key: vec![b'x', 0, 0, 0, b'k'],
+            secondaries: vec![vec![], vec![b'x', 0, 0, 0, b's']],
+            ..Default::default()
+        };
+        let out = vec![malformed].truncate_keyspace(keyspace);
+        assert_eq!(out[0].secondaries, vec![vec![], vec![b's']]);
+    }
 }

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;