feat: Infrastructure to intercept network messages in python (#2864)

This change introduces a large percentage of the `PeerMessage` schema ported to python, and a library that allows wrapping nodes of the cluster in such a way that the messages between them are intercepted by the test, de-borshified, and passed to a handler that can decide to drop them, deliver them, or change them on the fly.
near · Jun 18, 2020 · fa8e664 · fa8e664
1 parent e6b464e
commit fa8e664
Show file tree

Hide file tree

Showing 12 changed files with 768 additions and 20 deletions.
diff --git a/nightly/nightly.txt b/nightly/nightly.txt
@@ -37,6 +37,7 @@ pytest --timeout=300 sanity/gc_sync_after_sync.py swap_nodes
 pytest --timeout=300 sanity/large_messages.py
 pytest --timeout=300 sanity/upgradable.py
 pytest --timeout=240 sanity/validator_switch_key.py
+pytest sanity/nodes_proxy.py
 
 # python tests for smart contract deployment and invocation
 pytest contracts/deploy_call_smart_contract.py

diff --git a/pytest/lib/cluster.py b/pytest/lib/cluster.py
@@ -17,6 +17,7 @@
 from rc import gcloud
 import uuid
 import network
+from proxy import NodesProxy
 
 os.environ["ADVERSARY_CONSENT"] = "1"
 
@@ -72,7 +73,11 @@ def from_json_file(self, jf):
             return Key.from_json(json.loads(f.read()))
 
     def to_json(self):
-        return {'account_id': self.account_id, 'public_key': self.pk, 'secret_key': self.sk}
+        return {
+            'account_id': self.account_id,
+            'public_key': self.pk,
+            'secret_key': self.sk
+        }
 
 
 class BaseNode(object):
@@ -201,11 +206,14 @@ def validators(self):
                 self.get_status()['validators']))
 
     def stop_checking_refmap(self):
-        print("WARN: Stopping checking Reference Map for inconsistency for %s:%s" % self.addr())
+        print(
+            "WARN: Stopping checking Reference Map for inconsistency for %s:%s"
+            % self.addr())
         self.is_check_refmap = False
 
     def stop_checking_store(self):
-        print("WARN: Stopping checking Storage for inconsistency for %s:%s" % self.addr())
+        print("WARN: Stopping checking Storage for inconsistency for %s:%s" %
+              self.addr())
         self.is_check_store = False
 
     def check_refmap(self):
@@ -230,10 +238,13 @@ def check_store(self):
                 pass
             else:
                 if res['result'] == 0:
-                    print("ERROR: Storage for %s:%s in inconsistent state, stopping" % self.addr())
+                    print(
+                        "ERROR: Storage for %s:%s in inconsistent state, stopping"
+                        % self.addr())
                     self.kill()
                 self.store_tests += res['result']
 
+
 class RpcNode(BaseNode):
     """ A running node only interact by rpc queries """
 
@@ -489,7 +500,8 @@ def spin_up_node(config,
                  ordinal,
                  boot_key,
                  boot_addr,
-                 blacklist=[]):
+                 blacklist=[],
+                 proxy=None):
     is_local = config['local']
 
     print("Starting node %s %s" % (ordinal,
@@ -519,6 +531,9 @@ def spin_up_node(config,
             remote_nodes.append(node)
         print(f"node {ordinal} machine created")
 
+    if proxy is not None:
+        proxy.proxify_node(node)
+
     node.start(boot_key, boot_addr)
     time.sleep(3)
     print(f"node {ordinal} started")
@@ -622,8 +637,13 @@ def apply_config_changes(node_dir, client_config_change):
         f.write(json.dumps(config_json, indent=2))
 
 
-def start_cluster(num_nodes, num_observers, num_shards, config,
-                  genesis_config_changes, client_config_changes):
+def start_cluster(num_nodes,
+                  num_observers,
+                  num_shards,
+                  config,
+                  genesis_config_changes,
+                  client_config_changes,
+                  message_handler=None):
     if not config:
         config = load_config()
 
@@ -641,9 +661,11 @@ def start_cluster(num_nodes, num_observers, num_shards, config,
             filter(lambda n: not n.endswith('_finished'), node_dirs))
     ret = []
 
+    proxy = NodesProxy(message_handler) if message_handler is not None else None
+
     def spin_up_node_and_push(i, boot_key, boot_addr):
         node = spin_up_node(config, near_root, node_dirs[i], i, boot_key,
-                            boot_addr)
+                            boot_addr, [], proxy)
         while len(ret) < i:
             time.sleep(0.01)
         ret.append(node)

diff --git a/pytest/lib/messages/block.py b/pytest/lib/messages/block.py
@@ -0,0 +1,259 @@
+from messages.crypto import PublicKey, Signature, MerklePath, ShardProof
+from messages.tx import Receipt
+
+class Block:
+    pass
+
+
+class BlockV1:
+    pass
+
+
+class BlockHeader:
+    pass
+
+
+class BlockHeaderV1:
+    pass
+
+
+class BlockHeaderInnerLite:
+    pass
+
+
+class BlockHeaderInnerRest:
+    pass
+
+
+class ShardChunkHeader:
+    pass
+
+
+class ShardChunkHeaderInner:
+    pass
+
+
+class PartialEncodedChunkPart:
+    pass
+
+
+class ReceiptProof:
+    pass
+
+
+class PartialEncodedChunk:
+    pass
+
+
+class PartialEncodedChunkRequestMsg:
+    pass
+
+
+class PartialEncodedChunkResponseMsg:
+    pass
+
+
+class ValidatorStake:
+    pass
+
+
+class Approval:
+    pass
+
+
+class ApprovalInner:
+    pass
+
+
+block_schema = [
+    [
+        Block, {
+            'kind': 'enum',
+            'field': 'enum',
+            'values': [
+                ['BlockV1', BlockV1]
+            ]
+        }
+    ],
+    [
+        BlockV1, {
+            'kind': 'struct',
+            'fields': [
+                ['header', BlockHeader],
+                ['chunks', [ShardChunkHeader]],
+                ['challenges', [()]], # TODO
+
+                ['vrf_value', [32]],
+                ['vrf_proof', [64]],
+            ]
+        }
+    ],
+    [
+        BlockHeader, {
+            'kind': 'enum',
+            'field': 'enum',
+            'values': [
+                ['BlockHeaderV1', BlockHeaderV1]
+            ]
+        }
+    ],
+    [
+        BlockHeaderV1, {
+            'kind': 'struct',
+            'fields': [
+                ['prev_hash', [32]],
+                ['inner_lite', BlockHeaderInnerLite],
+                ['inner_rest', BlockHeaderInnerRest],
+                ['signature', Signature],
+            ]
+        }
+    ],
+    [
+        BlockHeaderInnerLite, {
+            'kind': 'struct',
+            'fields': [
+                ['height', 'u64'],
+                ['epoch_id', [32]],
+                ['next_epoch_id', [32]],
+                ['prev_state_root', [32]],
+                ['outcome_root', [32]],
+                ['timestamp', 'u64'],
+                ['next_bp_hash', [32]],
+                ['block_merkle_root', [32]],
+            ]
+        }
+    ],
+    [
+        BlockHeaderInnerRest, {
+            'kind': 'struct',
+            'fields': [
+                ['chunk_receipts_root', [32]],
+                ['chunk_headers_root', [32]],
+                ['chunk_tx_root', [32]],
+                ['chunks_included', 'u64'],
+                ['challenges_root', [32]],
+                ['random_value', [32]],
+                ['validator_proposals', [ValidatorStake]],
+                ['chunk_mask', ['u8']],
+                ['gas_price', 'u128'],
+                ['total_supply', 'u128'],
+                ['challenges_result', [()]], # TODO
+                ['last_final_block', [32]],
+                ['last_ds_final_block', [32]],
+                ['approvals', [{'kind': 'option', 'type': Signature}]],
+                ['latest_protocol_verstion', 'u32'],
+            ]
+        }
+    ],
+    [
+        ShardChunkHeader, {
+            'kind': 'struct',
+            'fields': [
+                ['inner', ShardChunkHeaderInner],
+                ['height_included', 'u64'],
+                ['signature', Signature],
+            ]
+        }
+    ],
+    [
+        ShardChunkHeaderInner, {
+            'kind': 'struct',
+            'fields': [
+                ['prev_block_hash', [32]],
+                ['prev_state_root', [32]],
+                ['outcome_root', [32]],
+                ['encoded_merkle_root', [32]],
+                ['encoded_length', 'u64'],
+                ['height_created', 'u64'],
+                ['shard_id', 'u64'],
+                ['gas_used', 'u64'],
+                ['gas_limit', 'u64'],
+                ['balance_burnt', 'u128'],
+                ['outgoing_receipt_root', [32]],
+                ['tx_root', [32]],
+                ['validator_proposals', [ValidatorStake]],
+            ]
+        }
+    ],
+    [
+        PartialEncodedChunkPart, {
+            'kind': 'struct',
+            'fields': [
+                ['part_ord', 'u64'],
+                ['part', ['u8']],
+                ['merkle_proof', MerklePath],
+            ]
+        }
+    ],
+    [
+        ReceiptProof, {
+            'kind': 'struct',
+            'fields': [
+                ['f1', [Receipt]],
+                ['f2', ShardProof],
+            ]
+        }
+    ],
+    [
+        PartialEncodedChunk, {
+            'kind': 'struct',
+            'fields': [
+                ['header', ShardChunkHeader],
+                ['parts', [PartialEncodedChunkPart]],
+                ['receipts', [ReceiptProof]]
+            ]
+        }
+    ],
+    [
+        PartialEncodedChunkRequestMsg, {
+            'kind': 'struct',
+            'fields': [
+                ['chunk_hash', [32]],
+                ['part_ords', ['u64']],
+                ['tracking_shards', ['u64']]
+            ]
+        }
+    ],
+    [
+        PartialEncodedChunkResponseMsg, {
+            'kind': 'struct',
+            'fields': [
+                ['chunk_hash', [32]],
+                ['parts', [PartialEncodedChunkPart]],
+                ['receipts', [ReceiptProof]]
+            ]
+        }
+    ],
+    [
+        ValidatorStake, {
+            'kind': 'struct',
+            'fields': [
+                ['account_id', 'string'],
+                ['public_key', PublicKey],
+                ['stake', 'u128'],
+            ]
+        }
+    ],
+    [
+        Approval, {
+            'kind': 'struct',
+            'fields': [
+                ['inner', ApprovalInner],
+                ['target_height', 'u64'],
+                ['signature', Signature],
+                ['account_id', 'string'],
+            ]
+        }
+    ],
+    [
+        ApprovalInner, {
+            'kind': 'enum',
+            'field': 'enum',
+            'values': [
+                ['Endorsement', [32]],
+                ['Skip', 'u64'],
+            ]
+        }
+    ],
+]
+