webrtc: correctly report incoming packet address on muxed connection

Author: sukunrt

p2p/transport/webrtc/udpmux/mux.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"strings"
 	"sync"
 
 	logging "github.com/ipfs/go-log/v2"
@@ -42,9 +43,15 @@ type UDPMux struct {
 
 	queue chan Candidate
 
-	mx       sync.Mutex
+	mx sync.Mutex
+	// ufragMap allows us to multiplex incoming STUN packets based on ufrag
 	ufragMap map[ufragConnKey]*muxedConnection
-	addrMap  map[string]*muxedConnection
+	// addrMap allows us to correctly direct incoming packets after the connection
+	// is established and ufrag isn't available on all packets
+	addrMap map[string]*muxedConnection
+	// ufragAddrMap allows us to clean up all addresses from the addrMap once
+	// connection is closed
+	ufragAddrMap map[ufragConnKey][]net.Addr
 
 	// the context controls the lifecycle of the mux
 	wg     sync.WaitGroup
@@ -57,12 +64,13 @@ var _ ice.UDPMux = &UDPMux{}
 func NewUDPMux(socket net.PacketConn) *UDPMux {
 	ctx, cancel := context.WithCancel(context.Background())
 	mux := &UDPMux{
-		ctx:      ctx,
-		cancel:   cancel,
-		socket:   socket,
-		ufragMap: make(map[ufragConnKey]*muxedConnection),
-		addrMap:  make(map[string]*muxedConnection),
-		queue:    make(chan Candidate, 32),
+		ctx:          ctx,
+		cancel:       cancel,
+		socket:       socket,
+		ufragMap:     make(map[ufragConnKey]*muxedConnection),
+		addrMap:      make(map[string]*muxedConnection),
+		ufragAddrMap: make(map[ufragConnKey][]net.Addr),
+		queue:        make(chan Candidate, 32),
 	}
 
 	return mux
@@ -130,7 +138,11 @@ func (mux *UDPMux) readLoop() {
 
 		n, addr, err := mux.socket.ReadFrom(buf)
 		if err != nil {
-			log.Errorf("error reading from socket: %v", err)
+			if strings.Contains(err.Error(), "use of closed network connection") {
+				log.Debugf("readLoop exiting: socket %s closed", mux.socket.LocalAddr())
+			} else {
+				log.Errorf("error reading from socket %s: %v", mux.socket.LocalAddr(), err)
+			}
 			pool.Put(buf)
 			return
 		}
@@ -157,7 +169,7 @@ func (mux *UDPMux) processPacket(buf []byte, addr net.Addr) (processed bool) {
 	conn, ok := mux.addrMap[addr.String()]
 	mux.mx.Unlock()
 	if ok {
-		if err := conn.Push(buf); err != nil {
+		if err := conn.Push(buf, addr); err != nil {
 			log.Debugf("could not push packet: %v", err)
 			return false
 		}
@@ -196,7 +208,7 @@ func (mux *UDPMux) processPacket(buf []byte, addr net.Addr) (processed bool) {
 		}
 	}
 
-	if err := conn.Push(buf); err != nil {
+	if err := conn.Push(buf, addr); err != nil {
 		log.Debugf("could not push packet: %v", err)
 		return false
 	}
@@ -250,9 +262,12 @@ func (mux *UDPMux) RemoveConnByUfrag(ufrag string) {
 
 	for _, isIPv6 := range [...]bool{true, false} {
 		key := ufragConnKey{ufrag: ufrag, isIPv6: isIPv6}
-		if conn, ok := mux.ufragMap[key]; ok {
+		if _, ok := mux.ufragMap[key]; ok {
 			delete(mux.ufragMap, key)
-			delete(mux.addrMap, conn.RemoteAddr().String())
+			for _, addr := range mux.ufragAddrMap[key] {
+				delete(mux.addrMap, addr.String())
+			}
+			delete(mux.ufragAddrMap, key)
 		}
 	}
 }
@@ -264,12 +279,14 @@ func (mux *UDPMux) getOrCreateConn(ufrag string, isIPv6 bool, _ *UDPMux, addr ne
 	defer mux.mx.Unlock()
 
 	if conn, ok := mux.ufragMap[key]; ok {
+		mux.addrMap[addr.String()] = conn
+		mux.ufragAddrMap[key] = append(mux.ufragAddrMap[key], addr)
 		return false, conn
 	}
 
-	conn := newMuxedConnection(mux, func() { mux.RemoveConnByUfrag(ufrag) }, addr)
+	conn := newMuxedConnection(mux, func() { mux.RemoveConnByUfrag(ufrag) })
 	mux.ufragMap[key] = conn
 	mux.addrMap[addr.String()] = conn
-
+	mux.ufragAddrMap[key] = append(mux.ufragAddrMap[key], addr)
 	return true, conn
 }

p2p/transport/webrtc/udpmux/muxed_connection.go

@@ -9,6 +9,11 @@ import (
 	pool "github.com/libp2p/go-buffer-pool"
 )
 
+type packet struct {
+	buf  []byte
+	addr net.Addr
+}
+
 var _ net.PacketConn = &muxedConnection{}
 
 const queueLen = 128
@@ -21,48 +26,46 @@ type muxedConnection struct {
 	ctx     context.Context
 	cancel  context.CancelFunc
 	onClose func()
-	queue   chan []byte
-	remote  net.Addr
+	queue   chan packet
 	mux     *UDPMux
 }
 
 var _ net.PacketConn = &muxedConnection{}
 
-func newMuxedConnection(mux *UDPMux, onClose func(), remote net.Addr) *muxedConnection {
+func newMuxedConnection(mux *UDPMux, onClose func()) *muxedConnection {
 	ctx, cancel := context.WithCancel(mux.ctx)
 	return &muxedConnection{
 		ctx:     ctx,
 		cancel:  cancel,
-		queue:   make(chan []byte, queueLen),
+		queue:   make(chan packet, queueLen),
 		onClose: onClose,
-		remote:  remote,
 		mux:     mux,
 	}
 }
 
-func (c *muxedConnection) Push(buf []byte) error {
+func (c *muxedConnection) Push(buf []byte, addr net.Addr) error {
 	select {
 	case <-c.ctx.Done():
 		return errors.New("closed")
 	default:
 	}
 	select {
-	case c.queue <- buf:
+	case c.queue <- packet{buf: buf, addr: addr}:
 		return nil
 	default:
 		return errors.New("queue full")
 	}
 }
 
-func (c *muxedConnection) ReadFrom(p []byte) (int, net.Addr, error) {
+func (c *muxedConnection) ReadFrom(buf []byte) (int, net.Addr, error) {
 	select {
-	case buf := <-c.queue:
-		n := copy(p, buf) // This might discard parts of the packet, if p is too short
-		if n < len(buf) {
-			log.Debugf("short read, had %d, read %d", len(buf), n)
+	case p := <-c.queue:
+		n := copy(buf, p.buf) // This might discard parts of the packet, if p is too short
+		if n < len(p.buf) {
+			log.Debugf("short read, had %d, read %d", len(p.buf), n)
 		}
-		pool.Put(buf)
-		return n, c.remote, nil
+		pool.Put(p.buf)
+		return n, p.addr, nil
 	case <-c.ctx.Done():
 		return 0, nil, c.ctx.Err()
 	}
@@ -83,15 +86,15 @@ func (c *muxedConnection) Close() error {
 	// drain the packet queue
 	for {
 		select {
-		case <-c.queue:
+		case p := <-c.queue:
+			pool.Put(p.buf)
 		default:
 			return nil
 		}
 	}
 }
 
-func (c *muxedConnection) LocalAddr() net.Addr  { return c.mux.socket.LocalAddr() }
-func (c *muxedConnection) RemoteAddr() net.Addr { return c.remote }
+func (c *muxedConnection) LocalAddr() net.Addr { return c.mux.socket.LocalAddr() }
 
 func (*muxedConnection) SetDeadline(t time.Time) error {
 	// no deadline is desired here