trojan,vless: use proxy.RelayUDP to support ReadDeadline

proxy/tproxy/tproxy_linux.go

@@ -86,20 +86,19 @@ func (s *TProxy) ListenAndServeUDP() {
 			continue
 		}
 
-		var session *Session
+		s.handleMsg(srcAddr, dstAddr, buf[:n])
-		sessionKey := srcAddr.String()
+	}
+}
 
-		v, ok := nm.Load(sessionKey)
+// handleMsg handles an udp message.
-		if !ok && v == nil {
+func (s *TProxy) handleMsg(srcAddr, dstAddr *net.UDPAddr, data []byte) {
-			session = newSession(sessionKey, srcAddr, dstAddr)
+	var session *Session
-			nm.Store(sessionKey, session)
+	sessionKey := srcAddr.String()
-			go s.ServeSession(session)
-			session.msgCh <- buf[:n]
-			continue
-		}
 
+	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.
+// serveSession serves a udp session.
-func (s *TProxy) ServeSession(session *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)
 	}

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)
+	go proxy.RelayUDP(rc, tgtAddr, pc, 2*time.Minute)
-	defer pool.PutBuffer(buf)
+	proxy.RelayUDP(pc, nil, rc, 2*time.Minute)
-
-	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
-		}
-	}
 }

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)
+	go proxy.RelayUDP(rc, tgtAddr, pc, 2*time.Minute)
-	defer pool.PutBuffer(buf)
+	proxy.RelayUDP(pc, nil, rc, 2*time.Minute)
-
-	for {
-		n, _, err := rc.ReadFrom(buf)
-		if err != nil {
-			break
-		}
-
-		_, err = pc.WriteTo(buf[:n], nil)
-		if err != nil {
-			break
-		}
-	}
 }
 
 // ServerConn is a vless client connection.