From b33b0bc89d0308bc3287bec109b0941db9b75edc Mon Sep 17 00:00:00 2001
From: RPRX <63339210+RPRX@users.noreply.github.com>
Date: Wed, 12 Apr 2023 23:20:38 +0800
Subject: [PATCH] Allow multiple XUDP in Mux when using XTLS Vision (server
 side)

---
 common/mux/server.go           |  6 ++++++
 common/session/context.go      | 13 +++++++++++++
 proxy/vless/inbound/inbound.go |  4 +++-
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/common/mux/server.go b/common/mux/server.go
index 55e3a674..47a6d3dc 100644
--- a/common/mux/server.go
+++ b/common/mux/server.go
@@ -132,6 +132,12 @@ func (w *ServerWorker) handleStatusNew(ctx context.Context, meta *FrameMetadata,
 		ctx = log.ContextWithAccessMessage(ctx, msg)
 	}
 
+	if network := session.AllowedNetworkFromContext(ctx); network != net.Network_Unknown {
+		if meta.Target.Network != network {
+			return newError("unexpected network ", meta.Target.Network) // it will break the whole Mux connection
+		}
+	}
+
 	if meta.GlobalID != [8]byte{} {
 		mb, err := NewPacketReader(reader, &meta.Target).ReadMultiBuffer()
 		if err != nil {
diff --git a/common/session/context.go b/common/session/context.go
index 71e4b154..329a5a65 100644
--- a/common/session/context.go
+++ b/common/session/context.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	_ "unsafe"
 
+	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/features/routing"
 )
 
@@ -22,6 +23,7 @@ const (
 	trackedConnectionErrorKey
 	dispatcherKey
 	timeoutOnlyKey
+	allowedNetworkKey
 )
 
 // ContextWithID returns a new context with the given ID.
@@ -147,3 +149,14 @@ func TimeoutOnlyFromContext(ctx context.Context) bool {
 	}
 	return false
 }
+
+func ContextWithAllowedNetwork(ctx context.Context, network net.Network) context.Context {
+	return context.WithValue(ctx, allowedNetworkKey, network)
+}
+
+func AllowedNetworkFromContext(ctx context.Context) net.Network {
+	if val, ok := ctx.Value(allowedNetworkKey).(net.Network); ok {
+		return val
+	}
+	return net.Network_Unknown
+}
diff --git a/proxy/vless/inbound/inbound.go b/proxy/vless/inbound/inbound.go
index c8a69444..8d9b9b43 100644
--- a/proxy/vless/inbound/inbound.go
+++ b/proxy/vless/inbound/inbound.go
@@ -495,7 +495,7 @@ func (h *Handler) Process(ctx context.Context, network net.Network, connection s
 			return newError(account.ID.String() + " is not able to use " + requestAddons.Flow).AtWarning()
 		}
 	case "":
-		if account.Flow == vless.XRV && (request.Command == protocol.RequestCommandTCP || isMuxAndNotXUDP(request, first)) {
+		if account.Flow == vless.XRV && request.Command == protocol.RequestCommandTCP {
 			return newError(account.ID.String() + " is not able to use \"\". Note that the pure TLS proxy has certain TLS in TLS characters.").AtWarning()
 		}
 	default:
@@ -510,6 +510,8 @@ func (h *Handler) Process(ctx context.Context, network net.Network, connection s
 			Reason: "",
 			Email:  request.User.Email,
 		})
+	} else if account.Flow == vless.XRV {
+		ctx = session.ContextWithAllowedNetwork(ctx, net.Network_UDP)
 	}
 
 	sessionPolicy = h.policyManager.ForLevel(request.User.Level)