fix tests

Author: Gabriel-Trintinalia

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

@@ -42,18 +42,20 @@ public void handle(Envelope envelope) {
 
     NodeSession session = envelope.get(Field.SESSION);
 
-    // If already awaiting handshake completion, resend the original WHOAREYOU so the
-    // initiator can complete it using the same challenge nonce, rather than issuing a
-    // new challenge with the retransmitted packet's nonce.
+    OrdinaryMessagePacket unknownPacket = envelope.get(Field.UNAUTHORIZED_PACKET_MESSAGE);
+    Bytes12 msgNonce = unknownPacket.getHeader().getStaticHeader().getNonce();
+
+    // If already awaiting handshake completion, resend the original WHOAREYOU for retransmissions
+    // of the same packet (same nonce). For a new nonce, fall through and issue a fresh WHOAREYOU
+    // so the initiator's nonce check can pass.
     if (session.getState() == SessionState.WHOAREYOU_SENT) {
-      session.resendOutgoingWhoAreYou();
-      return;
+      if (session.getPendingWhoAreYouNonce().map(msgNonce::equals).orElse(false)) {
+        session.resendOutgoingWhoAreYou();
+        return;
+      }
     }
 
-    OrdinaryMessagePacket unknownPacket = envelope.get(Field.UNAUTHORIZED_PACKET_MESSAGE);
     try {
-      // packet it either random or message packet if session is expired
-      Bytes12 msgNonce = unknownPacket.getHeader().getStaticHeader().getNonce();
       Bytes16 idNonce = Bytes16.random(Functions.getRandom());
 
       Header<WhoAreYouAuthData> header =

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

@@ -170,6 +170,11 @@ public synchronized void sendOutgoingWhoAreYou(final WhoAreYouPacket packet) {
     dispatchWhoAreYou(maskingIV, packet);
   }
 
+  public synchronized Optional<Bytes12> getPendingWhoAreYouNonce() {
+    return pendingWhoAreYouPacket.map(
+        p -> p.getHeader().getStaticHeader().getNonce());
+  }
+
   public synchronized void resendOutgoingWhoAreYou() {
     pendingWhoAreYouPacket.ifPresent(
         packet -> {

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

@@ -35,8 +35,11 @@ void shouldResendExistingWhoAreYouWhenInWhoAreYouSentState() {
     final NodeSession session = mock(NodeSession.class);
     when(session.getState()).thenReturn(SessionState.WHOAREYOU_SENT);
 
-    final Envelope envelope = envelopeWith(session, createOrdinaryPacket());
-    handler.handle(envelope);
+    final OrdinaryMessagePacket packet = createOrdinaryPacket();
+    final Bytes12 nonce = packet.getHeader().getStaticHeader().getNonce();
+    when(session.getPendingWhoAreYouNonce()).thenReturn(Optional.of(nonce));
+
+    handler.handle(envelopeWith(session, packet));
 
     verify(session).resendOutgoingWhoAreYou();
     verify(session, never()).sendOutgoingWhoAreYou(any());
@@ -48,11 +51,34 @@ void shouldNotChangeStateWhenResendingInWhoAreYouSentState() {
     final NodeSession session = mock(NodeSession.class);
     when(session.getState()).thenReturn(SessionState.WHOAREYOU_SENT);
 
-    handler.handle(envelopeWith(session, createOrdinaryPacket()));
+    final OrdinaryMessagePacket packet = createOrdinaryPacket();
+    final Bytes12 nonce = packet.getHeader().getStaticHeader().getNonce();
+    when(session.getPendingWhoAreYouNonce()).thenReturn(Optional.of(nonce));
+
+    handler.handle(envelopeWith(session, packet));
 
     verify(session, never()).setState(any());
   }
 
+  @Test
+  void shouldSendNewWhoAreYouWhenInWhoAreYouSentStateButDifferentNonce() {
+    final NodeSession session = mock(NodeSession.class);
+    when(session.getState()).thenReturn(SessionState.WHOAREYOU_SENT);
+    when(session.getNodeRecord()).thenReturn(Optional.empty());
+    // Pending WhoAreYou was for a different nonce
+    when(session.getPendingWhoAreYouNonce())
+        .thenReturn(Optional.of(Bytes12.wrap(Bytes.random(12))));
+
+    final OrdinaryMessagePacket packet = createOrdinaryPacket();
+    handler.handle(envelopeWith(session, packet));
+
+    verify(session, never()).resendOutgoingWhoAreYou();
+    final ArgumentCaptor<WhoAreYouPacket> captor = ArgumentCaptor.forClass(WhoAreYouPacket.class);
+    verify(session).sendOutgoingWhoAreYou(captor.capture());
+    final Bytes12 expectedNonce = packet.getHeader().getStaticHeader().getNonce();
+    assertThat(captor.getValue().getHeader().getStaticHeader().getNonce()).isEqualTo(expectedNonce);
+  }
+
   @Test
   void shouldSendNewWhoAreYouWithIncomingNonceWhenInInitialState() {
     final NodeSession session = mock(NodeSession.class);