From a08c939dacbccba0d500e9b8ec0d19895254e8ec Mon Sep 17 00:00:00 2001
From: nadoo <287492+nadoo@users.noreply.github.com>
Date: Mon, 26 Jul 2021 23:14:35 +0800
Subject: [PATCH] trojan,vless: use proxy.RelayUDP to support ReadDeadline

---
 proxy/tproxy/tproxy_linux.go | 34 ++++++++++++++++++++--------------
 proxy/trojan/server.go       | 35 +++--------------------------------
 proxy/vless/server.go        | 34 +++-------------------------------
 3 files changed, 26 insertions(+), 77 deletions(-)

diff --git a/proxy/tproxy/tproxy_linux.go b/proxy/tproxy/tproxy_linux.go
index 4f9fff4..28afeac 100644
--- a/proxy/tproxy/tproxy_linux.go
+++ b/proxy/tproxy/tproxy_linux.go
@@ -86,20 +86,19 @@ func (s *TProxy) ListenAndServeUDP() {
 			continue
 		}
 
-		var session *Session
-		sessionKey := srcAddr.String()
+		s.handleMsg(srcAddr, dstAddr, buf[:n])
+	}
+}
 
-		v, ok := nm.Load(sessionKey)
-		if !ok && v == nil {
-			session = newSession(sessionKey, srcAddr, dstAddr)
-			nm.Store(sessionKey, session)
-			go s.ServeSession(session)
-			session.msgCh <- buf[:n]
-			continue
-		}
+// handleMsg handles an udp message.
+func (s *TProxy) handleMsg(srcAddr, dstAddr *net.UDPAddr, data []byte) {
+	var session *Session
+	sessionKey := srcAddr.String()
 
+	v, ok := nm.Load(sessionKey)
+	if ok && v != nil {
 		session = v.(*Session)
-		session.msgCh <- buf[:n]
+		session.msgCh <- data
 
 		select {
 		case <-session.finCh:
@@ -107,12 +106,19 @@ func (s *TProxy) ListenAndServeUDP() {
 			close(session.msgCh)
 			close(session.finCh)
 		default:
+			return
 		}
 	}
+
+	session = newSession(sessionKey, srcAddr, dstAddr)
+	nm.Store(sessionKey, session)
+
+	go s.serveSession(session)
+	session.msgCh <- data
 }
 
-// ServeSession serves a udp session.
-func (s *TProxy) ServeSession(session *Session) {
+// serveSession serves a udp session.
+func (s *TProxy) serveSession(session *Session) {
 	dstPC, dialer, writeTo, err := s.proxy.DialUDP("udp", session.dst.String())
 	if err != nil {
 		log.F("[tproxyu] dial to %s error: %v", session.dst, err)
@@ -137,7 +143,7 @@ func (s *TProxy) ServeSession(session *Session) {
 	for data := range session.msgCh {
 		_, err = dstPC.WriteTo(data, writeTo)
 		if err != nil {
-			log.F("[tproxyu] writeTo %s error: %v", writeTo, err)
+			log.F("[tproxyu] writeTo error: %v", err)
 		}
 		pool.PutBuffer(data)
 	}
diff --git a/proxy/trojan/server.go b/proxy/trojan/server.go
index ac75604..f3f0ff7 100644
--- a/proxy/trojan/server.go
+++ b/proxy/trojan/server.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"net"
 	"strings"
+	"time"
 
 	"github.com/nadoo/glider/log"
 	"github.com/nadoo/glider/pool"
@@ -203,38 +204,8 @@ func (s *Trojan) ServeUoT(c net.Conn, tgt socks.Addr) {
 	}
 
 	pc := NewPktConn(c, tgt)
-
-	go func() {
-		buf := pool.GetBuffer(proxy.UDPBufSize)
-		defer pool.PutBuffer(buf)
-		for {
-			n, _, err := pc.ReadFrom(buf)
-			if err != nil {
-				return
-			}
-
-			_, err = rc.WriteTo(buf[:n], tgtAddr)
-			if err != nil {
-				return
-			}
-		}
-	}()
-
 	log.F("[trojan] %s <-tcp-> %s - %s <-udp-> %s", c.RemoteAddr(), c.LocalAddr(), rc.LocalAddr(), tgt)
 
-	buf := pool.GetBuffer(proxy.UDPBufSize)
-	defer pool.PutBuffer(buf)
-
-	for {
-		n, _, err := rc.ReadFrom(buf)
-		if err != nil {
-			break
-		}
-
-		// WriteTo addr can be nil because the PktConn has it's own target, see packet.go
-		_, err = pc.WriteTo(buf[:n], nil)
-		if err != nil {
-			break
-		}
-	}
+	go proxy.RelayUDP(rc, tgtAddr, pc, 2*time.Minute)
+	proxy.RelayUDP(pc, nil, rc, 2*time.Minute)
 }
diff --git a/proxy/vless/server.go b/proxy/vless/server.go
index 78c80c2..97de8aa 100644
--- a/proxy/vless/server.go
+++ b/proxy/vless/server.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"net"
 	"strings"
+	"time"
 
 	"github.com/nadoo/glider/log"
 	"github.com/nadoo/glider/pool"
@@ -162,39 +163,10 @@ func (s *VLess) ServeUoT(c net.Conn, tgt string) {
 	}
 
 	pc := NewPktConn(c)
-
-	go func() {
-		buf := pool.GetBuffer(proxy.UDPBufSize)
-		defer pool.PutBuffer(buf)
-		for {
-			n, _, err := pc.ReadFrom(buf)
-			if err != nil {
-				return
-			}
-
-			_, err = rc.WriteTo(buf[:n], tgtAddr)
-			if err != nil {
-				return
-			}
-		}
-	}()
-
 	log.F("[vless] %s <-tcp-> %s - %s <-udp-> %s", c.RemoteAddr(), c.LocalAddr(), rc.LocalAddr(), tgt)
 
-	buf := pool.GetBuffer(proxy.UDPBufSize)
-	defer pool.PutBuffer(buf)
-
-	for {
-		n, _, err := rc.ReadFrom(buf)
-		if err != nil {
-			break
-		}
-
-		_, err = pc.WriteTo(buf[:n], nil)
-		if err != nil {
-			break
-		}
-	}
+	go proxy.RelayUDP(rc, tgtAddr, pc, 2*time.Minute)
+	proxy.RelayUDP(pc, nil, rc, 2*time.Minute)
 }
 
 // ServerConn is a vless client connection.