From 59ab4fe745e2ecbc3b4676d2b7d54c50d4f290ca Mon Sep 17 00:00:00 2001
From: Larvan2 <78135608+Larvan2@users.noreply.github.com>
Date: Tue, 19 Dec 2023 19:59:37 +0800
Subject: [PATCH] chore: better Reject-Drop for UDP

---
 adapter/outbound/reject.go | 22 ----------------------
 tunnel/tunnel.go           |  5 ++++-
 2 files changed, 4 insertions(+), 23 deletions(-)

diff --git a/adapter/outbound/reject.go b/adapter/outbound/reject.go
index b564e28d..62f4aaa5 100644
--- a/adapter/outbound/reject.go
+++ b/adapter/outbound/reject.go
@@ -30,9 +30,6 @@ func (r *Reject) DialContext(ctx context.Context, metadata *C.Metadata, opts ...
 
 // ListenPacketContext implements C.ProxyAdapter
 func (r *Reject) ListenPacketContext(ctx context.Context, metadata *C.Metadata, opts ...dialer.Option) (C.PacketConn, error) {
-	if r.drop {
-		return newPacketConn(&dropPacketConn{}, r), nil
-	}
 	return newPacketConn(&nopPacketConn{}, r), nil
 }
 
@@ -129,22 +126,3 @@ func (rw dropConn) RemoteAddr() net.Addr                 { return nil }
 func (rw dropConn) SetDeadline(time.Time) error          { return nil }
 func (rw dropConn) SetReadDeadline(time.Time) error      { return nil }
 func (rw dropConn) SetWriteDeadline(time.Time) error     { return nil }
-
-type dropPacketConn struct{}
-
-func (npc dropPacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
-	time.Sleep(C.DefaultDropTime)
-	return len(b), nil
-}
-func (npc dropPacketConn) ReadFrom(b []byte) (int, net.Addr, error) {
-	time.Sleep(C.DefaultDropTime)
-	return 0, nil, io.EOF
-}
-func (npc dropPacketConn) WaitReadFrom() ([]byte, func(), net.Addr, error) {
-	return nil, nil, nil, io.EOF
-}
-func (npc dropPacketConn) Close() error                     { return nil }
-func (npc dropPacketConn) LocalAddr() net.Addr              { return udpAddrIPv4Unspecified }
-func (npc dropPacketConn) SetDeadline(time.Time) error      { return nil }
-func (npc dropPacketConn) SetReadDeadline(time.Time) error  { return nil }
-func (npc dropPacketConn) SetWriteDeadline(time.Time) error { return nil }
diff --git a/tunnel/tunnel.go b/tunnel/tunnel.go
index 9891d96f..71772e71 100644
--- a/tunnel/tunnel.go
+++ b/tunnel/tunnel.go
@@ -355,7 +355,6 @@ func handleUDPConn(packet C.PacketAdapter) {
 		if !metadata.Resolved() && proxy.Type() != C.Reject && proxy.Type() != C.RejectDrop {
 			ip, err := resolver.ResolveIP(context.Background(), metadata.Host)
 			if err != nil {
-				packet.Drop()
 				return
 			}
 			metadata.DstIP = ip
@@ -409,6 +408,10 @@ func handleUDPConn(packet C.PacketAdapter) {
 		case rule != nil:
 			if rule.Payload() != "" {
 				log.Infoln("[UDP] %s --> %s match %s using %s", metadata.SourceDetail(), metadata.RemoteAddress(), fmt.Sprintf("%s(%s)", rule.RuleType().String(), rule.Payload()), rawPc.Chains().String())
+				if rawPc.Chains().Last() == "REJECT-DROP" {
+					pc.Close()
+					return
+				}
 			} else {
 				log.Infoln("[UDP] %s --> %s match %s using %s", metadata.SourceDetail(), metadata.RemoteAddress(), rule.Payload(), rawPc.Chains().String())
 			}