From f5875c11130a74359706e87e199921b071b74d20 Mon Sep 17 00:00:00 2001
From: Daniel Mangum <georgedanielmangum@gmail.com>
Date: Sun, 27 Aug 2023 15:25:51 -0400
Subject: [PATCH] Set UDP routing if CID is enabled

Sets UDP routing functions in the default DTLS listener if a connection
ID generator is provided. Also updates to accept a dtls
net.PacketListener when a caller wishes to provide their own.

Signed-off-by: Daniel Mangum <georgedanielmangum@gmail.com>
---
 listener.go | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/listener.go b/listener.go
index 0d281fc4d..90dbbb427 100644
--- a/listener.go
+++ b/listener.go
@@ -6,10 +6,10 @@ package dtls
 import (
 	"net"
 
-	"github.com/pion/dtls/v2/internal/util"
+	"github.com/pion/dtls/v2/internal/net/udp"
+	dtlsnet "github.com/pion/dtls/v2/pkg/net"
 	"github.com/pion/dtls/v2/pkg/protocol"
 	"github.com/pion/dtls/v2/pkg/protocol/recordlayer"
-	"github.com/pion/transport/v2/udp"
 )
 
 // Listen creates a DTLS listener
@@ -31,6 +31,12 @@ func Listen(network string, laddr *net.UDPAddr, config *Config) (net.Listener, e
 			return h.ContentType == protocol.ContentTypeHandshake
 		},
 	}
+	// If connection ID support is enabled, then they must be supported in
+	// routing.
+	if config.ConnectionIDGenerator != nil {
+		lc.DatagramRouter = cidDatagramRouter(len(config.ConnectionIDGenerator()))
+		lc.ConnectionIdentifier = cidConnIdentifier()
+	}
 	parent, err := lc.Listen(network, laddr)
 	if err != nil {
 		return nil, err
@@ -42,7 +48,7 @@ func Listen(network string, laddr *net.UDPAddr, config *Config) (net.Listener, e
 }
 
 // NewListener creates a DTLS listener which accepts connections from an inner Listener.
-func NewListener(inner net.Listener, config *Config) (net.Listener, error) {
+func NewListener(inner dtlsnet.PacketListener, config *Config) (net.Listener, error) {
 	if err := validateConfig(config); err != nil {
 		return nil, err
 	}
@@ -56,7 +62,7 @@ func NewListener(inner net.Listener, config *Config) (net.Listener, error) {
 // listener represents a DTLS listener
 type listener struct {
 	config *Config
-	parent net.Listener
+	parent dtlsnet.PacketListener
 }
 
 // Accept waits for and returns the next connection to the listener.
@@ -64,11 +70,11 @@ type listener struct {
 // Connection handshake will timeout using ConnectContextMaker in the Config.
 // If you want to specify the timeout duration, set ConnectContextMaker.
 func (l *listener) Accept() (net.Conn, error) {
-	c, err := l.parent.Accept()
+	c, raddr, err := l.parent.Accept()
 	if err != nil {
 		return nil, err
 	}
-	return Server(util.FromConn(c), c.RemoteAddr(), l.config)
+	return Server(c, raddr, l.config)
 }
 
 // Close closes the listener.