diff --git a/agent_udpmux_test.go b/agent_udpmux_test.go index b561d06..b2d1929 100644 --- a/agent_udpmux_test.go +++ b/agent_udpmux_test.go @@ -20,18 +20,22 @@ func TestMuxAgent(t *testing.T) { lim := test.TimeOut(time.Second * 30) defer lim.Stop() - loggerFactory := logging.NewDefaultLoggerFactory() - udpMux := NewUDPMuxDefault(UDPMuxParams{ - Logger: loggerFactory.NewLogger("ice"), - }) - muxPort := 7686 + const muxPort = 7686 + c, err := net.ListenUDP(udp, &net.UDPAddr{ Port: muxPort, }) + + loggerFactory := logging.NewDefaultLoggerFactory() + udpMux := NewUDPMuxDefault(UDPMuxParams{ + Logger: loggerFactory.NewLogger("ice"), + UDPConn: c, + }) + require.NoError(t, err) - require.NoError(t, udpMux.Start(c)) defer func() { _ = udpMux.Close() + _ = c.Close() }() muxedA, err := NewAgent(&AgentConfig{ diff --git a/gather.go b/gather.go index bd8acc9..3506a1f 100644 --- a/gather.go +++ b/gather.go @@ -178,7 +178,7 @@ func (a *Agent) gatherCandidatesLocal(ctx context.Context, networkTypes []Networ // accessible from the current interface. case udp: if a.udpMux != nil { - conn, err = a.udpMux.GetConn(a.localUfrag, network) + conn, err = a.udpMux.GetConn(a.localUfrag) if err != nil { a.log.Warnf("could not get udp muxed connection: %v\n", err) continue diff --git a/udp_mux.go b/udp_mux.go index f8688b0..69dd5d0 100644 --- a/udp_mux.go +++ b/udp_mux.go @@ -2,11 +2,9 @@ package ice import ( "errors" - "fmt" "io" "net" "os" - "strings" "sync" "github.com/pion/logging" @@ -15,15 +13,13 @@ import ( // UDPMux allows multiple connections to go over a single UDP port type UDPMux interface { io.Closer - GetConn(ufrag, network string) (net.PacketConn, error) + GetConn(ufrag string) (net.PacketConn, error) RemoveConnByUfrag(ufrag string) - Start(conn net.PacketConn) error } // UDPMuxDefault is an implementation of the interface type UDPMuxDefault struct { - params UDPMuxParams - udpConn net.PacketConn + params UDPMuxParams closedChan chan struct{} closeOnce sync.Once @@ -44,12 +40,13 @@ const maxAddrSize = 512 // UDPMuxParams are parameters for UDPMux. type UDPMuxParams struct { - Logger logging.LeveledLogger + Logger logging.LeveledLogger + UDPConn *net.UDPConn } // NewUDPMuxDefault creates an implementation of UDPMux func NewUDPMuxDefault(params UDPMuxParams) *UDPMuxDefault { - return &UDPMuxDefault{ + m := &UDPMuxDefault{ params: params, conns: make(map[string]*udpMuxedConn), closedChan: make(chan struct{}, 1), @@ -60,35 +57,20 @@ func NewUDPMuxDefault(params UDPMuxParams) *UDPMuxDefault { }, }, } -} - -// Start starts the mux. Before the UDPMux is usable, it must be started -// the mux will read/write data on the underlying net.PacketConn. It must -// conn.ReadFrom *MUST* return a *net.UDPAddr -func (m *UDPMuxDefault) Start(conn net.PacketConn) error { - if m.udpConn != nil { - return ErrMultipleStart - } - m.udpConn = conn go m.connWorker() - return nil + + return m } // LocalAddr returns the listening address of this UDPMuxDefault func (m *UDPMuxDefault) LocalAddr() net.Addr { - return m.udpConn.LocalAddr() + return m.params.UDPConn.LocalAddr() } // GetConn returns a PacketConn given the connection's ufrag and network // creates the connection if an existing one can't be found -func (m *UDPMuxDefault) GetConn(ufrag, network string) (net.PacketConn, error) { - if m.udpConn == nil { - return nil, ErrMuxNotStarted - } - - key := fmt.Sprintf("%s|%s", ufrag, network) - +func (m *UDPMuxDefault) GetConn(ufrag string) (net.PacketConn, error) { m.mu.Lock() defer m.mu.Unlock() @@ -96,16 +78,16 @@ func (m *UDPMuxDefault) GetConn(ufrag, network string) (net.PacketConn, error) { return nil, io.ErrClosedPipe } - if c, ok := m.conns[key]; ok { + if c, ok := m.conns[ufrag]; ok { return c, nil } - c := m.createMuxedConn(key) + c := m.createMuxedConn(ufrag) go func() { <-c.CloseChannel() - m.removeConn(key) + m.removeConn(ufrag) }() - m.conns[key] = c + m.conns[ufrag] = c return c, nil } @@ -114,9 +96,10 @@ func (m *UDPMuxDefault) RemoveConnByUfrag(ufrag string) { m.mu.Lock() removedConns := make([]*udpMuxedConn, 0) for key := range m.conns { - if !strings.HasPrefix(key, ufrag) { + if key != ufrag { continue } + c := m.conns[key] delete(m.conns, key) if c != nil { @@ -151,9 +134,6 @@ func (m *UDPMuxDefault) Close() error { m.mu.Lock() defer m.mu.Unlock() - // close udp conn and prevent packets coming in - err = m.udpConn.Close() - for _, c := range m.conns { _ = c.Close() } @@ -181,7 +161,7 @@ func (m *UDPMuxDefault) removeConn(key string) { } func (m *UDPMuxDefault) writeTo(buf []byte, raddr net.Addr) (n int, err error) { - return m.udpConn.WriteTo(buf, raddr) + return m.params.UDPConn.WriteTo(buf, raddr) } func (m *UDPMuxDefault) registerConnForAddress(conn *udpMuxedConn, addr string) { @@ -214,13 +194,16 @@ func (m *UDPMuxDefault) connWorker() { }() buf := make([]byte, receiveMTU) for { - n, addr, err := m.udpConn.ReadFrom(buf) - if err != nil { + n, addr, err := m.params.UDPConn.ReadFrom(buf) + if m.IsClosed() { + return + } else if err != nil { if errors.Is(err, os.ErrDeadlineExceeded) { continue } else if err != io.EOF { logger.Errorf("could not read udp packet: %v", err) } + return } diff --git a/udp_mux_test.go b/udp_mux_test.go index 192b85e..f523e48 100644 --- a/udp_mux_test.go +++ b/udp_mux_test.go @@ -25,18 +25,20 @@ func TestUDPMux(t *testing.T) { lim := test.TimeOut(time.Second * 30) defer lim.Stop() - loggerFactory := logging.NewDefaultLoggerFactory() - udpMux := NewUDPMuxDefault(UDPMuxParams{ - Logger: loggerFactory.NewLogger("ice"), - }) - conn, err := net.ListenUDP(udp, &net.UDPAddr{}) require.NoError(t, err) - err = udpMux.Start(conn) + + loggerFactory := logging.NewDefaultLoggerFactory() + udpMux := NewUDPMuxDefault(UDPMuxParams{ + Logger: loggerFactory.NewLogger("ice"), + UDPConn: conn, + }) + require.NoError(t, err) defer func() { _ = udpMux.Close() + _ = conn.Close() }() require.NotNil(t, udpMux.LocalAddr(), "tcpMux.LocalAddr() is nil") @@ -65,7 +67,7 @@ func TestUDPMux(t *testing.T) { require.NoError(t, udpMux.Close()) // can't create more connections - _, err = udpMux.GetConn("failufrag", udp) + _, err = udpMux.GetConn("failufrag") require.Error(t, err) } @@ -110,7 +112,7 @@ func TestAddressEncoding(t *testing.T) { } func testMuxConnection(t *testing.T, udpMux *UDPMuxDefault, ufrag string, network string) { - pktConn, err := udpMux.GetConn(ufrag, network) + pktConn, err := udpMux.GetConn(ufrag) require.NoError(t, err, "error retrieving muxed connection for ufrag") defer func() { _ = pktConn.Close()