diff --git a/agent.go b/agent.go index effe86d..7d449f3 100644 --- a/agent.go +++ b/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 = "" diff --git a/udp_mux_test.go b/udp_mux_test.go index d99abcc..3dd47f9 100644 --- a/udp_mux_test.go +++ b/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()) +}