Fix get_custody_columns

specs/_features/eip7594/das-core.md

@@ -106,18 +106,19 @@ def get_custody_columns(node_id: NodeID, custody_subnet_count: uint64) -> Sequen
     assert custody_subnet_count <= DATA_COLUMN_SIDECAR_SUBNET_COUNT
 
     subnet_ids: List[uint64] = []
-    i = 0
+    current_id = uint256(node_id)
     while len(subnet_ids) < custody_subnet_count:
-        if node_id == UINT256_MAX:
-            node_id = NodeID(0)
-
         subnet_id = (
-            bytes_to_uint64(hash(uint_to_bytes(uint256(node_id + i)))[0:8])
+            bytes_to_uint64(hash(uint_to_bytes(uint256(current_id)))[0:8])
             % DATA_COLUMN_SIDECAR_SUBNET_COUNT
         )
         if subnet_id not in subnet_ids:
             subnet_ids.append(subnet_id)
-        i += 1
+        if current_id == UINT256_MAX:
+            # Overflow prevention
+            current_id = NodeID(0)
+        current_id += 1
+
     assert len(subnet_ids) == len(set(subnet_ids))
 
     columns_per_subnet = NUMBER_OF_COLUMNS // DATA_COLUMN_SIDECAR_SUBNET_COUNT

tests/core/pyspec/eth2spec/test/eip7594/networking/test_get_custody_columns.py

@@ -65,6 +65,25 @@ def test_get_custody_columns__max_node_id_max_custody_subnet_count(spec):
     )
 
 
+@with_eip7594_and_later
+@spec_test
+@single_phase
+def test_get_custody_columns__max_node_id_max_custody_subnet_count_minus_1(spec):
+    rng = random.Random(1111)
+    yield from _run_get_custody_columns(
+        spec, rng, node_id=2**256 - 2,
+        custody_subnet_count=spec.config.DATA_COLUMN_SIDECAR_SUBNET_COUNT,
+    )
+
+
+@with_eip7594_and_later
+@spec_test
+@single_phase
+def test_get_custody_columns__short_node_id(spec):
+    rng = random.Random(1111)
+    yield from _run_get_custody_columns(spec, rng, node_id=1048576, custody_subnet_count=1)
+
+
 @with_eip7594_and_later
 @spec_test
 @single_phase

tests/formats/networking/get_custody_columns.md

@@ -8,7 +8,7 @@
 
 ```yaml
 description: string            -- optional: description of test case, purely for debugging purposes.
-node_id: int                   -- argument: the NodeId input.
+node_id: int                   -- argument: the NodeID input.
 custody_subnet_count: int      -- argument: the count of custody subnets.
 result: list of int            -- output: the list of resulting column indices.
 ```