wip

2 files changed, 17 insertions, 10 deletions

diff --git a/lib/connections/quic_listen.go b/lib/connections/quic_listen.go
index 4146ebd1c0..eddc2dc5de 100644
--- a/lib/connections/quic_listen.go
+++ b/lib/connections/quic_listen.go
@@ -97,16 +97,21 @@ func (t *quicListener) serve(ctx context.Context) error {
 	}
 	defer func() { _ = udpConn.Close() }()
 
-	// XXX: HAX TO BE REVERTED
-	// svc, conn := stun.New(t.cfg, t, udpConn)
-	// defer conn.Close()
-
-	// go svc.Serve(ctx)
+	// Wrap the UDP connection in a STUN service if STUN is enabled.
+	var pktConn net.PacketConn
+	if t.cfg.Options().IsStunDisabled() {
+		pktConn = udpConn
+	} else {
+		svc, stunConn := stun.New(t.cfg, t, udpConn)
+		defer stunConn.Close()
+		go svc.Serve(ctx)
+		pktConn = stunConn
+	}
 
-	t.registry.Register(t.uri.Scheme, udpConn)
-	defer t.registry.Unregister(t.uri.Scheme, udpConn)
+	t.registry.Register(t.uri.Scheme, pktConn)
+	defer t.registry.Unregister(t.uri.Scheme, pktConn)
 
-	listener, err := quic.Listen(udpConn, t.tlsCfg, quicConfig)
+	listener, err := quic.Listen(pktConn, t.tlsCfg, quicConfig)
 	if err != nil {
 		l.Infoln("Listen (BEP/quic):", err)
 		return err
diff --git a/lib/stun/stun.go b/lib/stun/stun.go
index db96502cd7..8b7457358d 100644
--- a/lib/stun/stun.go
+++ b/lib/stun/stun.go
@@ -21,8 +21,10 @@ import (
 
 const stunRetryInterval = 5 * time.Minute
 
-type Host = stun.Host
-type NATType = stun.NATType
+type (
+	Host    = stun.Host
+	NATType = stun.NATType
+)
 
 // NAT types.