Skip to content

Commit eac42de

Browse files
Gabriel-TrintinaliaGabriel-Trintinalia
committed
resend who are you
1 parent 114bd04 commit eac42de

4 files changed

Lines changed: 218 additions & 1 deletion

File tree

‎src/main/java/org/ethereum/beacon/discovery/pipeline/handler/UnauthorizedMessagePacketHandler.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,15 @@ public void handle(Envelope envelope) {
4141
envelope.getIdString()));
4242

4343
NodeSession session = envelope.get(Field.SESSION);
44+
45+
// If already awaiting handshake completion, resend the original WHOAREYOU so the
46+
// initiator can complete it using the same challenge nonce, rather than issuing a
47+
// new challenge with the retransmitted packet's nonce.
48+
if (session.getState() == SessionState.WHOAREYOU_SENT) {
49+
session.resendOutgoingWhoAreYou();
50+
return;
51+
}
52+
4453
OrdinaryMessagePacket unknownPacket = envelope.get(Field.UNAUTHORIZED_PACKET_MESSAGE);
4554
try {
4655
// packet it either random or message packet if session is expired

‎src/main/java/org/ethereum/beacon/discovery/schema/NodeSession.java‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ public class NodeSession {
6969
private final Signer signer;
7070
private Optional<InetSocketAddress> reportedExternalAddress = Optional.empty();
7171
private Optional<Bytes> whoAreYouChallenge = Optional.empty();
72+
private Optional<WhoAreYouPacket> pendingWhoAreYouPacket = Optional.empty();
7273
private Optional<Bytes12> lastOutboundNonce = Optional.empty();
7374
private boolean active = true;
7475
private final Function<Random, Bytes12> nonceGenerator;
@@ -161,10 +162,29 @@ public void sendOutgoingRandom(final Bytes randomData) {
161162
sendOutgoing(generateMaskingIV(), packet);
162163
}
163164

164-
public void sendOutgoingWhoAreYou(final WhoAreYouPacket packet) {
165+
public synchronized void sendOutgoingWhoAreYou(final WhoAreYouPacket packet) {
165166
LOG.trace(
166167
() -> String.format("Sending outgoing WhoAreYou message %s in session %s", packet, this));
167168
Bytes16 maskingIV = generateMaskingIV();
169+
pendingWhoAreYouPacket = Optional.of(packet);
170+
dispatchWhoAreYou(maskingIV, packet);
171+
}
172+
173+
public synchronized void resendOutgoingWhoAreYou() {
174+
pendingWhoAreYouPacket.ifPresent(
175+
packet -> {
176+
LOG.trace(
177+
() ->
178+
String.format(
179+
"Resending outgoing WhoAreYou message %s in session %s", packet, this));
180+
// Reuse the original maskingIV so the stored challenge remains stable; the initiator
181+
// may have already signed against it.
182+
Bytes16 maskingIV = Bytes16.wrap(whoAreYouChallenge.orElseThrow().slice(0, 16));
183+
sendOutgoing(maskingIV, packet);
184+
});
185+
}
186+
187+
private void dispatchWhoAreYou(final Bytes16 maskingIV, final WhoAreYouPacket packet) {
168188
whoAreYouChallenge = Optional.of(Bytes.wrap(maskingIV, packet.getHeader().getBytes()));
169189
sendOutgoing(maskingIV, packet);
170190
}
@@ -228,6 +248,7 @@ public synchronized RequestInfo createNextRequest(final Request<?> request) {
228248

229249
private synchronized void resetHandshakeState() {
230250
if (state == SessionState.WHOAREYOU_SENT || state == SessionState.RANDOM_PACKET_SENT) {
251+
pendingWhoAreYouPacket = Optional.empty();
231252
setState(SessionState.INITIAL);
232253
}
233254
}

‎src/test/java/org/ethereum/beacon/discovery/pipeline/handler/UnauthorizedMessagePacketHandlerTest.java‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*/
4+
5+
package org.ethereum.beacon.discovery.pipeline.handler;
6+
7+
import static org.assertj.core.api.Assertions.assertThat;
8+
import static org.mockito.ArgumentMatchers.any;
9+
import static org.mockito.Mockito.mock;
10+
import static org.mockito.Mockito.never;
11+
import static org.mockito.Mockito.verify;
12+
import static org.mockito.Mockito.when;
13+
14+
import java.util.Optional;
15+
import org.apache.tuweni.bytes.Bytes;
16+
import org.apache.tuweni.bytes.Bytes32;
17+
import org.ethereum.beacon.discovery.packet.Header;
18+
import org.ethereum.beacon.discovery.packet.OrdinaryMessagePacket;
19+
import org.ethereum.beacon.discovery.packet.OrdinaryMessagePacket.OrdinaryAuthData;
20+
import org.ethereum.beacon.discovery.packet.WhoAreYouPacket;
21+
import org.ethereum.beacon.discovery.pipeline.Envelope;
22+
import org.ethereum.beacon.discovery.pipeline.Field;
23+
import org.ethereum.beacon.discovery.schema.NodeSession;
24+
import org.ethereum.beacon.discovery.schema.NodeSession.SessionState;
25+
import org.ethereum.beacon.discovery.type.Bytes12;
26+
import org.junit.jupiter.api.Test;
27+
import org.mockito.ArgumentCaptor;
28+
29+
class UnauthorizedMessagePacketHandlerTest {
30+
31+
private final UnauthorizedMessagePacketHandler handler = new UnauthorizedMessagePacketHandler();
32+
33+
@Test
34+
void shouldResendExistingWhoAreYouWhenInWhoAreYouSentState() {
35+
final NodeSession session = mock(NodeSession.class);
36+
when(session.getState()).thenReturn(SessionState.WHOAREYOU_SENT);
37+
38+
final Envelope envelope = envelopeWith(session, createOrdinaryPacket());
39+
handler.handle(envelope);
40+
41+
verify(session).resendOutgoingWhoAreYou();
42+
verify(session, never()).sendOutgoingWhoAreYou(any());
43+
verify(session, never()).setState(any());
44+
}
45+
46+
@Test
47+
void shouldNotChangeStateWhenResendingInWhoAreYouSentState() {
48+
final NodeSession session = mock(NodeSession.class);
49+
when(session.getState()).thenReturn(SessionState.WHOAREYOU_SENT);
50+
51+
handler.handle(envelopeWith(session, createOrdinaryPacket()));
52+
53+
verify(session, never()).setState(any());
54+
}
55+
56+
@Test
57+
void shouldSendNewWhoAreYouWithIncomingNonceWhenInInitialState() {
58+
final NodeSession session = mock(NodeSession.class);
59+
when(session.getState()).thenReturn(SessionState.INITIAL);
60+
when(session.getNodeRecord()).thenReturn(Optional.empty());
61+
62+
final OrdinaryMessagePacket packet = createOrdinaryPacket();
63+
handler.handle(envelopeWith(session, packet));
64+
65+
final ArgumentCaptor<WhoAreYouPacket> captor = ArgumentCaptor.forClass(WhoAreYouPacket.class);
66+
verify(session).sendOutgoingWhoAreYou(captor.capture());
67+
verify(session, never()).resendOutgoingWhoAreYou();
68+
verify(session).setState(SessionState.WHOAREYOU_SENT);
69+
70+
// The WHOAREYOU nonce must echo the incoming packet's nonce so the initiator can
71+
// match it to their pending request.
72+
final Bytes12 expectedNonce = packet.getHeader().getStaticHeader().getNonce();
73+
assertThat(captor.getValue().getHeader().getStaticHeader().getNonce()).isEqualTo(expectedNonce);
74+
}
75+
76+
@Test
77+
void shouldSkipWhenUnauthorizedPacketMessageFieldAbsent() {
78+
final NodeSession session = mock(NodeSession.class);
79+
final Envelope envelope = new Envelope();
80+
envelope.put(Field.SESSION, session);
81+
82+
handler.handle(envelope);
83+
84+
verify(session, never()).resendOutgoingWhoAreYou();
85+
verify(session, never()).sendOutgoingWhoAreYou(any());
86+
}
87+
88+
@Test
89+
void shouldSkipWhenSessionFieldAbsent() {
90+
final Envelope envelope = new Envelope();
91+
envelope.put(Field.UNAUTHORIZED_PACKET_MESSAGE, createOrdinaryPacket());
92+
93+
// Should not throw even without a session.
94+
handler.handle(envelope);
95+
}
96+
97+
private static OrdinaryMessagePacket createOrdinaryPacket() {
98+
final Bytes12 nonce = Bytes12.wrap(Bytes.random(12));
99+
final Header<OrdinaryAuthData> header = Header.createOrdinaryHeader(Bytes32.ZERO, nonce);
100+
return OrdinaryMessagePacket.createRandom(header, Bytes.random(20));
101+
}
102+
103+
private static Envelope envelopeWith(
104+
final NodeSession session, final OrdinaryMessagePacket packet) {
105+
final Envelope envelope = new Envelope();
106+
envelope.put(Field.UNAUTHORIZED_PACKET_MESSAGE, packet);
107+
envelope.put(Field.SESSION, session);
108+
return envelope;
109+
}
110+
}

‎src/test/java/org/ethereum/beacon/discovery/schema/NodeSessionTest.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,13 @@
55
package org.ethereum.beacon.discovery.schema;
66

77
import static org.assertj.core.api.Assertions.assertThat;
8+
import static org.mockito.ArgumentMatchers.any;
89
import static org.mockito.ArgumentMatchers.eq;
910
import static org.mockito.Mockito.doReturn;
1011
import static org.mockito.Mockito.mock;
12+
import static org.mockito.Mockito.never;
1113
import static org.mockito.Mockito.spy;
14+
import static org.mockito.Mockito.times;
1215
import static org.mockito.Mockito.verify;
1316
import static org.mockito.Mockito.when;
1417

@@ -22,11 +25,16 @@
2225
import java.util.function.Consumer;
2326
import org.apache.tuweni.bytes.Bytes;
2427
import org.apache.tuweni.bytes.Bytes32;
28+
import org.apache.tuweni.units.bigints.UInt64;
2529
import org.ethereum.beacon.discovery.SimpleIdentitySchemaInterpreter;
2630
import org.ethereum.beacon.discovery.crypto.DefaultSigner;
2731
import org.ethereum.beacon.discovery.crypto.Signer;
2832
import org.ethereum.beacon.discovery.message.V5Message;
2933
import org.ethereum.beacon.discovery.network.NetworkParcel;
34+
import org.ethereum.beacon.discovery.network.NetworkParcelV5;
35+
import org.ethereum.beacon.discovery.packet.Header;
36+
import org.ethereum.beacon.discovery.packet.WhoAreYouPacket;
37+
import org.ethereum.beacon.discovery.packet.WhoAreYouPacket.WhoAreYouAuthData;
3038
import org.ethereum.beacon.discovery.pipeline.handler.NodeSessionManager;
3139
import org.ethereum.beacon.discovery.pipeline.info.Request;
3240
import org.ethereum.beacon.discovery.pipeline.info.RequestInfo;
@@ -36,6 +44,8 @@
3644
import org.ethereum.beacon.discovery.storage.LocalNodeRecordStore;
3745
import org.ethereum.beacon.discovery.storage.NewAddressHandler;
3846
import org.ethereum.beacon.discovery.storage.NodeRecordListener;
47+
import org.ethereum.beacon.discovery.type.Bytes12;
48+
import org.ethereum.beacon.discovery.type.Bytes16;
3949
import org.ethereum.beacon.discovery.util.Functions;
4050
import org.junit.jupiter.api.Test;
4151
import org.junit.jupiter.params.ParameterizedTest;
@@ -187,6 +197,73 @@ void createNextRequest_shouldNotResetAuthenticatedStatesWhenRequestTimesOut() {
187197
assertThat(session.getState()).isEqualTo(SessionState.AUTHENTICATED);
188198
}
189199

200+
@Test
201+
void resendOutgoingWhoAreYou_shouldSendPacketWhenPendingPacketExists() {
202+
session.sendOutgoingWhoAreYou(createWhoAreYouPacket(Bytes12.wrap(Bytes.random(12))));
203+
204+
final ArgumentCaptor<NetworkParcelV5> firstCaptor =
205+
ArgumentCaptor.forClass(NetworkParcelV5.class);
206+
verify(outgoingPipeline).accept(firstCaptor.capture());
207+
208+
session.resendOutgoingWhoAreYou();
209+
210+
final ArgumentCaptor<NetworkParcelV5> secondCaptor =
211+
ArgumentCaptor.forClass(NetworkParcelV5.class);
212+
verify(outgoingPipeline, times(2)).accept(secondCaptor.capture());
213+
// A packet must actually be sent on resend.
214+
assertThat(secondCaptor.getAllValues()).hasSize(2);
215+
}
216+
217+
@Test
218+
void resendOutgoingWhoAreYou_shouldDoNothingWhenNoPendingPacket() {
219+
session.resendOutgoingWhoAreYou();
220+
221+
verify(outgoingPipeline, never()).accept(any());
222+
}
223+
224+
@Test
225+
void resendOutgoingWhoAreYou_shouldDoNothingAfterHandshakeStateReset() {
226+
final Request<?> request = createRequestMock();
227+
final RequestInfo requestInfo = session.createNextRequest(request);
228+
229+
final ArgumentCaptor<Runnable> timeoutHandlerCaptor = ArgumentCaptor.forClass(Runnable.class);
230+
verify(expirationScheduler).put(eq(requestInfo.getRequestId()), timeoutHandlerCaptor.capture());
231+
232+
session.sendOutgoingWhoAreYou(createWhoAreYouPacket(Bytes12.wrap(Bytes.random(12))));
233+
session.setState(SessionState.WHOAREYOU_SENT);
234+
235+
// Simulate request timeout which resets the handshake state.
236+
timeoutHandlerCaptor.getValue().run();
237+
assertThat(session.getState()).isEqualTo(SessionState.INITIAL);
238+
239+
session.resendOutgoingWhoAreYou();
240+
241+
// sendOutgoingWhoAreYou was called once above; resend should not add another send.
242+
verify(outgoingPipeline, times(1)).accept(any());
243+
}
244+
245+
@Test
246+
void resendOutgoingWhoAreYou_shouldPreserveOriginalNonce() {
247+
final Bytes12 originalNonce = Bytes12.wrap(Bytes.random(12));
248+
final WhoAreYouPacket originalPacket = createWhoAreYouPacket(originalNonce);
249+
session.sendOutgoingWhoAreYou(originalPacket);
250+
251+
final Bytes challengeAfterSend = session.getWhoAreYouChallenge().orElseThrow();
252+
253+
session.resendOutgoingWhoAreYou();
254+
255+
// Challenge must be unchanged after resend so a handshake signed against the original
256+
// challenge remains valid.
257+
assertThat(session.getWhoAreYouChallenge()).contains(challengeAfterSend);
258+
}
259+
260+
private static WhoAreYouPacket createWhoAreYouPacket(final Bytes12 nonce) {
261+
final Bytes16 idNonce = Bytes16.wrap(Bytes.random(16));
262+
final Header<WhoAreYouAuthData> header =
263+
Header.createWhoAreYouHeader(nonce, idNonce, UInt64.ZERO);
264+
return WhoAreYouPacket.create(header);
265+
}
266+
190267
private Request<?> createRequestMock() {
191268
final Request<?> request = mock(Request.class);
192269
when(request.getResultPromise()).thenReturn(new CompletableFuture<>());

0 commit comments

Comments
 (0)