Fix UDPMux memory leaks

Clients performing an early termination/cancelation
would cause `udpMuxedConn` to leak addresses
into `UDPMuxDefault`

gather.go

@@ -275,7 +275,7 @@ func (a *Agent) gatherCandidatesLocalUDPMux(ctx context.Context) error {
 		udpAddr, ok := conn.LocalAddr().(*net.UDPAddr)
 		if !ok {
 			closeConnAndLog(conn, a.log, fmt.Sprintf("Failed to create host mux candidate: %s failed to cast", candidateIP))
-			return nil
+			continue
 		}
 
 		hostConfig := CandidateHostConfig{
@@ -288,15 +288,16 @@ func (a *Agent) gatherCandidatesLocalUDPMux(ctx context.Context) error {
 		c, err := NewCandidateHost(&hostConfig)
 		if err != nil {
 			closeConnAndLog(conn, a.log, fmt.Sprintf("Failed to create host mux candidate: %s %d: %v", candidateIP, udpAddr.Port, err))
-			// already logged error
-			return nil
+			continue
 		}
 
 		if err := a.addCandidate(ctx, c, conn); err != nil {
 			if closeErr := c.close(); closeErr != nil {
 				a.log.Warnf("Failed to close candidate: %v", closeErr)
 			}
-			return err
+
+			closeConnAndLog(conn, a.log, fmt.Sprintf("Failed to add candidate: %s %d: %v", candidateIP, udpAddr.Port, err))
+			continue
 		}
 	}
 

udp_mux.go

@@ -93,7 +93,7 @@ func (m *UDPMuxDefault) GetConn(ufrag string, isIPv6 bool) (net.PacketConn, erro
 	c := m.createMuxedConn(ufrag)
 	go func() {
 		<-c.CloseChannel()
-		m.removeConn(ufrag)
+		m.RemoveConnByUfrag(ufrag)
 	}()
 
 	if isIPv6 {
@@ -121,6 +121,11 @@ func (m *UDPMuxDefault) RemoveConnByUfrag(ufrag string) {
 	}
 	m.mu.Unlock()
 
+	if len(removedConns) == 0 {
+		// No need to lock if no connection was found
+		return
+	}
+
 	m.addressMapMu.Lock()
 	defer m.addressMapMu.Unlock()
 
@@ -164,38 +169,6 @@ func (m *UDPMuxDefault) Close() error {
 	return err
 }
 
-func (m *UDPMuxDefault) removeConn(key string) {
-	// keep lock section small to avoid deadlock with conn lock
-	c := func() *udpMuxedConn {
-		m.mu.Lock()
-		defer m.mu.Unlock()
-
-		if c, ok := m.connsIPv4[key]; ok {
-			delete(m.connsIPv4, key)
-			return c
-		}
-
-		if c, ok := m.connsIPv6[key]; ok {
-			delete(m.connsIPv6, key)
-			return c
-		}
-
-		return nil
-	}()
-
-	if c == nil {
-		return
-	}
-
-	m.addressMapMu.Lock()
-	defer m.addressMapMu.Unlock()
-
-	addresses := c.getAddresses()
-	for _, addr := range addresses {
-		delete(m.addressMap, addr)
-	}
-}
-
 func (m *UDPMuxDefault) writeTo(buf []byte, raddr net.Addr) (n int, err error) {
 	return m.params.UDPConn.WriteTo(buf, raddr)
 }

udp_muxed_conn.go

@@ -112,9 +112,6 @@ func (c *udpMuxedConn) Close() error {
 		err = c.buffer.Close()
 		close(c.closedChan)
 	})
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	c.addresses = nil
 	return err
 }