From 01040ca1fff5b43da845c730ee5baef5b4f8f784 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Tue, 27 Oct 2015 22:22:08 +0800
Subject: [PATCH] socks5 udp tunnel

---
 conn.go  |  5 +++--
 socks.go |  9 +++++----
 udp.go   | 31 +++++++++++++++++++++++++++++--
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/conn.go b/conn.go
index cefeb56..baf9be8 100644
--- a/conn.go
+++ b/conn.go
@@ -15,6 +15,7 @@ import (
 	"strconv"
 	"strings"
 	"sync/atomic"
+	"time"
 )
 
 var (
@@ -204,7 +205,7 @@ func Connect(addr string) (conn net.Conn, err error) {
 		addr += ":80"
 	}
 	if len(forwardArgs) == 0 {
-		return net.Dial("tcp", addr)
+		return net.DialTimeout("tcp", addr, time.Second*30)
 	}
 
 	var end Args
@@ -224,7 +225,7 @@ func Connect(addr string) (conn net.Conn, err error) {
 
 func forwardChain(chain ...Args) (conn net.Conn, end Args, err error) {
 	end = chain[0]
-	if conn, err = net.Dial("tcp", end.Addr); err != nil {
+	if conn, err = net.DialTimeout("tcp", end.Addr, time.Second*30); err != nil {
 		return
 	}
 	c, err := forward(conn, end)
diff --git a/socks.go b/socks.go
index 3bce504..e01402c 100644
--- a/socks.go
+++ b/socks.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"net"
 	"strconv"
+	"time"
 )
 
 const (
@@ -235,18 +236,18 @@ func handleSocks5Request(req *gosocks5.Request, conn net.Conn) {
 			return
 		}
 
-		if err = sc.WriteUDP(dgram); err != nil {
+		if err = sc.WriteUDPTimeout(dgram, time.Second*30); err != nil {
 			glog.V(LWARNING).Infoln("socks5 udp:", err)
 			return
 		}
-		dgram, err = sc.ReadUDP()
+		dgram, err = sc.ReadUDPTimeout(time.Second * 30)
 		if err != nil {
 			glog.V(LWARNING).Infoln("socks5 udp:", err)
 			return
 		}
 		glog.V(LDEBUG).Infof("[udp] from %s, length %d", dgram.Header.Addr, len(dgram.Data))
 
-		if err = cc.WriteUDP(dgram); err != nil {
+		if err = cc.WriteUDPTimeout(dgram, time.Second*30); err != nil {
 			glog.V(LWARNING).Infoln("socks5 udp:", err)
 			return
 		}
@@ -468,7 +469,7 @@ func PipeUDP(src, dst *UDPConn, ch chan<- error) {
 		if err != nil {
 			break
 		}
-		glog.V(LDEBUG).Infof("[udp] addr %s, length %d", dgram.Header.Addr, len(dgram.Data))
+		// glog.V(LDEBUG).Infof("[udp] addr %s, length %d", dgram.Header.Addr, len(dgram.Data))
 
 		if err = dst.WriteUDP(dgram); err != nil {
 			break
diff --git a/udp.go b/udp.go
index eb53f49..e3ce064 100644
--- a/udp.go
+++ b/udp.go
@@ -49,6 +49,20 @@ func (c *UDPConn) ReadUDP() (*gosocks5.UDPDatagram, error) {
 	return c.readUDPServer()
 }
 
+func (c *UDPConn) ReadUDPTimeout(timeout time.Duration) (*gosocks5.UDPDatagram, error) {
+	if c.udp != nil {
+		c.udp.SetReadDeadline(time.Now().Add(timeout))
+		defer c.udp.SetReadDeadline(time.Time{})
+	} else {
+		c.tcp.SetReadDeadline(time.Now().Add(timeout))
+		defer c.tcp.SetReadDeadline(time.Time{})
+	}
+	if c.isClient {
+		return c.readUDPClient()
+	}
+	return c.readUDPServer()
+}
+
 func (c *UDPConn) readUDPClient() (*gosocks5.UDPDatagram, error) {
 	if c.udp != nil {
 		return gosocks5.ReadUDPDatagram(c.udp)
@@ -59,7 +73,7 @@ func (c *UDPConn) readUDPClient() (*gosocks5.UDPDatagram, error) {
 func (c *UDPConn) readUDPServer() (*gosocks5.UDPDatagram, error) {
 	if c.udp != nil {
 		b := make([]byte, 65535)
-		n, addr, err := c.udp.ReadFromUDP(b)
+		n, addr, err := c.udp.ReadFrom(b)
 		if err != nil {
 			return nil, err
 		}
@@ -67,7 +81,6 @@ func (c *UDPConn) readUDPServer() (*gosocks5.UDPDatagram, error) {
 			gosocks5.NewUDPHeader(0, 0, ToSocksAddr(addr)), b[:n])
 		return dgram, nil
 	}
-
 	return gosocks5.ReadUDPDatagram(c.tcp)
 }
 
@@ -78,6 +91,20 @@ func (c *UDPConn) WriteUDP(dgram *gosocks5.UDPDatagram) error {
 	return c.writeUDPServer(dgram)
 }
 
+func (c *UDPConn) WriteUDPTimeout(dgram *gosocks5.UDPDatagram, timeout time.Duration) error {
+	if c.udp != nil {
+		c.udp.SetWriteDeadline(time.Now().Add(timeout))
+		defer c.udp.SetWriteDeadline(time.Time{})
+	} else {
+		c.tcp.SetWriteDeadline(time.Now().Add(timeout))
+		defer c.tcp.SetWriteDeadline(time.Time{})
+	}
+	if c.isClient {
+		return c.writeUDPClient(dgram)
+	}
+	return c.writeUDPServer(dgram)
+}
+
 func (c *UDPConn) writeUDPClient(dgram *gosocks5.UDPDatagram) error {
 	if c.udp != nil {
 		dgram.Header.Rsv = 0