diff --git a/src/ethereum_spec_tools/sync.py b/src/ethereum_spec_tools/sync.py index 4124aa522a..416de50f62 100644 --- a/src/ethereum_spec_tools/sync.py +++ b/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):