feat(spec-tools): support arbitrary headers in sync

Author: SamWilsn

src/ethereum_spec_tools/sync.py

@@ -13,7 +13,7 @@
 import time
 from queue import Empty, Full, Queue
 from threading import Thread
-from typing import Any, Dict, List, Optional, TypeVar, Union, cast
+from typing import Any, Dict, Final, List, Optional, TypeVar, Union, cast
 from urllib import request
 
 from ethereum_rlp import rlp
@@ -111,6 +111,7 @@ class BlockDownloader(ForkTracking):
     log: logging.Logger
     rpc_url: str
     geth: bool
+    headers: Final[dict[str, str]]
 
     def __init__(
         self,
@@ -120,6 +121,8 @@ def __init__(
         geth: bool,
         first_block: Uint,
         first_block_timestamp: U256,
+        *,
+        headers: dict[str, str] | None = None,
     ) -> None:
         ForkTracking.__init__(self, forks, first_block, first_block_timestamp)
 
@@ -136,6 +139,11 @@ def __init__(
         self.rpc_url = rpc_url
         self.geth = geth
 
+        if headers is None:
+            headers = {}
+
+        self.headers = headers
+
         Thread(target=self.download, name="download", daemon=True).start()
 
     def take_block(self) -> Optional[Any]:
@@ -221,14 +229,18 @@ def fetch_blocks_debug(
 
         self.log.debug("fetching blocks [%d, %d)...", first, first + count)
 
+        headers = {
+            "Content-Length": str(len(data)),
+            "Content-Type": "application/json",
+            "User-Agent": "ethereum-spec-sync",
+        }
+
+        headers.update(self.headers)
+
         post = request.Request(
             self.rpc_url,
             data=data,
-            headers={
-                "Content-Length": str(len(data)),
-                "Content-Type": "application/json",
-                "User-Agent": "ethereum-spec-sync",
-            },
+            headers=headers,
         )
 
         with request.urlopen(post) as response:
@@ -412,14 +424,18 @@ def fetch_blocks_eth(
 
         self.log.debug("fetching blocks [%d, %d)...", first, first + count)
 
+        headers = {
+            "Content-Length": str(len(data)),
+            "Content-Type": "application/json",
+            "User-Agent": "ethereum-spec-sync",
+        }
+
+        headers.update(self.headers)
+
         post = request.Request(
             self.rpc_url,
             data=data,
-            headers={
-                "Content-Length": str(len(data)),
-                "Content-Type": "application/json",
-                "User-Agent": "ethereum-spec-sync",
-            },
+            headers=headers,
         )
 
         with request.urlopen(post) as response:
@@ -489,14 +505,18 @@ def fetch_ommers(self, ommers_needed: Dict[Uint, int]) -> Dict[Uint, Any]:
             max(ommers_needed),
         )
 
+        headers = {
+            "Content-Length": str(len(data)),
+            "Content-Type": "application/json",
+            "User-Agent": "ethereum-spec-sync",
+        }
+
+        headers.update(self.headers)
+
         post = request.Request(
             self.rpc_url,
             data=data,
-            headers={
-                "Content-Length": str(len(data)),
-                "Content-Type": "application/json",
-                "User-Agent": "ethereum-spec-sync",
-            },
+            headers=headers,
         )
 
         with request.urlopen(post) as response:
@@ -663,18 +683,26 @@ def parse_arguments() -> argparse.Namespace:
         )
         parser.add_argument(
             "--zhejiang",
-            help="Set the chain to mainnet",
+            help="Set the chain to zhejiang",
             action="store_const",
             dest="chain",
             const="zhejiang",
         )
         parser.add_argument(
             "--sepolia",
-            help="Set the chain to mainnet",
+            help="Set the chain to sepolia",
             action="store_const",
             dest="chain",
             const="sepolia",
         )
+        parser.add_argument(
+            "--header",
+            "-H",
+            help="Add a header to RPC requests",
+            nargs=2,
+            action="append",
+            dest="headers",
+        )
 
         return parser.parse_args()
 
@@ -687,6 +715,10 @@ def __init__(self) -> None:
         self.log = logging.getLogger(__name__)
         self.options = self.parse_arguments()
 
+        headers = None
+        if self.options.headers is not None:
+            headers = dict(self.options.headers)
+
         if not self.options.unoptimized:
             import ethereum_optimized
 
@@ -805,6 +837,7 @@ def __init__(self) -> None:
                 self.options.geth,
                 Uint(0),
                 genesis_configuration.timestamp,
+                headers=headers,
             )
             self.set_block(Uint(0), genesis_configuration.timestamp)
         else:
@@ -820,6 +853,7 @@ def __init__(self) -> None:
                 self.options.geth,
                 persisted_block - initial_blocks_length,
                 persisted_block_timestamp,
+                headers=headers,
             )
             blocks = []
             for _ in range(initial_blocks_length):