From 58ee6625fa9918c6bbbbbf356d0cf5cf63ff4176 Mon Sep 17 00:00:00 2001
From: acheron <acheroncrypto@gmail.com>
Date: Thu, 14 Nov 2024 18:18:44 +0100
Subject: [PATCH] idl: Disallow account discriminators that can conflict with
 the `zero` constraint

---
 idl/src/build.rs | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/idl/src/build.rs b/idl/src/build.rs
index d6dd9c2477..c5dbe4a91a 100644
--- a/idl/src/build.rs
+++ b/idl/src/build.rs
@@ -404,5 +404,35 @@ fn verify(idl: &Idl) -> Result<()> {
         ));
     }
 
+    // Disallow account discriminators that can conflict with the `zero` constraint.
+    //
+    // Problematic scenario:
+    //
+    // 1. Account 1's discriminator starts with 0 (but not all 0s, since that's disallowed)
+    // 2. Account 2's discriminator is a 1-byte custom discriminator
+    // 3. Account 2 gets initialized using the `zero` constraint.
+    //
+    // In this case, it's possible to pass an already initialized Account 1 to a place that expects
+    // non-initialized Account 2, because the first byte of Account 1 is also 0, which is what the
+    // `zero` constraint checks.
+    for account in &idl.accounts {
+        let zero_count = account
+            .discriminator
+            .iter()
+            .take_while(|b| **b == 0)
+            .count();
+        if let Some(account2) = idl
+            .accounts
+            .iter()
+            .find(|acc| acc.discriminator.len() <= zero_count)
+        {
+            return Err(anyhow!(
+                "Accounts may allow substitution when used with the `zero` constraint: `{}` `{}`",
+                account.name,
+                account2.name
+            ));
+        }
+    }
+
     Ok(())
 }