fix: resend original WHOAREYOU for any new ordinary packet during handshake

When a responder has already issued a WHOAREYOU challenge and receives a
second ordinary packet with a different nonce, it now resends the original
WHOAREYOU rather than issuing a fresh challenge. This matches the discv5
spec intent: hold one challenge per session attempt so the initiator can
complete the handshake it already started.

On the initiator side, WHOAREYOU validation is relaxed from 'must match the
last sent nonce' to 'must match any recently sent nonce' (bounded window of
8). The nonce-to-session lookup in NodeSessionManager is updated to keep all
recent nonces (not just the last), cleared together when the session is
deleted, reset on handshake timeout, and cleared on AUTHENTICATED transition.

Author: usmansaleem

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

@@ -128,8 +128,9 @@ private void deleteSession(final SessionKey sessionKey) {
     if (removedSession != null) {
       // Mark inactive to prevent registering any new nonces
       removedSession.markInactive();
-      // And then clean up the last recorded nonce, if any
-      removedSession.getLastOutboundNonce().ifPresent(lastNonceToSession::remove);
+      // Remove all nonces associated with this session (we keep multiple recent nonces per session
+      // so a single-nonce removal is no longer sufficient).
+      lastNonceToSession.values().removeIf(s -> s == removedSession); // identity check intentional
     }
   }
 
@@ -147,7 +148,8 @@ public Optional<NodeSession> getNodeSessionByLastOutboundNonce(final Bytes12 non
 
   public void onSessionLastNonceUpdate(
       final NodeSession session, final Optional<Bytes12> previousNonce, final Bytes12 newNonce) {
-    previousNonce.ifPresent(lastNonceToSession::remove);
+    // Keep the previous nonce in the map — a WHOAREYOU echoing it may still arrive if the remote
+    // resends an earlier challenge. All nonces for a session are removed together on session delete.
     lastNonceToSession.put(newNonce, session);
   }
 

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

@@ -46,14 +46,21 @@ protected void handlePacket(Envelope envelope) {
     Bytes12 msgNonce = unknownPacket.getHeader().getStaticHeader().getNonce();
 
     // If we have already issued a WHOAREYOU for this exact ordinary packet nonce, resend that
-    // same challenge — the initiator may have already signed against it. For any other nonce,
-    // fall through and issue a fresh WHOAREYOU; previous in-flight challenges remain queued so a
-    // delayed handshake signed against one of them can still be validated.
+    // same challenge — the initiator may have already signed against it.
     if (session.hasPendingWhoAreYouForNonce(msgNonce)) {
       session.resendOutgoingWhoAreYouFor(msgNonce);
       return;
     }
 
+    // If we're already awaiting handshake completion (a different nonce triggered the challenge),
+    // resend the original WHOAREYOU rather than issuing a new one. The spec requires the responder
+    // to hold the challenge open until the handshake times out; issuing a fresh nonce would race
+    // with the initiator signing the earlier one.
+    if (session.getState() == SessionState.WHOAREYOU_SENT) {
+      session.resendFirstPendingWhoAreYou();
+      return;
+    }
+
     try {
       Bytes16 idNonce = Bytes16.random(Functions.getRandom());
 

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

@@ -65,8 +65,9 @@ protected void handlePacket(final Envelope envelope) {
       final NodeRecord nodeRecord = session.getNodeRecord().orElseThrow();
 
       Bytes12 whoAreYouNonce = whoAreYouPacket.getHeader().getStaticHeader().getNonce();
-      boolean nonceMatches =
-          session.getLastOutboundNonce().map(whoAreYouNonce::equals).orElse(false);
+      // Accept WHOAREYOU for any recently sent nonce, not just the last one. The remote may echo
+      // back an earlier nonce if it resends the original challenge rather than issuing a new one.
+      boolean nonceMatches = session.hasRecentOutboundNonce(whoAreYouNonce);
       if (!nonceMatches) {
         LOG.trace(
             "Verification not passed for message [{}] from node {} in status {}",

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

@@ -14,6 +14,8 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Random;
@@ -63,6 +65,13 @@ public class NodeSession {
    */
   @VisibleForTesting static final int MAX_PENDING_WHOAREYOU = 5;
 
+  /**
+   * Maximum number of recent outbound nonces tracked per session on the initiator side. Allows
+   * WHOAREYOU validation to succeed even when the remote echoes back the nonce of an earlier
+   * ordinary packet rather than the most recent one.
+   */
+  @VisibleForTesting static final int MAX_RECENT_OUTBOUND_NONCES = 8;
+
   private final Bytes32 homeNodeId;
   private final LocalNodeRecordStore localNodeRecordStore;
   private final NodeSessionManager nodeSessionManager;
@@ -94,6 +103,10 @@ protected boolean removeEldestEntry(final Map.Entry<Bytes12, PendingWhoAreYou> e
             }
           });
 
+  // Bounded set of recent outbound nonces in insertion order. Used by the initiator to accept
+  // WHOAREYOU packets that echo an earlier nonce rather than the most recent one.
+  private final LinkedHashSet<Bytes12> recentOutboundNonces = new LinkedHashSet<>();
+
   private Optional<Bytes12> lastOutboundNonce = Optional.empty();
   private boolean active = true;
   private final Function<Random, Bytes12> nonceGenerator;
@@ -196,6 +209,17 @@ public boolean hasPendingWhoAreYouForNonce(final Bytes12 nonce) {
     return pendingWhoAreYou.containsKey(nonce);
   }
 
+  public void resendFirstPendingWhoAreYou() {
+    final Bytes12 firstNonce;
+    synchronized (pendingWhoAreYou) {
+      if (pendingWhoAreYou.isEmpty()) {
+        return;
+      }
+      firstNonce = pendingWhoAreYou.keySet().iterator().next();
+    }
+    resendOutgoingWhoAreYouFor(firstNonce);
+  }
+
   public void resendOutgoingWhoAreYouFor(final Bytes12 nonce) {
     final PendingWhoAreYou pending = pendingWhoAreYou.get(nonce);
     if (pending == null) {
@@ -287,6 +311,7 @@ public synchronized RequestInfo createNextRequest(final Request<?> request) {
   private synchronized void resetHandshakeState() {
     if (state == SessionState.WHOAREYOU_SENT || state == SessionState.RANDOM_PACKET_SENT) {
       pendingWhoAreYou.clear();
+      recentOutboundNonces.clear();
       setState(SessionState.INITIAL);
     }
   }
@@ -321,6 +346,10 @@ public synchronized Bytes12 generateNonce() {
     final Bytes12 newNonce = nonceGenerator.apply(rnd);
     final Optional<Bytes12> oldNonce = lastOutboundNonce;
     lastOutboundNonce = Optional.of(newNonce);
+    recentOutboundNonces.add(newNonce);
+    if (recentOutboundNonces.size() > MAX_RECENT_OUTBOUND_NONCES) {
+      recentOutboundNonces.remove(recentOutboundNonces.iterator().next());
+    }
     if (active) {
       // Update while synchronized to ensure only one update in flight at a time. Otherwise the
       // previous nonce may not have been recorded before we try to remove it leading to a memory
@@ -330,6 +359,14 @@ public synchronized Bytes12 generateNonce() {
     return newNonce;
   }
 
+  public synchronized boolean hasRecentOutboundNonce(final Bytes12 nonce) {
+    return recentOutboundNonces.contains(nonce);
+  }
+
+  public synchronized Collection<Bytes12> getRecentOutboundNonces() {
+    return List.copyOf(recentOutboundNonces);
+  }
+
   public synchronized Optional<Bytes12> getLastOutboundNonce() {
     return lastOutboundNonce;
   }
@@ -456,6 +493,9 @@ public synchronized void setState(final SessionState newStatus) {
     LOG.trace(
         () -> String.format("Switching status of node %s from %s to %s", nodeId, state, newStatus));
     this.state = newStatus;
+    if (newStatus == SessionState.AUTHENTICATED) {
+      recentOutboundNonces.clear();
+    }
   }
 
   public Signer getSigner() {

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

@@ -61,22 +61,22 @@ void shouldNotChangeStateWhenResendingForKnownNonce() {
   }
 
   @Test
-  void shouldSendNewWhoAreYouWhenIncomingNonceIsUnknown() {
+  void shouldResendFirstPendingWhoAreYouForDifferentNonceDuringHandshake() {
+    // When WHOAREYOU_SENT and the incoming nonce doesn't match any pending challenge, resend the
+    // original challenge so the initiator can complete the handshake it already started.
     final NodeSession session = mock(NodeSession.class);
     when(session.getState()).thenReturn(SessionState.WHOAREYOU_SENT);
-    when(session.getNodeRecord()).thenReturn(Optional.empty());
 
     final OrdinaryMessagePacket packet = createOrdinaryPacket();
     final Bytes12 incomingNonce = packet.getHeader().getStaticHeader().getNonce();
-    // No pending WhoAreYou for this nonce — an earlier challenge was for a different nonce.
     when(session.hasPendingWhoAreYouForNonce(incomingNonce)).thenReturn(false);
 
     handler.handle(envelopeWith(session, packet));
 
+    verify(session).resendFirstPendingWhoAreYou();
     verify(session, never()).resendOutgoingWhoAreYouFor(any());
-    final ArgumentCaptor<WhoAreYouPacket> captor = ArgumentCaptor.forClass(WhoAreYouPacket.class);
-    verify(session).sendOutgoingWhoAreYou(captor.capture());
-    assertThat(captor.getValue().getHeader().getStaticHeader().getNonce()).isEqualTo(incomingNonce);
+    verify(session, never()).sendOutgoingWhoAreYou(any());
+    verify(session, never()).setState(any());
   }
 
   @Test