Clean ufrag from UDP/TCP Mux during ICERestart

TCP/UDP Mux would have old entries of closed PacketConns from restarted
ICE Agents. Move cleanup code from Closing the ICE Agent into a shared
function and also call it during Restart.

agent.go

@@ -884,6 +884,13 @@ func (a *Agent) GetRemoteUserCredentials() (frag string, pwd string, err error)
 	return
 }
 
+func (a *Agent) removeUfragFromMux() {
+	a.tcpMux.RemoveConnByUfrag(a.localUfrag)
+	if a.udpMux != nil {
+		a.udpMux.RemoveConnByUfrag(a.localUfrag)
+	}
+}
+
 // Close cleans up the Agent
 func (a *Agent) Close() error {
 	if err := a.ok(); err != nil {
@@ -899,10 +906,7 @@ func (a *Agent) Close() error {
 
 	a.err.Store(ErrClosed)
 
-	a.tcpMux.RemoveConnByUfrag(a.localUfrag)
-	if a.udpMux != nil {
-		a.udpMux.RemoveConnByUfrag(a.localUfrag)
-	}
+	a.removeUfragFromMux()
 
 	close(a.done)
 
@@ -1216,6 +1220,7 @@ func (a *Agent) Restart(ufrag, pwd string) error {
 		}
 
 		// Clear all agent needed to take back to fresh state
+		a.removeUfragFromMux()
 		agent.localUfrag = ufrag
 		agent.localPwd = pwd
 		agent.remoteUfrag = ""

udp_mux_test.go

@@ -219,3 +219,41 @@ func verifyPacket(t *testing.T, b []byte, nextSeq uint32) {
 	h := sha1.Sum(b[24:]) //nolint:gosec
 	require.Equal(t, h[:], b[4:24])
 }
+
+func TestUDPMux_Agent_Restart(t *testing.T) {
+	oneSecond := time.Second
+	connA, connB := pipe(&AgentConfig{
+		DisconnectedTimeout: &oneSecond,
+		FailedTimeout:       &oneSecond,
+	})
+
+	aNotifier, aConnected := onConnected()
+	require.NoError(t, connA.agent.OnConnectionStateChange(aNotifier))
+
+	bNotifier, bConnected := onConnected()
+	require.NoError(t, connB.agent.OnConnectionStateChange(bNotifier))
+
+	// Maintain Credentials across restarts
+	ufragA, pwdA, err := connA.agent.GetLocalUserCredentials()
+	require.NoError(t, err)
+
+	ufragB, pwdB, err := connB.agent.GetLocalUserCredentials()
+	require.NoError(t, err)
+
+	require.NoError(t, err)
+
+	// Restart and Re-Signal
+	require.NoError(t, connA.agent.Restart(ufragA, pwdA))
+	require.NoError(t, connB.agent.Restart(ufragB, pwdB))
+
+	require.NoError(t, connA.agent.SetRemoteCredentials(ufragB, pwdB))
+	require.NoError(t, connB.agent.SetRemoteCredentials(ufragA, pwdA))
+	gatherAndExchangeCandidates(connA.agent, connB.agent)
+
+	// Wait until both have gone back to connected
+	<-aConnected
+	<-bConnected
+
+	require.NoError(t, connA.agent.Close())
+	require.NoError(t, connB.agent.Close())
+}