Skip to content

Commit cbd2487

Browse files
committed
feat(spec-tools): support arbitrary headers in sync
1 parent 85c6fff commit cbd2487

File tree

1 file changed

+43
-13
lines changed

1 file changed

+43
-13
lines changed

src/ethereum_spec_tools/sync.py

Lines changed: 43 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
import time
1414
from queue import Empty, Full, Queue
1515
from threading import Thread
16-
from typing import Any, Dict, List, Optional, TypeVar, Union, cast
16+
from typing import Any, Dict, Final, List, Optional, TypeVar, Union, cast
1717
from urllib import request
1818

1919
from ethereum_rlp import rlp
@@ -111,6 +111,7 @@ class BlockDownloader(ForkTracking):
111111
log: logging.Logger
112112
rpc_url: str
113113
geth: bool
114+
headers: Final[dict[str, str]]
114115

115116
def __init__(
116117
self,
@@ -120,6 +121,8 @@ def __init__(
120121
geth: bool,
121122
first_block: Uint,
122123
first_block_timestamp: U256,
124+
*,
125+
headers: dict[str, str] | None = None,
123126
) -> None:
124127
ForkTracking.__init__(self, forks, first_block, first_block_timestamp)
125128

@@ -136,6 +139,11 @@ def __init__(
136139
self.rpc_url = rpc_url
137140
self.geth = geth
138141

142+
if headers is None:
143+
headers = {}
144+
145+
self.headers = headers
146+
139147
Thread(target=self.download, name="download", daemon=True).start()
140148

141149
def take_block(self) -> Optional[Any]:
@@ -221,14 +229,18 @@ def fetch_blocks_debug(
221229

222230
self.log.debug("fetching blocks [%d, %d)...", first, first + count)
223231

232+
headers = {
233+
"Content-Length": str(len(data)),
234+
"Content-Type": "application/json",
235+
"User-Agent": "ethereum-spec-sync",
236+
}
237+
238+
headers.update(self.headers)
239+
224240
post = request.Request(
225241
self.rpc_url,
226242
data=data,
227-
headers={
228-
"Content-Length": str(len(data)),
229-
"Content-Type": "application/json",
230-
"User-Agent": "ethereum-spec-sync",
231-
},
243+
headers=headers,
232244
)
233245

234246
with request.urlopen(post) as response:
@@ -412,14 +424,18 @@ def fetch_blocks_eth(
412424

413425
self.log.debug("fetching blocks [%d, %d)...", first, first + count)
414426

427+
headers = {
428+
"Content-Length": str(len(data)),
429+
"Content-Type": "application/json",
430+
"User-Agent": "ethereum-spec-sync",
431+
}
432+
433+
headers.update(self.headers)
434+
415435
post = request.Request(
416436
self.rpc_url,
417437
data=data,
418-
headers={
419-
"Content-Length": str(len(data)),
420-
"Content-Type": "application/json",
421-
"User-Agent": "ethereum-spec-sync",
422-
},
438+
headers=headers,
423439
)
424440

425441
with request.urlopen(post) as response:
@@ -663,18 +679,26 @@ def parse_arguments() -> argparse.Namespace:
663679
)
664680
parser.add_argument(
665681
"--zhejiang",
666-
help="Set the chain to mainnet",
682+
help="Set the chain to zhejiang",
667683
action="store_const",
668684
dest="chain",
669685
const="zhejiang",
670686
)
671687
parser.add_argument(
672688
"--sepolia",
673-
help="Set the chain to mainnet",
689+
help="Set the chain to sepolia",
674690
action="store_const",
675691
dest="chain",
676692
const="sepolia",
677693
)
694+
parser.add_argument(
695+
"--header",
696+
"-H",
697+
help="Add a header to RPC requests",
698+
nargs=2,
699+
action="append",
700+
dest="headers",
701+
)
678702

679703
return parser.parse_args()
680704

@@ -687,6 +711,10 @@ def __init__(self) -> None:
687711
self.log = logging.getLogger(__name__)
688712
self.options = self.parse_arguments()
689713

714+
headers = None
715+
if self.options.headers is not None:
716+
headers = dict(self.options.headers)
717+
690718
if not self.options.unoptimized:
691719
import ethereum_optimized
692720

@@ -805,6 +833,7 @@ def __init__(self) -> None:
805833
self.options.geth,
806834
Uint(0),
807835
genesis_configuration.timestamp,
836+
headers=headers,
808837
)
809838
self.set_block(Uint(0), genesis_configuration.timestamp)
810839
else:
@@ -820,6 +849,7 @@ def __init__(self) -> None:
820849
self.options.geth,
821850
persisted_block - initial_blocks_length,
822851
persisted_block_timestamp,
852+
headers=headers,
823853
)
824854
blocks = []
825855
for _ in range(initial_blocks_length):

0 commit comments

Comments
 (0)