From 9ee9a0634e5d789d7f014596d87fd770d6ab03d4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=A3=8E=E6=89=87=E6=BB=91=E7=BF=94=E7=BF=BC?=
 <Fangliding.fshxy@outlook.com>
Date: Wed, 22 May 2024 11:02:45 +0800
Subject: [PATCH] Add UDPFilter to Socks5 server when `auth == password`
 (#3371)

Co-authored-by: RPRX <63339210+RPRX@users.noreply.github.com>
---
 proxy/socks/server.go    | 11 +++++++++++
 proxy/socks/udpfilter.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)
 create mode 100644 proxy/socks/udpfilter.go

diff --git a/proxy/socks/server.go b/proxy/socks/server.go
index 0109d5b4..18084864 100644
--- a/proxy/socks/server.go
+++ b/proxy/socks/server.go
@@ -27,6 +27,7 @@ type Server struct {
 	config        *ServerConfig
 	policyManager policy.Manager
 	cone          bool
+	udpFilter     *UDPFilter
 }
 
 // NewServer creates a new Server object.
@@ -37,6 +38,9 @@ func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) {
 		policyManager: v.GetFeature(policy.ManagerType()).(policy.Manager),
 		cone:          ctx.Value("cone").(bool),
 	}
+	if config.AuthType == AuthType_PASSWORD {
+		s.udpFilter = new(UDPFilter) // We only use this when auth is enabled
+	}
 	return s, nil
 }
 
@@ -135,6 +139,9 @@ func (s *Server) processTCP(ctx context.Context, conn stat.Connection, dispatche
 	}
 
 	if request.Command == protocol.RequestCommandUDP {
+		if s.udpFilter != nil {
+			s.udpFilter.Add(conn.RemoteAddr())
+		}
 		return s.handleUDP(conn)
 	}
 
@@ -193,6 +200,10 @@ func (s *Server) transport(ctx context.Context, reader io.Reader, writer io.Writ
 }
 
 func (s *Server) handleUDPPayload(ctx context.Context, conn stat.Connection, dispatcher routing.Dispatcher) error {
+	if s.udpFilter != nil && !s.udpFilter.Check(conn.RemoteAddr()) {
+		newError("Unauthorized UDP access from ", conn.RemoteAddr().String()).AtDebug().WriteToLog(session.ExportIDToError(ctx))
+		return nil
+	}
 	udpServer := udp.NewDispatcher(dispatcher, func(ctx context.Context, packet *udp_proto.Packet) {
 		payload := packet.Payload
 		newError("writing back UDP response with ", payload.Len(), " bytes").AtDebug().WriteToLog(session.ExportIDToError(ctx))
diff --git a/proxy/socks/udpfilter.go b/proxy/socks/udpfilter.go
new file mode 100644
index 00000000..25f8a416
--- /dev/null
+++ b/proxy/socks/udpfilter.go
@@ -0,0 +1,31 @@
+package socks
+
+import (
+	"net"
+	"sync"
+)
+
+/*
+In the sock implementation of * ray, UDP authentication is flawed and can be bypassed.
+Tracking a UDP connection may be a bit troublesome.
+Here is a simple solution.
+We creat a filter, add remote IP to the pool when it try to establish a UDP connection with auth.
+And drop UDP packets from unauthorized IP.
+After discussion, we believe it is not necessary to add a timeout mechanism to this filter.
+*/
+
+type UDPFilter struct {
+	ips sync.Map
+}
+
+func (f *UDPFilter) Add(addr net.Addr) bool {
+	ip, _, _ := net.SplitHostPort(addr.String())
+	f.ips.Store(ip, true)
+	return true
+}
+
+func (f *UDPFilter) Check(addr net.Addr) bool {
+	ip, _, _ := net.SplitHostPort(addr.String())
+	_, ok := f.ips.Load(ip)
+	return ok
+}