Implement SetRemoteCredentials and Restart

Add simple tests that check it works E2E

Resolves #190

agent.go

@@ -176,26 +176,6 @@ func NewAgent(config *AgentConfig) (*Agent, error) {
 		return nil, ErrPort
 	}
 
-	// local username fragment and password
-	localUfrag := randSeq(16)
-	localPwd := randSeq(32)
-
-	if config.LocalUfrag != "" {
-		if len([]rune(config.LocalUfrag))*8 < 24 {
-			return nil, ErrLocalUfragInsufficientBits
-		}
-
-		localUfrag = config.LocalUfrag
-	}
-
-	if config.LocalPwd != "" {
-		if len([]rune(config.LocalPwd))*8 < 128 {
-			return nil, ErrLocalPwdInsufficientBits
-		}
-
-		localPwd = config.LocalPwd
-	}
-
 	mDNSName := config.MulticastDNSHostName
 	if mDNSName == "" {
 		if mDNSName, err = generateMulticastDNSName(); err != nil {
@@ -234,29 +214,25 @@ func NewAgent(config *AgentConfig) (*Agent, error) {
 	}
 
 	a := &Agent{
-		tieBreaker:             rand.New(rand.NewSource(time.Now().UnixNano())).Uint64(),
-		lite:                   config.Lite,
-		gatheringState:         GatheringStateNew,
-		connectionState:        ConnectionStateNew,
-		localCandidates:        make(map[NetworkType][]Candidate),
-		remoteCandidates:       make(map[NetworkType][]Candidate),
-		pendingBindingRequests: make([]bindingRequest, 0),
-		checklist:              make([]*candidatePair, 0),
-		urls:                   config.Urls,
-		networkTypes:           config.NetworkTypes,
-		localUfrag:             localUfrag,
-		localPwd:               localPwd,
-		onConnected:            make(chan struct{}),
-		buffer:                 packetio.NewBuffer(),
-		done:                   make(chan struct{}),
-		chanState:              make(chan ConnectionState, 1),
-		portmin:                config.PortMin,
-		portmax:                config.PortMax,
-		trickle:                config.Trickle,
-		loggerFactory:          loggerFactory,
-		log:                    log,
-		net:                    config.Net,
-		muChan:                 make(chan struct{}, 1),
+		tieBreaker:       rand.New(rand.NewSource(time.Now().UnixNano())).Uint64(),
+		lite:             config.Lite,
+		gatheringState:   GatheringStateNew,
+		connectionState:  ConnectionStateNew,
+		localCandidates:  make(map[NetworkType][]Candidate),
+		remoteCandidates: make(map[NetworkType][]Candidate),
+		urls:             config.Urls,
+		networkTypes:     config.NetworkTypes,
+		onConnected:      make(chan struct{}),
+		buffer:           packetio.NewBuffer(),
+		done:             make(chan struct{}),
+		chanState:        make(chan ConnectionState, 1),
+		portmin:          config.PortMin,
+		portmax:          config.PortMax,
+		trickle:          config.Trickle,
+		loggerFactory:    loggerFactory,
+		log:              log,
+		net:              config.Net,
+		muChan:           make(chan struct{}, 1),
 
 		mDNSMode: mDNSMode,
 		mDNSName: mDNSName,
@@ -301,14 +277,11 @@ func NewAgent(config *AgentConfig) (*Agent, error) {
 		return nil, err
 	}
 
-	go func() {
-		for s := range a.chanState {
-			hdlr, ok := a.onConnectionStateChangeHdlr.Load().(func(ConnectionState))
-			if ok {
-				hdlr(s)
-			}
-		}
-	}()
+	// Restart is also used to initialize the agent for the first time
+	if err := a.Restart(config.LocalUfrag, config.LocalPwd); err != nil {
+		closeMDNSConn()
+		return nil, err
+	}
 
 	// Initialize local candidates
 	if !a.trickle {
@@ -345,17 +318,26 @@ func (a *Agent) onSelectedCandidatePairChange(p *candidatePair) {
 	}
 }
 
+func (a *Agent) startOnConnectionStateChangeRoutine() {
+	go func() {
+		for s := range a.chanState {
+			if hdlr, ok := a.onConnectionStateChangeHdlr.Load().(func(ConnectionState)); ok {
+				hdlr(s)
+			}
+		}
+	}()
+}
+
 func (a *Agent) startConnectivityChecks(isControlling bool, remoteUfrag, remotePwd string) error {
-	switch {
-	case a.haveStarted.Load():
+	if a.haveStarted.Load().(bool) {
 		return ErrMultipleStart
-	case remoteUfrag == "":
-		return ErrRemoteUfragEmpty
-	case remotePwd == "":
-		return ErrRemotePwdEmpty
+	}
+	if err := a.SetRemoteCredentials(remoteUfrag, remotePwd); err != nil {
+		return err
 	}
 
 	a.haveStarted.Store(true)
+	a.startOnConnectionStateChangeRoutine()
 	a.log.Debugf("Started agent: isControlling? %t, remoteUfrag: %q, remotePwd: %q", isControlling, remoteUfrag, remotePwd)
 
 	return a.run(func(agent *Agent) {
@@ -451,19 +433,16 @@ func (a *Agent) setSelectedPair(p *candidatePair) {
 	// Notify when the selected pair changes
 	a.onSelectedCandidatePairChange(p)
 
-	if p != nil {
-		p.nominated = true
-		a.selectedPair.Store(p)
-	} else {
+	if p == nil {
 		var nilPair *candidatePair
 		a.selectedPair.Store(nilPair)
+		return
 	}
 
-	a.updateConnectionState(ConnectionStateConnected)
+	p.nominated = true
+	a.selectedPair.Store(p)
 
-	// Close mDNS Conn. We don't need to do anymore querying
-	// and no reason to respond to others traffic
-	a.closeMulticastConn()
+	a.updateConnectionState(ConnectionStateConnected)
 
 	// Signal connected
 	a.onConnectedOnce.Do(func() { close(a.onConnected) })
@@ -969,6 +948,7 @@ func (a *Agent) validateNonSTUNTraffic(local Candidate, remote net.Addr) bool {
 
 func (a *Agent) getSelectedPair() *candidatePair {
 	selectedPair := a.selectedPair.Load()
+
 	if selectedPair == nil {
 		return nil
 	}
@@ -983,3 +963,72 @@ func (a *Agent) closeMulticastConn() {
 		}
 	}
 }
+
+// SetRemoteCredentials sets the credentials of the remote agent
+func (a *Agent) SetRemoteCredentials(remoteUfrag, remotePwd string) error {
+	switch {
+	case remoteUfrag == "":
+		return ErrRemoteUfragEmpty
+	case remotePwd == "":
+		return ErrRemotePwdEmpty
+	}
+
+	return a.run(func(agent *Agent) {
+		agent.remoteUfrag = remoteUfrag
+		agent.remotePwd = remotePwd
+	})
+}
+
+// Restart restarts the ICE Agent with the provided ufrag/pwd
+// If no ufrag/pwd is provided the Agent will generate one itself
+//
+// Restart must only be called when GatheringState is GatheringStateComplete
+// a user must then call GatherCandidates explicitly to start generating new ones
+func (a *Agent) Restart(ufrag, pwd string) error {
+	if ufrag == "" {
+		ufrag = randSeq(16)
+	}
+	if pwd == "" {
+		pwd = randSeq(32)
+	}
+
+	if len([]rune(ufrag))*8 < 24 {
+		return ErrLocalUfragInsufficientBits
+	}
+	if len([]rune(pwd))*8 < 128 {
+		return ErrLocalPwdInsufficientBits
+	}
+
+	err := make(chan error, 1)
+	if runErr := a.run(func(agent *Agent) {
+		if agent.gatheringState == GatheringStateGathering {
+			err <- ErrRestartWhenGathering
+			return
+		}
+
+		// Clear all agent needed to take back to fresh state
+		agent.localUfrag = ufrag
+		agent.localPwd = pwd
+		agent.remoteUfrag = ""
+		agent.remotePwd = ""
+		a.gatheringState = GatheringStateNew
+		a.checklist = make([]*candidatePair, 0)
+		a.pendingBindingRequests = make([]bindingRequest, 0)
+		a.setSelectedPair(nil)
+		a.deleteAllCandidates()
+		if a.selector != nil {
+			a.selector.Start()
+		}
+
+		// Restart is used by NewAgent. Accept/Connect should be used to move to checking
+		// for new Agents
+		if a.connectionState != ConnectionStateNew {
+			a.updateConnectionState(ConnectionStateChecking)
+		}
+
+		close(err)
+	}); runErr != nil {
+		return runErr
+	}
+	return <-err
+}

agent_test.go

@@ -5,6 +5,7 @@ package ice
 import (
 	"context"
 	"net"
+	"strconv"
 	"sync"
 	"testing"
 	"time"
@@ -148,6 +149,8 @@ func TestOnSelectedCandidatePairChange(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Failed to create agent: %s", err)
 	}
+	a.startOnConnectionStateChangeRoutine()
+
 	callbackCalled := make(chan struct{}, 1)
 	if err = a.OnSelectedCandidatePairChange(func(local, remote Candidate) {
 		close(callbackCalled)
@@ -1381,3 +1384,110 @@ func TestConnectionStateConnectingToFailed(t *testing.T) {
 	assert.NoError(t, aAgent.Close())
 	assert.NoError(t, bAgent.Close())
 }
+
+func TestAgentRestart(t *testing.T) {
+	lim := test.TimeOut(time.Second * 30)
+	defer lim.Stop()
+
+	report := test.CheckRoutines(t)
+	defer report()
+
+	t.Run("Restart During Gather", func(t *testing.T) {
+		agent, err := NewAgent(&AgentConfig{})
+		assert.NoError(t, err)
+
+		agent.gatheringState = GatheringStateGathering
+
+		assert.Equal(t, ErrRestartWhenGathering, agent.Restart("", ""))
+		assert.NoError(t, agent.Close())
+	})
+
+	t.Run("Restart When Closed", func(t *testing.T) {
+		agent, err := NewAgent(&AgentConfig{})
+		assert.NoError(t, err)
+		assert.NoError(t, agent.Close())
+
+		assert.Equal(t, ErrClosed, agent.Restart("", ""))
+	})
+
+	t.Run("Restart One Side", func(t *testing.T) {
+		oneSecond := time.Second
+		connA, connB := pipe(&AgentConfig{
+			DisconnectTimeout: &oneSecond,
+			FailedTimeout:     &oneSecond,
+		})
+
+		ctx, cancel := context.WithCancel(context.Background())
+		assert.NoError(t, connB.agent.OnConnectionStateChange(func(c ConnectionState) {
+			if c == ConnectionStateFailed || c == ConnectionStateDisconnected {
+				cancel()
+			}
+		}))
+		assert.NoError(t, connA.agent.Restart("", ""))
+
+		<-ctx.Done()
+		assert.NoError(t, connA.agent.Close())
+		assert.NoError(t, connB.agent.Close())
+	})
+
+	t.Run("Restart Both Sides", func(t *testing.T) {
+		// Get all addresses of candidates concatenated
+		generateCandidateAddressStrings := func(candidates []Candidate, err error) (out string) {
+			assert.NoError(t, err)
+
+			for _, c := range candidates {
+				out += c.Address() + ":"
+				out += strconv.Itoa(c.Port())
+			}
+			return
+		}
+
+		// Store the original candidates, confirm that after we reconnect we have new pairs
+		connA, connB := pipe(nil)
+		connAFirstCandidates := generateCandidateAddressStrings(connA.agent.GetLocalCandidates())
+		connBFirstCandidates := generateCandidateAddressStrings(connB.agent.GetLocalCandidates())
+
+		aNotifier, aConnected := onConnected()
+		assert.NoError(t, connA.agent.OnConnectionStateChange(aNotifier))
+
+		bNotifier, bConnected := onConnected()
+		assert.NoError(t, connB.agent.OnConnectionStateChange(bNotifier))
+
+		// Restart and Re-Signal
+		assert.NoError(t, connA.agent.Restart("", ""))
+		assert.NoError(t, connB.agent.Restart("", ""))
+
+		// Gather, and block until both sides are done
+		var gatherWaitGroup sync.WaitGroup
+		gatherWaitGroup.Add(2)
+		onCandidate := func(c Candidate) {
+			if c == nil {
+				gatherWaitGroup.Done()
+			}
+		}
+
+		assert.NoError(t, connA.agent.OnCandidate(onCandidate))
+		assert.NoError(t, connB.agent.OnCandidate(onCandidate))
+
+		assert.NoError(t, connA.agent.GatherCandidates())
+		assert.NoError(t, connB.agent.GatherCandidates())
+
+		gatherWaitGroup.Wait()
+
+		// Exchange Candidates and Credentials
+		assert.NoError(t, connA.agent.SetRemoteCredentials(connB.agent.GetLocalUserCredentials()))
+		assert.NoError(t, connB.agent.SetRemoteCredentials(connA.agent.GetLocalUserCredentials()))
+		signalAgents(connA.agent, connB.agent)
+
+		// Wait until both have gone back to connected
+		<-aConnected
+		<-bConnected
+
+		// Assert that we have new candiates each time
+		assert.NotEqual(t, connAFirstCandidates, generateCandidateAddressStrings(connA.agent.GetLocalCandidates()))
+		assert.NotEqual(t, connBFirstCandidates, generateCandidateAddressStrings(connB.agent.GetLocalCandidates()))
+
+		assert.NoError(t, connA.agent.Close())
+		assert.NoError(t, connB.agent.Close())
+	})
+}

errors.go

@@ -97,4 +97,7 @@ var (
 
 	// ErrInvalidMulticastDNSHostName indicates an invalid MulticastDNSHostName
 	ErrInvalidMulticastDNSHostName = errors.New("invalid mDNS HostName, must end with .local and can only contain a single '.'")
+
+	// ErrRestartWhenGathering indicates Restart was called when Agent is in GatheringStateGathering
+	ErrRestartWhenGathering = errors.New("ICE Agent can not be restarted when gathering")
 )

selection.go

@@ -26,6 +26,8 @@ type controllingSelector struct {
 
 func (s *controllingSelector) Start() {
 	s.startTime = time.Now()
+	s.nominatedPair = nil
+	s.nominationRequestCount = 0
 }
 
 func (s *controllingSelector) isNominatable(c Candidate) bool {

transport_test.go

@@ -64,7 +64,7 @@ func TestTimeout(t *testing.T) {
 		t.Skip("skipping test in short mode.")
 	}
 
-	ca, cb := pipe()
+	ca, cb := pipe(nil)
 	err := cb.Close()
 
 	if err != nil {
@@ -86,7 +86,7 @@ func TestTimeout(t *testing.T) {
 }
 
 func TestReadClosed(t *testing.T) {
-	ca, cb := pipe()
+	ca, cb := pipe(nil)
 
 	err := ca.Close()
 	if err != nil {
@@ -108,7 +108,7 @@ func TestReadClosed(t *testing.T) {
 }
 
 func stressDuplex(t *testing.T) {
-	ca, cb := pipe()
+	ca, cb := pipe(nil)
 
 	defer func() {
 		err := ca.Close()
@@ -133,7 +133,7 @@ func stressDuplex(t *testing.T) {
 }
 
 func Benchmark(b *testing.B) {
-	ca, cb := pipe()
+	ca, cb := pipe(nil)
 	defer func() {
 		err := ca.Close()
 		check(err)
@@ -193,7 +193,7 @@ func connect(aAgent, bAgent *Agent) (*Conn, *Conn) {
 	return aConn, bConn
 }
 
-func pipe() (*Conn, *Conn) {
+func pipe(defaultConfig *AgentConfig) (*Conn, *Conn) {
 	var urls []*URL
 
 	aNotifier, aConnected := onConnected()
@@ -202,12 +202,15 @@ func pipe() (*Conn, *Conn) {
 	var wg sync.WaitGroup
 	wg.Add(2)
 
-	cfg := &AgentConfig{
-		Urls:         urls,
-		Trickle:      true,
-		NetworkTypes: supportedNetworkTypes,
+	cfg := &AgentConfig{}
+	if defaultConfig != nil {
+		*cfg = *defaultConfig
 	}
 
+	cfg.Urls = urls
+	cfg.Trickle = true
+	cfg.NetworkTypes = supportedNetworkTypes
+
 	aAgent, err := NewAgent(cfg)
 	if err != nil {
 		panic(err)
@@ -404,7 +407,7 @@ func randomPort(t testing.TB) int {
 }
 
 func TestConnStats(t *testing.T) {
-	ca, cb := pipe()
+	ca, cb := pipe(nil)
 	if _, err := ca.Write(make([]byte, 10)); err != nil {
 		t.Fatal("unexpected error trying to write")
 	}