From d64ae4c32f1a78dc2549f507e351ef4a4c44ee69 Mon Sep 17 00:00:00 2001
From: Atsushi Watanabe <atsushi.w@ieee.org>
Date: Fri, 26 Jun 2020 14:14:21 +0900
Subject: [PATCH] Pause candidate recvLoop until init

Avoid using uninitialized object on receiving packet.
---
 agent.go          | 19 ++++++++++++-------
 candidate.go      |  2 +-
 candidate_base.go | 12 +++++++++---
 transport.go      |  3 ---
 4 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/agent.go b/agent.go
index ca2d190..0ce004c 100644
--- a/agent.go
+++ b/agent.go
@@ -35,9 +35,6 @@ type Agent struct {
 	onSelectedCandidatePairChangeHdlr atomic.Value // func(Candidate, Candidate)
 	onCandidateHdlr                   atomic.Value // func(Candidate)
 
-	// Used to block double Dial/Accept
-	opened bool
-
 	// State owned by the taskLoop
 	onConnected     chan struct{}
 	onConnectedOnce sync.Once
@@ -56,8 +53,9 @@ type Agent struct {
 	mDNSName string
 	mDNSConn *mdns.Conn
 
-	haveStarted   bool
 	muHaveStarted sync.Mutex
+	startedCh     <-chan struct{}
+	startedFn     func()
 	isControlling bool
 
 	maxBindingRequests uint16
@@ -217,6 +215,8 @@ func NewAgent(config *AgentConfig) (*Agent, error) {
 		}
 	}
 
+	startedCtx, startedFn := context.WithCancel(context.Background())
+
 	a := &Agent{
 		tieBreaker:       rand.New(rand.NewSource(time.Now().UnixNano())).Uint64(),
 		lite:             config.Lite,
@@ -229,6 +229,8 @@ func NewAgent(config *AgentConfig) (*Agent, error) {
 		onConnected:      make(chan struct{}),
 		buffer:           packetio.NewBuffer(),
 		done:             make(chan struct{}),
+		startedCh:        startedCtx.Done(),
+		startedFn:        startedFn,
 		chanState:        make(chan ConnectionState, 1),
 		portmin:          config.PortMin,
 		portmax:          config.PortMax,
@@ -329,13 +331,14 @@ func (a *Agent) startOnConnectionStateChangeRoutine() {
 func (a *Agent) startConnectivityChecks(isControlling bool, remoteUfrag, remotePwd string) error {
 	a.muHaveStarted.Lock()
 	defer a.muHaveStarted.Unlock()
-	if a.haveStarted {
+	select {
+	case <-a.startedCh:
 		return ErrMultipleStart
+	default:
 	}
 	if err := a.SetRemoteCredentials(remoteUfrag, remotePwd); err != nil {
 		return err
 	}
-	a.haveStarted = true
 
 	a.startOnConnectionStateChangeRoutine()
 	a.log.Debugf("Started agent: isControlling? %t, remoteUfrag: %q, remotePwd: %q", isControlling, remoteUfrag, remotePwd)
@@ -356,6 +359,7 @@ func (a *Agent) startConnectivityChecks(isControlling bool, remoteUfrag, remoteP
 		}
 
 		a.selector.Start()
+		a.startedFn()
 
 		agent.updateConnectionState(ConnectionStateChecking)
 
@@ -653,7 +657,7 @@ func (a *Agent) addRemoteCandidate(c Candidate) {
 
 func (a *Agent) addCandidate(c Candidate, candidateConn net.PacketConn) error {
 	return a.run(func(agent *Agent) {
-		c.start(a, candidateConn)
+		c.start(a, candidateConn, a.startedCh)
 
 		set := a.localCandidates[c.NetworkType()]
 		for _, candidate := range set {
@@ -715,6 +719,7 @@ func (a *Agent) Close() error {
 		close(agent.done)
 
 		a.deleteAllCandidates()
+		a.startedFn()
 
 		if err := a.buffer.Close(); err != nil {
 			a.log.Warnf("failed to close buffer: %v", err)
diff --git a/candidate.go b/candidate.go
index daeb8fa..58f00cb 100644
--- a/candidate.go
+++ b/candidate.go
@@ -37,6 +37,6 @@ type Candidate interface {
 
 	close() error
 	seen(outbound bool)
-	start(a *Agent, conn net.PacketConn)
+	start(a *Agent, conn net.PacketConn, initializedCh <-chan struct{})
 	writeTo(raw []byte, dst Candidate) (int, error)
 }
diff --git a/candidate_base.go b/candidate_base.go
index e216367..04077a2 100644
--- a/candidate_base.go
+++ b/candidate_base.go
@@ -72,20 +72,26 @@ func (c *candidateBase) RelatedAddress() *CandidateRelatedAddress {
 }
 
 // start runs the candidate using the provided connection
-func (c *candidateBase) start(a *Agent, conn net.PacketConn) {
+func (c *candidateBase) start(a *Agent, conn net.PacketConn, initializedCh <-chan struct{}) {
 	c.currAgent = a
 	c.conn = conn
 	c.closeCh = make(chan struct{})
 	c.closedCh = make(chan struct{})
 
-	go c.recvLoop()
+	go c.recvLoop(initializedCh)
 }
 
-func (c *candidateBase) recvLoop() {
+func (c *candidateBase) recvLoop(initializedCh <-chan struct{}) {
 	defer func() {
 		close(c.closedCh)
 	}()
 
+	select {
+	case <-initializedCh:
+	case <-c.closeCh:
+		return
+	}
+
 	log := c.agent().log
 	buffer := make([]byte, receiveMTU)
 	for {
diff --git a/transport.go b/transport.go
index e07ae60..ed728bf 100644
--- a/transport.go
+++ b/transport.go
@@ -45,9 +45,6 @@ func (a *Agent) connect(ctx context.Context, isControlling bool, remoteUfrag, re
 	if err != nil {
 		return nil, err
 	}
-	if a.opened {
-		return nil, errors.New("a connection is already opened")
-	}
 	err = a.startConnectivityChecks(isControlling, remoteUfrag, remotePwd)
 	if err != nil {
 		return nil, err