From 888c0d2e1f925c3b163dd4b0f4884aa25c95231b Mon Sep 17 00:00:00 2001
From: LjhAUMEM <llnu14702@gmail.com>
Date: Wed, 4 Feb 2026 08:29:45 +0800
Subject: [PATCH] Finalmask UDP: Support WireGuard & Shadowsocks AEAD/2022
 (#5643)

https://github.com/XTLS/Xray-core/pull/5633#issuecomment-3833910076
---
 transport/internet/kcp/dialer.go | 16 ++++++++--------
 transport/internet/udp/dialer.go | 20 +++++++++++++++++++-
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/transport/internet/kcp/dialer.go b/transport/internet/kcp/dialer.go
index 577ef27b..a4915588 100644
--- a/transport/internet/kcp/dialer.go
+++ b/transport/internet/kcp/dialer.go
@@ -54,15 +54,15 @@ func DialKCP(ctx context.Context, dest net.Destination, streamSettings *internet
 		return nil, errors.New("failed to dial to dest: ", err).AtWarning().Base(err)
 	}
 
-	wrapper, ok := rawConn.(*internet.PacketConnWrapper)
-	if !ok {
-		rawConn.Close()
-		return nil, errors.New("raw is not PacketConnWrapper")
-	}
-
-	raw := wrapper.Conn
-
 	if streamSettings.UdpmaskManager != nil {
+		wrapper, ok := rawConn.(*internet.PacketConnWrapper)
+		if !ok {
+			rawConn.Close()
+			return nil, errors.New("raw is not PacketConnWrapper")
+		}
+
+		raw := wrapper.Conn
+
 		wrapper.Conn, err = streamSettings.UdpmaskManager.WrapPacketConnClient(raw)
 		if err != nil {
 			raw.Close()
diff --git a/transport/internet/udp/dialer.go b/transport/internet/udp/dialer.go
index 27f1505a..af25eb33 100644
--- a/transport/internet/udp/dialer.go
+++ b/transport/internet/udp/dialer.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/xtls/xray-core/common"
+	"github.com/xtls/xray-core/common/errors"
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/transport/internet"
 	"github.com/xtls/xray-core/transport/internet/stat"
@@ -20,7 +21,24 @@ func init() {
 			if err != nil {
 				return nil, err
 			}
+
+			if streamSettings != nil && streamSettings.UdpmaskManager != nil {
+				wrapper, ok := conn.(*internet.PacketConnWrapper)
+				if !ok {
+					conn.Close()
+					return nil, errors.New("conn is not PacketConnWrapper")
+				}
+
+				raw := wrapper.Conn
+
+				wrapper.Conn, err = streamSettings.UdpmaskManager.WrapPacketConnClient(raw)
+				if err != nil {
+					raw.Close()
+					return nil, errors.New("mask err").Base(err)
+				}
+			}
+
 			// TODO: handle dialer options
-			return stat.Connection(conn), nil
+			return conn, nil
 		}))
 }