From 3c3731a34d874221519f0c375e6f1c5fb8424ce2 Mon Sep 17 00:00:00 2001
From: sukun <sukunrt@gmail.com>
Date: Mon, 20 May 2024 23:22:22 +0530
Subject: [PATCH] webrtc: add a test for establishing many connections

---
 p2p/transport/webrtc/listener.go       |  4 +-
 p2p/transport/webrtc/transport_test.go | 91 ++++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/p2p/transport/webrtc/listener.go b/p2p/transport/webrtc/listener.go
index 1834fc812b..b6711d66ed 100644
--- a/p2p/transport/webrtc/listener.go
+++ b/p2p/transport/webrtc/listener.go
@@ -329,8 +329,8 @@ func (l *listener) Multiaddr() ma.Multiaddr {
 func addOnConnectionStateChangeCallback(pc *webrtc.PeerConnection) <-chan error {
 	errC := make(chan error, 1)
 	var once sync.Once
-	pc.OnConnectionStateChange(func(state webrtc.PeerConnectionState) {
-		switch state {
+	pc.OnConnectionStateChange(func(_ webrtc.PeerConnectionState) {
+		switch pc.ConnectionState() {
 		case webrtc.PeerConnectionStateConnected:
 			once.Do(func() { close(errC) })
 		case webrtc.PeerConnectionStateFailed:
diff --git a/p2p/transport/webrtc/transport_test.go b/p2p/transport/webrtc/transport_test.go
index a3054a82df..3b329702f9 100644
--- a/p2p/transport/webrtc/transport_test.go
+++ b/p2p/transport/webrtc/transport_test.go
@@ -17,6 +17,7 @@ import (
 	"github.com/libp2p/go-libp2p/core/crypto"
 	"github.com/libp2p/go-libp2p/core/network"
 	"github.com/libp2p/go-libp2p/core/peer"
+	tpt "github.com/libp2p/go-libp2p/core/transport"
 	ma "github.com/multiformats/go-multiaddr"
 	manet "github.com/multiformats/go-multiaddr/net"
 	"github.com/multiformats/go-multibase"
@@ -860,3 +861,93 @@ func TestMaxInFlightRequests(t *testing.T) {
 	require.Equal(t, count, int(success.Load()), "expected exactly 3 dial successes")
 	require.Equal(t, 1, int(fails.Load()), "expected exactly 1 dial failure")
 }
+
+func TestManyConnections(t *testing.T) {
+	const N = 200
+	errCh := make(chan error, 200)
+	successCh := make(chan struct{}, 1)
+
+	tr, lp := getTransport(t)
+	ln, err := tr.Listen(ma.StringCast("/ip4/127.0.0.1/udp/0/webrtc-direct"))
+	require.NoError(t, err)
+	defer ln.Close()
+
+	runListenConn := func(conn tpt.CapableConn) {
+		s, err := conn.AcceptStream()
+		if err != nil {
+			t.Errorf("accept stream failed for listener: %s", err)
+			errCh <- err
+			return
+		}
+		var b [4]byte
+		if _, err := s.Read(b[:]); err != nil {
+			t.Errorf("read stream failed for listener: %s", err)
+			errCh <- err
+			return
+		}
+		s.Write(b[:])
+		s.Read(b[:]) // peer will close the stream after read
+		successCh <- struct{}{}
+		conn.Close()
+	}
+
+	runDialConn := func(conn tpt.CapableConn) {
+		s, err := conn.OpenStream(context.Background())
+		if err != nil {
+			t.Errorf("accept stream failed for listener: %s", err)
+			errCh <- err
+			return
+		}
+		var b [4]byte
+		if _, err := s.Write(b[:]); err != nil {
+			t.Errorf("write stream failed for dialer: %s", err)
+		}
+		if _, err := s.Read(b[:]); err != nil {
+			t.Errorf("read stream failed for dialer: %s", err)
+			errCh <- err
+			return
+		}
+		s.Close()
+		conn.Close()
+	}
+
+	go func() {
+		for i := 0; i < N; i++ {
+			conn, err := ln.Accept()
+			if err != nil {
+				t.Errorf("listener failed to accept conneciton: %s", err)
+				return
+			}
+			runListenConn(conn)
+		}
+	}()
+
+	go func() {
+		tp, _ := getTransport(t)
+		for i := 0; i < N; i++ {
+			// This test aims to check for deadlocks. So keep a high timeout
+			ctx, cancel := context.WithTimeout(context.Background(), 300*time.Second)
+			conn, err := tp.Dial(ctx, ln.Multiaddr(), lp)
+			if err != nil {
+				t.Errorf("dial failed: %s", err)
+				errCh <- err
+				cancel()
+				return
+			}
+			runDialConn(conn)
+			conn.Close()
+			cancel()
+		}
+	}()
+
+	for i := 0; i < N; i++ {
+		select {
+		case <-successCh:
+			t.Log(i)
+		case err := <-errCh:
+			t.Fatalf("failed: %s", err)
+		case <-time.After(100 * time.Second):
+			t.Fatalf("timed out")
+		}
+	}
+}