80
80
p2p_port ,
81
81
wait_until_helper_internal ,
82
82
)
83
+ from test_framework .v2_p2p import (
84
+ EncryptedP2PState ,
85
+ )
83
86
84
87
logger = logging .getLogger ("TestFramework.p2p" )
85
88
@@ -159,11 +162,16 @@ def __init__(self):
159
162
# The underlying transport of the connection.
160
163
# Should only call methods on this from the NetworkThread, c.f. call_soon_threadsafe
161
164
self ._transport = None
165
+ self .v2_state = None # EncryptedP2PState object needed for v2 p2p connections
162
166
163
167
@property
164
168
def is_connected (self ):
165
169
return self ._transport is not None
166
170
171
+ @property
172
+ def supports_v2_p2p (self ):
173
+ return self .v2_state is not None
174
+
167
175
def peer_connect_helper (self , dstaddr , dstport , net , timeout_factor ):
168
176
assert not self .is_connected
169
177
self .timeout_factor = timeout_factor
@@ -174,16 +182,20 @@ def peer_connect_helper(self, dstaddr, dstport, net, timeout_factor):
174
182
self .recvbuf = b""
175
183
self .magic_bytes = MAGIC_BYTES [net ]
176
184
177
- def peer_connect (self , dstaddr , dstport , * , net , timeout_factor ):
185
+ def peer_connect (self , dstaddr , dstport , * , net , timeout_factor , supports_v2_p2p ):
178
186
self .peer_connect_helper (dstaddr , dstport , net , timeout_factor )
187
+ if supports_v2_p2p :
188
+ self .v2_state = EncryptedP2PState (initiating = True , net = net )
179
189
180
190
loop = NetworkThread .network_event_loop
181
191
logger .debug ('Connecting to Bitcoin Node: %s:%d' % (self .dstaddr , self .dstport ))
182
192
coroutine = loop .create_connection (lambda : self , host = self .dstaddr , port = self .dstport )
183
193
return lambda : loop .call_soon_threadsafe (loop .create_task , coroutine )
184
194
185
- def peer_accept_connection (self , connect_id , connect_cb = lambda : None , * , net , timeout_factor ):
195
+ def peer_accept_connection (self , connect_id , connect_cb = lambda : None , * , net , timeout_factor , supports_v2_p2p ):
186
196
self .peer_connect_helper ('0' , 0 , net , timeout_factor )
197
+ if supports_v2_p2p :
198
+ self .v2_state = EncryptedP2PState (initiating = False , net = net )
187
199
188
200
logger .debug ('Listening for Bitcoin Node with id: {}' .format (connect_id ))
189
201
return lambda : NetworkThread .listen (self , connect_cb , idx = connect_id )
@@ -199,7 +211,13 @@ def connection_made(self, transport):
199
211
assert not self ._transport
200
212
logger .debug ("Connected & Listening: %s:%d" % (self .dstaddr , self .dstport ))
201
213
self ._transport = transport
202
- if self .on_connection_send_msg :
214
+ # in an inbound connection to the TestNode with P2PConnection as the initiator, [TestNode <---- P2PConnection]
215
+ # send the initial handshake immediately
216
+ if self .supports_v2_p2p and self .v2_state .initiating and not self .v2_state .tried_v2_handshake :
217
+ send_handshake_bytes = self .v2_state .initiate_v2_handshake ()
218
+ self .send_raw_message (send_handshake_bytes )
219
+ # if v2 connection, send `on_connection_send_msg` after initial v2 handshake.
220
+ if self .on_connection_send_msg and not self .supports_v2_p2p :
203
221
self .send_message (self .on_connection_send_msg )
204
222
self .on_connection_send_msg = None # Never used again
205
223
self .on_open ()
0 commit comments