From 4cd6a1458e37c78ec172cf62b9168e3b92c65ca7 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Tue, 3 Aug 2021 11:21:39 +0200
Subject: [PATCH] cmd/devp2p: fix ping/pong race in discv4 tests (#23306)

This PR modifies the post-PING-send expectations to both be laxer and stricter: it doesn't care what order the packets arrive, but also verifies that exactly one PING and one PONG is returned.
---
 cmd/devp2p/internal/v4test/discv4tests.go | 57 ++++++++++++++++-------
 1 file changed, 41 insertions(+), 16 deletions(-)

diff --git a/cmd/devp2p/internal/v4test/discv4tests.go b/cmd/devp2p/internal/v4test/discv4tests.go
index 1b5e5304e..04ad67637 100644
--- a/cmd/devp2p/internal/v4test/discv4tests.go
+++ b/cmd/devp2p/internal/v4test/discv4tests.go
@@ -79,14 +79,44 @@ func BasicPing(t *utesting.T) {
 		To:         te.remoteEndpoint(),
 		Expiration: futureExpiration(),
 	})
-
-	reply, _, _ := te.read(te.l1)
-	if err := te.checkPong(reply, pingHash); err != nil {
+	if err := te.checkPingPong(pingHash); err != nil {
 		t.Fatal(err)
 	}
 }
 
-// checkPong verifies that reply is a valid PONG matching the given ping hash.
+// checkPingPong verifies that the remote side sends both a PONG with the
+// correct hash, and a PING.
+// The two packets do not have to be in any particular order.
+func (te *testenv) checkPingPong(pingHash []byte) error {
+	var (
+		pings int
+		pongs int
+	)
+	for i := 0; i < 2; i++ {
+		reply, _, err := te.read(te.l1)
+		if err != nil {
+			return err
+		}
+		switch reply.Kind() {
+		case v4wire.PongPacket:
+			if err := te.checkPong(reply, pingHash); err != nil {
+				return err
+			}
+			pongs++
+		case v4wire.PingPacket:
+			pings++
+		default:
+			return fmt.Errorf("expected PING or PONG, got %v %v", reply.Name(), reply)
+		}
+	}
+	if pongs == 1 && pings == 1 {
+		return nil
+	}
+	return fmt.Errorf("expected 1 PING  (got %d) and 1 PONG (got %d)", pings, pongs)
+}
+
+// checkPong verifies that reply is a valid PONG matching the given ping hash,
+// and a PING. The two packets do not have to be in any particular order.
 func (te *testenv) checkPong(reply v4wire.Packet, pingHash []byte) error {
 	if reply == nil {
 		return fmt.Errorf("expected PONG reply, got nil")
@@ -119,9 +149,7 @@ func PingWrongTo(t *utesting.T) {
 		To:         wrongEndpoint,
 		Expiration: futureExpiration(),
 	})
-
-	reply, _, _ := te.read(te.l1)
-	if err := te.checkPong(reply, pingHash); err != nil {
+	if err := te.checkPingPong(pingHash); err != nil {
 		t.Fatal(err)
 	}
 }
@@ -139,8 +167,7 @@ func PingWrongFrom(t *utesting.T) {
 		Expiration: futureExpiration(),
 	})
 
-	reply, _, _ := te.read(te.l1)
-	if err := te.checkPong(reply, pingHash); err != nil {
+	if err := te.checkPingPong(pingHash); err != nil {
 		t.Fatal(err)
 	}
 }
@@ -161,8 +188,7 @@ func PingExtraData(t *utesting.T) {
 		JunkData2:  []byte{9, 8, 7, 6, 5, 4, 3, 2, 1},
 	})
 
-	reply, _, _ := te.read(te.l1)
-	if err := te.checkPong(reply, pingHash); err != nil {
+	if err := te.checkPingPong(pingHash); err != nil {
 		t.Fatal(err)
 	}
 }
@@ -183,8 +209,7 @@ func PingExtraDataWrongFrom(t *utesting.T) {
 		JunkData2:  []byte{9, 8, 7, 6, 5, 4, 3, 2, 1},
 	}
 	pingHash := te.send(te.l1, &req)
-	reply, _, _ := te.read(te.l1)
-	if err := te.checkPong(reply, pingHash); err != nil {
+	if err := te.checkPingPong(pingHash); err != nil {
 		t.Fatal(err)
 	}
 }
@@ -240,9 +265,9 @@ func BondThenPingWithWrongFrom(t *utesting.T) {
 		To:         te.remoteEndpoint(),
 		Expiration: futureExpiration(),
 	})
-
-	reply, _, _ := te.read(te.l1)
-	if err := te.checkPong(reply, pingHash); err != nil {
+	if reply, _, err := te.read(te.l1); err != nil {
+		t.Fatal(err)
+	} else if err := te.checkPong(reply, pingHash); err != nil {
 		t.Fatal(err)
 	}
 }