listener_test.go: Fix failing test

Author: lactyy2

conn_negotiation_test.go

@@ -197,6 +197,32 @@ func TestConcurrentDialersShareSignaling(t *testing.T) {
 	}
 }
 
+func TestListenerRejectsDuplicateConn(t *testing.T) {
+	client, server := newMemorySignalingPair("1", "2")
+	defer client.close()
+	defer server.close()
+
+	l, err := (ListenConfig{AllowAnonymous: true}).Listen(server)
+	if err != nil {
+		t.Fatalf("Listen() error = %v", err)
+	}
+	defer l.Close()
+
+	ctx, cancel := context.WithTimeout(t.Context(), time.Second*10)
+	defer cancel()
+
+	firstConn, err := Dialer{ConnectionID: 20}.DialContext(ctx, server.NetworkID(), client)
+	if err != nil {
+		t.Fatalf("DialContext = %s, expected nil", err)
+	}
+	defer firstConn.Close()
+
+	_, err = Dialer{ConnectionID: 20}.DialContext(ctx, server.NetworkID(), client)
+	if err == nil {
+		t.Fatal("DialContext = nil, expected to return error")
+	}
+}
+
 func testDialListener(t *testing.T, disableTrickle bool) {
 	t.Helper()
 

listener_test.go

@@ -38,10 +38,10 @@ func TestListenerConnectionOwnership(t *testing.T) {
 	duplicate := &Conn{id: 7, networkID: "remote"}
 	key := first.remoteAddr().String()
 
-	if !l.registerConnection(first) {
+	if _, exists := l.connections.LoadOrStore(first.remoteAddr().String(), first); exists {
 		t.Fatal("registerConnection(first) = false, want true")
 	}
-	if l.registerConnection(duplicate) {
+	if _, exists := l.connections.LoadOrStore(duplicate.remoteAddr().String(), duplicate); !exists {
 		t.Fatal("registerConnection(duplicate) = true, want false")
 	}
 	if got, ok := l.connections.Load(key); !ok || got != first {