From dc9ae1edd0abecf1f8028bf028d8a2bc0278862d Mon Sep 17 00:00:00 2001
From: gVisor bot <gvisor-bot@google.com>
Date: Thu, 16 Apr 2020 18:19:36 +0800
Subject: [PATCH] Improve: recycle buffer after packet used

---
 constant/adapters.go |  4 ++--
 proxy/redir/udp.go   |  9 ++++-----
 proxy/redir/utils.go | 14 ++++++--------
 proxy/socks/udp.go   | 10 +++++-----
 proxy/socks/utils.go | 17 ++++++++---------
 tunnel/connection.go |  2 ++
 6 files changed, 27 insertions(+), 29 deletions(-)

diff --git a/constant/adapters.go b/constant/adapters.go
index ac19de36..7900117c 100644
--- a/constant/adapters.go
+++ b/constant/adapters.go
@@ -135,8 +135,8 @@ type UDPPacket interface {
 	//   this is important when using Fake-IP.
 	WriteBack(b []byte, addr net.Addr) (n int, err error)
 
-	// Close closes the underlaying connection.
-	Close() error
+	// Drop call after packet is used, could recycle buffer in this function.
+	Drop()
 
 	// LocalAddr returns the source IP/Port of packet
 	LocalAddr() net.Addr
diff --git a/proxy/redir/udp.go b/proxy/redir/udp.go
index 89209aa9..9b342996 100644
--- a/proxy/redir/udp.go
+++ b/proxy/redir/udp.go
@@ -66,10 +66,9 @@ func (l *RedirUDPListener) Address() string {
 
 func handleRedirUDP(pc net.PacketConn, buf []byte, lAddr *net.UDPAddr, rAddr *net.UDPAddr) {
 	target := socks5.ParseAddrToSocksAddr(rAddr)
-	packet := &fakeConn{
-		PacketConn: pc,
-		lAddr:      lAddr,
-		buf:        buf,
+	pkt := &packet{
+		lAddr: lAddr,
+		buf:   buf,
 	}
-	tunnel.AddPacket(adapters.NewPacket(target, packet, C.REDIR))
+	tunnel.AddPacket(adapters.NewPacket(target, pkt, C.REDIR))
 }
diff --git a/proxy/redir/utils.go b/proxy/redir/utils.go
index f6a60d46..5c737889 100644
--- a/proxy/redir/utils.go
+++ b/proxy/redir/utils.go
@@ -6,18 +6,17 @@ import (
 	"github.com/Dreamacro/clash/common/pool"
 )
 
-type fakeConn struct {
-	net.PacketConn
+type packet struct {
 	lAddr *net.UDPAddr
 	buf   []byte
 }
 
-func (c *fakeConn) Data() []byte {
+func (c *packet) Data() []byte {
 	return c.buf
 }
 
 // WriteBack opens a new socket binding `addr` to wirte UDP packet back
-func (c *fakeConn) WriteBack(b []byte, addr net.Addr) (n int, err error) {
+func (c *packet) WriteBack(b []byte, addr net.Addr) (n int, err error) {
 	tc, err := dialUDP("udp", addr.(*net.UDPAddr), c.lAddr)
 	if err != nil {
 		n = 0
@@ -29,12 +28,11 @@ func (c *fakeConn) WriteBack(b []byte, addr net.Addr) (n int, err error) {
 }
 
 // LocalAddr returns the source IP/Port of UDP Packet
-func (c *fakeConn) LocalAddr() net.Addr {
+func (c *packet) LocalAddr() net.Addr {
 	return c.lAddr
 }
 
-func (c *fakeConn) Close() error {
-	err := c.PacketConn.Close()
+func (c *packet) Drop() {
 	pool.BufPool.Put(c.buf[:cap(c.buf)])
-	return err
+	return
 }
diff --git a/proxy/socks/udp.go b/proxy/socks/udp.go
index 1e623fb5..9d77139c 100644
--- a/proxy/socks/udp.go
+++ b/proxy/socks/udp.go
@@ -63,11 +63,11 @@ func handleSocksUDP(pc net.PacketConn, buf []byte, addr net.Addr) {
 		pool.BufPool.Put(buf[:cap(buf)])
 		return
 	}
-	packet := &fakeConn{
-		PacketConn: pc,
-		rAddr:      addr,
-		payload:    payload,
-		bufRef:     buf,
+	packet := &packet{
+		pc:      pc,
+		rAddr:   addr,
+		payload: payload,
+		bufRef:  buf,
 	}
 	tunnel.AddPacket(adapters.NewPacket(target, packet, C.SOCKS))
 }
diff --git a/proxy/socks/utils.go b/proxy/socks/utils.go
index b4ef429b..3ee52f5f 100644
--- a/proxy/socks/utils.go
+++ b/proxy/socks/utils.go
@@ -7,33 +7,32 @@ import (
 	"github.com/Dreamacro/clash/component/socks5"
 )
 
-type fakeConn struct {
-	net.PacketConn
+type packet struct {
+	pc      net.PacketConn
 	rAddr   net.Addr
 	payload []byte
 	bufRef  []byte
 }
 
-func (c *fakeConn) Data() []byte {
+func (c *packet) Data() []byte {
 	return c.payload
 }
 
 // WriteBack wirtes UDP packet with source(ip, port) = `addr`
-func (c *fakeConn) WriteBack(b []byte, addr net.Addr) (n int, err error) {
+func (c *packet) WriteBack(b []byte, addr net.Addr) (n int, err error) {
 	packet, err := socks5.EncodeUDPPacket(socks5.ParseAddrToSocksAddr(addr), b)
 	if err != nil {
 		return
 	}
-	return c.PacketConn.WriteTo(packet, c.rAddr)
+	return c.pc.WriteTo(packet, c.rAddr)
 }
 
 // LocalAddr returns the source IP/Port of UDP Packet
-func (c *fakeConn) LocalAddr() net.Addr {
+func (c *packet) LocalAddr() net.Addr {
 	return c.rAddr
 }
 
-func (c *fakeConn) Close() error {
-	err := c.PacketConn.Close()
+func (c *packet) Drop() {
 	pool.BufPool.Put(c.bufRef[:cap(c.bufRef)])
-	return err
+	return
 }
diff --git a/tunnel/connection.go b/tunnel/connection.go
index ce70a18b..f9fa5038 100644
--- a/tunnel/connection.go
+++ b/tunnel/connection.go
@@ -82,6 +82,8 @@ func handleHTTP(request *adapters.HTTPAdapter, outbound net.Conn) {
 }
 
 func handleUDPToRemote(packet C.UDPPacket, pc C.PacketConn, metadata *C.Metadata) {
+	defer packet.Drop()
+
 	if _, err := pc.WriteWithMetadata(packet.Data(), metadata); err != nil {
 		return
 	}