From d8af58cec756e0081dcc64f4a4f5681146403d19 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Sat, 29 Feb 2020 15:45:58 +0800
Subject: [PATCH] reduce buffer allocation

---
 relay.go  | 10 ++++++++--
 tuntap.go |  8 +++++++-
 udp.go    |  7 ++++++-
 3 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/relay.go b/relay.go
index d3157cd..9f3a30d 100644
--- a/relay.go
+++ b/relay.go
@@ -284,7 +284,7 @@ func (c *relayConn) Read(b []byte) (n int, err error) {
 	})
 
 	if err != nil {
-		log.Log("[relay] %s <- %s: %s", c.Conn.LocalAddr(), c.Conn.RemoteAddr(), err)
+		log.Logf("[relay] %s <- %s: %s", c.Conn.LocalAddr(), c.Conn.RemoteAddr(), err)
 		return
 	}
 
@@ -334,7 +334,13 @@ func (c *relayConn) Write(b []byte) (n int, err error) {
 	if !c.udp {
 		return c.Conn.Write(b)
 	}
-	buf := make([]byte, 2+len(b))
+	var buf []byte
+	if 2+len(b) <= mediumBufferSize {
+		buf = mPool.Get().([]byte)
+		defer mPool.Put(buf)
+	} else {
+		buf = make([]byte, 2+len(b))
+	}
 	binary.BigEndian.PutUint16(buf[:2], uint16(len(b)))
 	n = copy(buf[2:], b)
 	_, err = c.Conn.Write(buf)
diff --git a/tuntap.go b/tuntap.go
index 8e429b6..03feb4d 100644
--- a/tuntap.go
+++ b/tuntap.go
@@ -563,7 +563,13 @@ func (h *tapHandler) Handle(conn net.Conn) {
 				if err != nil {
 					return err
 				}
-				pc = cc.(net.PacketConn)
+				var ok bool
+				pc, ok = cc.(net.PacketConn)
+				if !ok {
+					err = errors.New("not a packet connection")
+					log.Logf("[tap] %s - %s: %s", conn.LocalAddr(), raddr, err)
+					return err
+				}
 			} else {
 				if h.options.TCPMode {
 					if raddr != nil {
diff --git a/udp.go b/udp.go
index 5533a3b..c2a71ab 100644
--- a/udp.go
+++ b/udp.go
@@ -89,7 +89,9 @@ func UDPListener(addr string, cfg *UDPListenConfig) (Listener, error) {
 
 func (l *udpListener) listenLoop() {
 	for {
-		b := make([]byte, mediumBufferSize)
+		// NOTE: this buffer will be released in the udpServerConn after read.
+		b := mPool.Get().([]byte)
+
 		n, raddr, err := l.ln.ReadFrom(b)
 		if err != nil {
 			log.Logf("[udp] peer -> %s : %s", l.Addr(), err)
@@ -240,6 +242,9 @@ func (c *udpServerConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
 	select {
 	case bb := <-c.rChan:
 		n = copy(b, bb)
+		if cap(bb) == mediumBufferSize {
+			mPool.Put(bb[:cap(bb)])
+		}
 	case <-c.closed:
 		err = errors.New("read from closed connection")
 		return