From b0ac6d9c379f84d93f033057bbdabb7abd15079e Mon Sep 17 00:00:00 2001
From: Chao Yuan <yuanchao0310@163.com>
Date: Mon, 30 Sep 2019 11:43:07 +0800
Subject: [PATCH] Refactor: Use random port to gather udp candidate

This speed up the gather progress if many connections existed
---
 README.md      |  1 +
 gather.go      | 19 ++++++++++++++++---
 gather_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index c12cdb3..c4733ea 100644
--- a/README.md
+++ b/README.md
@@ -44,6 +44,7 @@ Check out the **[contributing wiki](https://github.com/pion/webrtc/wiki/Contribu
 * [Sebastian Waisbrot](https://github.com/seppo0010)
 * [Zizheng Tai](https://github.com/ZizhengTai)
 * [Aaron France](https://github.com/AeroNotix)
+* [Chao Yuan](https://github.com/yuanchao0310)
 
 ### License
 MIT License - see [LICENSE](LICENSE) for full text
diff --git a/gather.go b/gather.go
index 1dc847a..828738d 100644
--- a/gather.go
+++ b/gather.go
@@ -2,6 +2,7 @@ package ice
 
 import (
 	"fmt"
+	"math/rand"
 	"net"
 	"sync"
 	"time"
@@ -92,14 +93,26 @@ func (a *Agent) listenUDP(portMax, portMin int, network string, laddr *net.UDPAd
 	if j == 0 {
 		j = 0xFFFF
 	}
-	for i <= j {
-		laddr = &net.UDPAddr{IP: laddr.IP, Port: i}
+	if i > j {
+		return nil, ErrPort
+	}
+
+	portStart := rand.Intn(j-i+1) + i
+	portCurrent := portStart
+	for {
+		laddr = &net.UDPAddr{IP: laddr.IP, Port: portCurrent}
 		c, e := a.net.ListenUDP(network, laddr)
 		if e == nil {
 			return c, e
 		}
 		a.log.Debugf("failed to listen %s: %v", laddr.String(), e)
-		i++
+		portCurrent++
+		if portCurrent > j {
+			portCurrent = i
+		}
+		if portCurrent == portStart {
+			break
+		}
 	}
 	return nil, ErrPort
 }
diff --git a/gather_test.go b/gather_test.go
index 662a0da..ecbd030 100644
--- a/gather_test.go
+++ b/gather_test.go
@@ -2,6 +2,9 @@ package ice
 
 import (
 	"net"
+	"reflect"
+	"sort"
+	"strconv"
 	"testing"
 )
 
@@ -48,4 +51,43 @@ func TestListenUDP(t *testing.T) {
 	} else if port != "5000" {
 		t.Fatalf("listenUDP with port restriction of 5000 listened on incorrect port (%s)", port)
 	}
+
+	portMin := 5100
+	portMax := 5109
+	total := portMax - portMin + 1
+	result := make([]int, 0, total)
+	portRange := make([]int, 0, total)
+	for i := 0; i < total; i++ {
+		conn, err = a.listenUDP(portMax, portMin, udp, &net.UDPAddr{IP: ip, Port: 0})
+		if err != nil {
+			t.Fatalf("listenUDP error with no port restriction %v", err)
+		} else if conn == nil {
+			t.Fatalf("listenUDP error with no port restriction return a nil conn")
+		}
+
+		_, port, err = net.SplitHostPort(conn.LocalAddr().String())
+		if err != nil {
+			t.Fatal(err)
+		}
+		p, _ := strconv.Atoi(port)
+		if p < portMin || p > portMax {
+			t.Fatalf("listenUDP with port restriction [%d, %d] listened on incorrect port (%s)", portMin, portMax, port)
+		}
+		result = append(result, p)
+		portRange = append(portRange, portMin+i)
+	}
+	if sort.IntsAreSorted(result) {
+		t.Fatalf("listenUDP with port restriction [%d, %d], ports result should be random", portMin, portMax)
+	}
+	sort.Ints(result)
+	if !reflect.DeepEqual(result, portRange) {
+		t.Fatalf("listenUDP with port restriction [%d, %d], got:%v, want:%v", portMin, portMax, result, portRange)
+	}
+	_, err = a.listenUDP(portMax, portMin, udp, &net.UDPAddr{IP: ip, Port: 0})
+	if err == nil {
+		t.Fatalf("listenUDP with port restriction [%d, %d], should return error", portMin, portMax)
+	}
+	if err != ErrPort {
+		t.Fatalf("listenUDP with port restriction [%d, %d], did not return ErrPort", portMin, portMax)
+	}
 }