feat: limit collisions per code hash and family

requirements.txt

@@ -1,5 +1,6 @@
-psycopg[binary]
+psycopg[binary]==3.2.1
 tqdm==4.66.1
 typing_extensions==4.10.0
 requests==2.31.0
-web3
+web3==6.19.0
+cbor2==5.6.4

tests/integration/snapshots/snap_test_snapshots.py

@@ -14,15 +14,15 @@
     "types": ["balance"],
 }
 
-snapshots["test_tod_attack_miner_e2e num_candidates"] = 161
+snapshots["test_tod_attack_miner_e2e num_candidates"] = 157
 
 snapshots["test_tod_attack_miner_e2e stats"] = {
     "accesses": {"balance": 4663, "code": 2248, "nonce": 4332, "storage": 8239},
-    "candidates": 161,
+    "candidates": 157,
     "candidates_filters": {
         "candidates": {
             "before_filters": 2195,
-            "final": 161,
+            "final": 157,
             "original_without_same_value": 264234,
         },
         "filtered": {
@@ -32,42 +32,43 @@
             "indirect_dependencies_quick": 47,
             "indirect_dependencies_recursive": 38,
             "limited_collisions_per_address": 208,
+            "limited_collisions_per_code_family": 2,
+            "limited_collisions_per_code_hash": 2,
             "nonces": 658,
             "recipient_eth_transfer": 61,
             "same_sender": 174,
         },
     },
-    "candidates_transactions_unique": 238,
-    "collision_addresses_unique": 75,
-    "collisions": {"balance": 69, "storage": 209},
+    "candidates_transactions_unique": 235,
+    "collision_addresses_unique": 70,
+    "collisions": {"balance": 63, "storage": 172},
     "collisions_before_filters": {"balance": 1526, "nonce": 879, "storage": 710},
     "frequencies": {
-        "candidates_transactions": [(3, 1), (2, 82), (1, 155)],
+        "candidates_transactions": [(2, 79), (1, 156)],
         "collisions_addresses": [
-            (10, 13),
-            (9, 1),
+            (10, 11),
             (6, 2),
             (5, 5),
-            (4, 6),
-            (3, 5),
-            (2, 20),
-            (1, 23),
+            (4, 4),
+            (3, 3),
+            (2, 18),
+            (1, 27),
         ],
     },
     "samples": {
         "candidates_transactions_frequent": [
-            ("0x618f4b0392ac8ee3cb3c447c6dbe1ec4b472e8e49d302a393fd23e6d36164ec6", 3),
+            ("0x367b13927c05d804d67b2809daade8baba114e6d5a0f2d95cf7318ffdb656f97", 2),
             ("0x767e0cc6d695c96fb2b7be8a28f075c5a0729cdd646e3892230a6138a33925df", 2),
             ("0x4fa5690f54408827f28253342d0bf8c616b9377a9917ab7706be23c79079df9a", 2),
             ("0x7d98a6b0fa1b22bdd88bbaed71178105c121010dae17a6d96084e7576c055316", 2),
-            ("0x367b13927c05d804d67b2809daade8baba114e6d5a0f2d95cf7318ffdb656f97", 2),
+            ("0x5406c057c0fecf3d9ca3a200f9c8f58a10d57c2b732df59c5607b45412f40af6", 2),
         ],
         "collision_addresses_frequent": [
-            ("0x0d7e906bd9cafa154b048cfa766cc1e54e39af9b", 10),
             ("0x0ec68c5b10f21effb74f2a5c61dfe6b08c0db6cb", 10),
             ("0x3328f7f4a1d1c57c35df56bbf0c9dcafca309c49", 10),
-            ("0x6774bcbd5cecef1336b5300fb5186a12ddd8b367", 10),
             ("0x7777777f279eba3d3ad8f4e708545291a6fdba8b", 10),
+            ("0x8390a1da07e376ef7add4be859ba74fb83aa02d5", 10),
+            ("0x8fa3b4570b4c96f8036c13b64971ba65867eeb48", 10),
         ],
     },
     "state_diffs": {"balance": 2577, "code": 3, "nonce": 880, "storage": 2594},

tests/integration/test_snapshots.py

@@ -22,6 +22,7 @@ def test_tod_attack_miner_e2e(postgresql: Connection, snapshot: PyTestSnapshotTe
 
     miner.fetch(block_range.start, block_range.end)
     miner.find_collisions()
+    miner.compute_skelcodes()
     miner.filter_candidates(window_size=3)
 
     candidates = miner.get_candidates()

tod_attack_miner/cli.py

@@ -46,6 +46,7 @@ def main():
             print(json.dumps(miner.get_stats()))
         else:
             miner.fetch(int(args.from_block), int(args.to_block))
+            miner.compute_skelcodes()
             miner.find_collisions()
             miner.filter_candidates(args.window_size)
             print(f"Found {miner.count_candidates()} candidates")

tod_attack_miner/db/db.py

@@ -3,6 +3,7 @@
 
 import psycopg
 import psycopg.sql
+from tod_attack_miner.ethutils.skeleton import skeletize
 from tod_attack_miner.rpc.types import BlockWithTransactions, TxPrestate, TxStateDiff
 
 _TABLES = {
@@ -12,6 +13,8 @@
     "state_diffs": "(block_number INTEGER, tx_index INTEGER, tx_hash TEXT, type TEXT, key TEXT, pre_value TEXT, post_value TEXT)",
     "collisions": "(tx_write_hash TEXT, tx_access_hash TEXT, type TEXT, key TEXT, block_dist INTEGER, PRIMARY KEY(tx_write_hash, tx_access_hash, type, key))",
     "candidates": "(tx_write_hash TEXT, tx_access_hash TEXT, PRIMARY KEY(tx_write_hash, tx_access_hash))",
+    "codes": "(addr TEXT, code TEXT, hash TEXT, PRIMARY KEY(addr))",
+    "skeletons": "(addr TEXT, family TEXT, hash TEXT, PRIMARY KEY(addr))",
 }
 _INDEXES = {
     "accesses_type_key": "accesses(type, key, value)",
@@ -60,6 +63,8 @@ def _setup_tables(self):
     def insert_prestate(self, block_number: int, tx_index: int, prestate: TxPrestate):
         with self._con.cursor() as cursor:
             accesses: list[tuple[int, int, str, ACCESS_TYPE, str, str]] = []
+            codes: list[tuple[str, str, str]] = []
+
             for addr, state in prestate["result"].items():
                 if (balance := state.get("balance")) is not None:
                     accesses.append(
@@ -83,6 +88,7 @@ def insert_prestate(self, block_number: int, tx_index: int, prestate: TxPrestate
                             hash_code(code),
                         )
                     )
+                    codes.append((addr.lower(), code, hash_code(code)))
                 if (nonce := state.get("nonce")) is not None:
                     accesses.append(
                         (
@@ -110,6 +116,10 @@ def insert_prestate(self, block_number: int, tx_index: int, prestate: TxPrestate
                 "INSERT INTO accesses VALUES (%s, %s, %s, %s, %s, %s)",
                 accesses,
             )
+            cursor.executemany(
+                "INSERT INTO codes VALUES (%s, %s, %s) ON CONFLICT (addr) DO UPDATE SET code=EXCLUDED.code, hash=EXCLUDED.hash",
+                codes,
+            )
         self._con.commit()
 
     def insert_state_diff(
@@ -118,6 +128,7 @@ def insert_state_diff(
         pre, post = state_diff["result"]["pre"], state_diff["result"]["post"]
         with self._con.cursor() as cursor:
             state_diffs: list[tuple[int, int, str, ACCESS_TYPE, str, str, str]] = []
+            codes: list[tuple[str, str, str]] = []
 
             for addr, prestate in pre.items():
                 poststate = post[addr]
@@ -153,6 +164,8 @@ def insert_state_diff(
                             hash_code(post_code),
                         )
                     )
+                    # we use the post_code, since pre_code must always be empty
+                    codes.append((addr.lower(), post_code, hash_code(post_code)))
                 if (pre_nonce := prestate.get("nonce")) is not None:
                     assert "nonce" in poststate
                     post_nonce = poststate["nonce"]
@@ -189,8 +202,28 @@ def insert_state_diff(
                 "INSERT INTO state_diffs VALUES (%s, %s, %s, %s, %s, %s, %s)",
                 state_diffs,
             )
+
+            cursor.executemany(
+                "INSERT INTO codes VALUES (%s, %s, %s) ON CONFLICT (addr) DO UPDATE SET code=EXCLUDED.code, hash=EXCLUDED.hash",
+                codes,
+            )
         self._con.commit()
 
+    def insert_skeletons(self):
+        codes = self.get_codes()
+        mapped_codes = [(addr, code_skeleton_hash(c), hash) for addr, c, hash in codes]
+        self._insert_code_skeletons(mapped_codes)
+
+    def get_codes(self) -> Sequence[tuple[str, str, str]]:
+        with self._con.cursor() as cursor:
+            return cursor.execute("SELECT addr, code, hash FROM codes").fetchall()
+
+    def _insert_code_skeletons(self, mapped_codes: Iterable[tuple[str, str, str]]):
+        with self._con.cursor() as cursor:
+            cursor.executemany(
+                "INSERT INTO skeletons VALUES (%s, %s, %s)", mapped_codes
+            )
+
     def insert_collisions(self):
         self.insert_read_write_collisions()
         self.insert_write_write_collisions()
@@ -433,3 +466,9 @@ def count_unique_collision_addresses(self):
 def hash_code(code: str) -> str:
     """Some hash of the code, not the EVM hash"""
     return hashlib.sha256(code.encode("utf-8")).hexdigest()
+
+
+def code_skeleton_hash(code: str) -> str:
+    code_without_prefix = code[2:]
+    skelcode = skeletize(bytes.fromhex(code_without_prefix)).hex()
+    return hash_code(skelcode)

tod_attack_miner/db/filters.py

@@ -139,3 +139,43 @@ def limit_collisions_per_address(db: DB, limit=10):
         cursor.execute("SELECT setseed(0)")
         cursor.execute(sql)  # type: ignore
     return db.remove_candidates_without_collision()
+
+
+def limit_collisions_per_code_hash(db: DB, limit=10):
+    sql = f"""
+DELETE FROM collisions c
+USING (
+    SELECT tx_write_hash, tx_access_hash, type, key, ROW_NUMBER() OVER (PARTITION BY hash ORDER BY RANDOM()) AS n
+    FROM collisions
+    INNER JOIN skeletons ON SUBSTR(collisions.key, 1, 42) = skeletons.addr
+) grouped
+WHERE c.tx_write_hash = grouped.tx_write_hash
+  AND c.tx_access_hash = grouped.tx_access_hash
+  AND c.type = grouped.type
+  AND c.key = grouped.key
+  AND n > {limit}
+"""
+    with db._con.cursor() as cursor:
+        cursor.execute("SELECT setseed(0)")
+        cursor.execute(sql)  # type: ignore
+    return db.remove_candidates_without_collision()
+
+
+def limit_collisions_per_code_family(db: DB, limit=10):
+    sql = f"""
+DELETE FROM collisions c
+USING (
+    SELECT tx_write_hash, tx_access_hash, type, key, ROW_NUMBER() OVER (PARTITION BY family ORDER BY RANDOM()) AS n
+    FROM collisions
+    INNER JOIN skeletons ON SUBSTR(collisions.key, 1, 42) = skeletons.addr
+) grouped
+WHERE c.tx_write_hash = grouped.tx_write_hash
+  AND c.tx_access_hash = grouped.tx_access_hash
+  AND c.type = grouped.type
+  AND c.key = grouped.key
+  AND n > {limit}
+"""
+    with db._con.cursor() as cursor:
+        cursor.execute("SELECT setseed(0)")
+        cursor.execute(sql)  # type: ignore
+    return db.remove_candidates_without_collision()

tod_attack_miner/ethutils/README.txt

@@ -0,0 +1,3 @@
+This module was downloaded from: https://github.com/gsalzer/ethutils/releases/tag/emse2024
+
+Only the imports were modified to match the code structure, the code was left as is except for automatic formatting.

tod_attack_miner/ethutils/__init__.py

@@ -0,0 +1 @@
+__all__ = ["opcodes", "section", "skeleton"]