From f7d4014f5af1f2d2232084c1fb3ba0352ba7ade2 Mon Sep 17 00:00:00 2001
From: cnderrauber <zengjie9004@gmail.com>
Date: Mon, 10 Oct 2022 10:48:27 +0800
Subject: [PATCH] Revert mixup ipv4 and ipv6 connections

Revert mixup ipv4 and ipv6 connections for unspecified address
---
 udp_mux.go           | 63 +++++++++++++++++++++++++++++++-------------
 udp_mux_universal.go |  2 +-
 2 files changed, 45 insertions(+), 20 deletions(-)

diff --git a/udp_mux.go b/udp_mux.go
index ac00626..f31d4be 100644
--- a/udp_mux.go
+++ b/udp_mux.go
@@ -28,8 +28,8 @@ type UDPMuxDefault struct {
 	closedChan chan struct{}
 	closeOnce  sync.Once
 
-	// conns are maps of all udpMuxedConn indexed by ufrag|network|candidateType
-	conns map[string]*udpMuxedConn
+	// connsIPv4 and connsIPv6 are maps of all udpMuxedConn indexed by ufrag|network|candidateType
+	connsIPv4, connsIPv6 map[string]*udpMuxedConn
 
 	addressMapMu sync.RWMutex
 	addressMap   map[string]*udpMuxedConn
@@ -92,7 +92,8 @@ func NewUDPMuxDefault(params UDPMuxParams) *UDPMuxDefault {
 	m := &UDPMuxDefault{
 		addressMap: map[string]*udpMuxedConn{},
 		params:     params,
-		conns:      make(map[string]*udpMuxedConn),
+		connsIPv4:  make(map[string]*udpMuxedConn),
+		connsIPv6:  make(map[string]*udpMuxedConn),
 		closedChan: make(chan struct{}, 1),
 		pool: &sync.Pool{
 			New: func() interface{} {
@@ -122,13 +123,18 @@ func (m *UDPMuxDefault) GetListenAddresses() []net.Addr {
 	return []net.Addr{m.LocalAddr()}
 }
 
-// GetConn returns a PacketConn given the connection's ufrag and network
+// GetConn returns a PacketConn given the connection's ufrag and network address
 // creates the connection if an existing one can't be found
 func (m *UDPMuxDefault) GetConn(ufrag string, addr net.Addr) (net.PacketConn, error) {
 	// don't check addr for mux using unspecified address
 	if len(m.localAddrsForUnspecified) == 0 && m.params.UDPConn.LocalAddr().String() != addr.String() {
 		return nil, errInvalidAddress
 	}
+
+	var isIPv6 bool
+	if udpAddr, _ := addr.(*net.UDPAddr); udpAddr != nil && udpAddr.IP.To4() == nil {
+		isIPv6 = true
+	}
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
@@ -136,7 +142,7 @@ func (m *UDPMuxDefault) GetConn(ufrag string, addr net.Addr) (net.PacketConn, er
 		return nil, io.ErrClosedPipe
 	}
 
-	if conn, ok := m.getConn(ufrag); ok {
+	if conn, ok := m.getConn(ufrag, isIPv6); ok {
 		return conn, nil
 	}
 
@@ -146,24 +152,32 @@ func (m *UDPMuxDefault) GetConn(ufrag string, addr net.Addr) (net.PacketConn, er
 		m.RemoveConnByUfrag(ufrag)
 	}()
 
-	m.conns[ufrag] = c
+	if isIPv6 {
+		m.connsIPv6[ufrag] = c
+	} else {
+		m.connsIPv4[ufrag] = c
+	}
 
 	return c, nil
 }
 
 // RemoveConnByUfrag stops and removes the muxed packet connection
 func (m *UDPMuxDefault) RemoveConnByUfrag(ufrag string) {
-	var removedConn *udpMuxedConn
+	removedConns := make([]*udpMuxedConn, 0, 2)
 
 	// Keep lock section small to avoid deadlock with conn lock
 	m.mu.Lock()
-	if c, ok := m.conns[ufrag]; ok {
-		delete(m.conns, ufrag)
-		removedConn = c
+	if c, ok := m.connsIPv4[ufrag]; ok {
+		delete(m.connsIPv4, ufrag)
+		removedConns = append(removedConns, c)
+	}
+	if c, ok := m.connsIPv6[ufrag]; ok {
+		delete(m.connsIPv6, ufrag)
+		removedConns = append(removedConns, c)
 	}
 	m.mu.Unlock()
 
-	if removedConn == nil {
+	if len(removedConns) == 0 {
 		// No need to lock if no connection was found
 		return
 	}
@@ -171,9 +185,11 @@ func (m *UDPMuxDefault) RemoveConnByUfrag(ufrag string) {
 	m.addressMapMu.Lock()
 	defer m.addressMapMu.Unlock()
 
-	addresses := removedConn.getAddresses()
-	for _, addr := range addresses {
-		delete(m.addressMap, addr)
+	for _, c := range removedConns {
+		addresses := c.getAddresses()
+		for _, addr := range addresses {
+			delete(m.addressMap, addr)
+		}
 	}
 }
 
@@ -194,11 +210,15 @@ func (m *UDPMuxDefault) Close() error {
 		m.mu.Lock()
 		defer m.mu.Unlock()
 
-		for _, c := range m.conns {
+		for _, c := range m.connsIPv4 {
+			_ = c.Close()
+		}
+		for _, c := range m.connsIPv6 {
 			_ = c.Close()
 		}
 
-		m.conns = make(map[string]*udpMuxedConn)
+		m.connsIPv4 = make(map[string]*udpMuxedConn)
+		m.connsIPv6 = make(map[string]*udpMuxedConn)
 
 		close(m.closedChan)
 
@@ -290,9 +310,10 @@ func (m *UDPMuxDefault) connWorker() {
 			}
 
 			ufrag := strings.Split(string(attr), ":")[0]
+			isIPv6 := udpAddr.IP.To4() == nil
 
 			m.mu.Lock()
-			destinationConn, _ = m.getConn(ufrag)
+			destinationConn, _ = m.getConn(ufrag, isIPv6)
 			m.mu.Unlock()
 		}
 
@@ -307,8 +328,12 @@ func (m *UDPMuxDefault) connWorker() {
 	}
 }
 
-func (m *UDPMuxDefault) getConn(ufrag string) (val *udpMuxedConn, ok bool) {
-	val, ok = m.conns[ufrag]
+func (m *UDPMuxDefault) getConn(ufrag string, isIPv6 bool) (val *udpMuxedConn, ok bool) {
+	if isIPv6 {
+		val, ok = m.connsIPv6[ufrag]
+	} else {
+		val, ok = m.connsIPv4[ufrag]
+	}
 	return
 }
 
diff --git a/udp_mux_universal.go b/udp_mux_universal.go
index 53e45d9..f1e05ac 100644
--- a/udp_mux_universal.go
+++ b/udp_mux_universal.go
@@ -122,7 +122,7 @@ func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
 			return
 		}
 	}
-	return
+	return n, addr, err
 }
 
 // isXORMappedResponse indicates whether the message is a XORMappedAddress and is coming from the known STUN server.