From 0bba2e8b00e964058dc50ca1f96a1e4d9c9d2b64 Mon Sep 17 00:00:00 2001
From: Alin Tomescu <tomescu.alin@gmail.com>
Date: Wed, 28 Feb 2024 19:58:23 -0800
Subject: [PATCH] clearer feature gating for keyless TXNs

---
 aptos-move/aptos-vm/src/aptos_vm.rs           | 15 ++++++++++++++-
 aptos-move/aptos-vm/src/keyless_validation.rs | 15 +++++----------
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/aptos-move/aptos-vm/src/aptos_vm.rs b/aptos-move/aptos-vm/src/aptos_vm.rs
index a9a99fc531fc6c..a7a10c42d36932 100644
--- a/aptos-move/aptos-vm/src/aptos_vm.rs
+++ b/aptos-move/aptos-vm/src/aptos_vm.rs
@@ -1298,7 +1298,20 @@ impl AptosVM {
 
         let authenticators = aptos_types::keyless::get_authenticators(transaction)
             .map_err(|_| VMStatus::error(StatusCode::INVALID_SIGNATURE, None))?;
-        keyless_validation::validate_authenticators(&authenticators, self.features(), resolver)?;
+
+        // If there are keyless TXN authenticators, validate them all.
+        if !authenticators.is_empty() {
+            // Feature-gating keyless TXNs: if they are *not* enabled, return `FEATURE_UNDER_GATING`,
+            // which will discard the TXN from being put on-chain.
+            if !self.features().is_keyless_enabled() {
+                return Err(VMStatus::error(StatusCode::FEATURE_UNDER_GATING, None));
+            }
+            keyless_validation::validate_authenticators(
+                &authenticators,
+                self.features(),
+                resolver,
+            )?;
+        }
 
         // The prologue MUST be run AFTER any validation. Otherwise you may run prologue and hit
         // SEQUENCE_NUMBER_TOO_NEW if there is more than one transaction from the same sender and
diff --git a/aptos-move/aptos-vm/src/keyless_validation.rs b/aptos-move/aptos-vm/src/keyless_validation.rs
index ad82b287b0ed11..7f5e410590b94a 100644
--- a/aptos-move/aptos-vm/src/keyless_validation.rs
+++ b/aptos-move/aptos-vm/src/keyless_validation.rs
@@ -103,22 +103,17 @@ pub(crate) fn validate_authenticators(
     features: &Features,
     resolver: &impl AptosMoveResolver,
 ) -> Result<(), VMStatus> {
-    // Feature gating
     for (_, sig) in authenticators {
-        if !features.is_keyless_enabled() && matches!(sig.sig, ZkpOrOpenIdSig::Groth16Zkp { .. }) {
-            return Err(VMStatus::error(StatusCode::FEATURE_UNDER_GATING, None));
-        }
-        if (!features.is_keyless_enabled() || !features.is_keyless_zkless_enabled())
-            && matches!(sig.sig, ZkpOrOpenIdSig::OpenIdSig { .. })
+        // Feature-gating for keyless-but-zkless TXNs: If keyless TXNs *are* enabled, and (1) this
+        // is a ZKless transaction but (2) ZKless TXNs are not yet enabled, discard the TXN from
+        // being put on-chain.
+        if matches!(sig.sig, ZkpOrOpenIdSig::OpenIdSig { .. })
+            && !features.is_keyless_zkless_enabled()
         {
             return Err(VMStatus::error(StatusCode::FEATURE_UNDER_GATING, None));
         }
     }
 
-    if authenticators.is_empty() {
-        return Ok(());
-    }
-
     let config = &get_configs_onchain(resolver)?;
     if authenticators.len() > config.max_signatures_per_txn as usize {
         return Err(invalid_signature!("Too many keyless authenticators"));