From 6d369f37aca35e20c0efc1a6329849edf654b8f3 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Sun, 18 Oct 2015 20:30:40 +0800
Subject: [PATCH] add udp tunnel

---
 socks.go | 85 ++++++++++++++++++++++++++++++++++++++++++++++++++------
 udp.go   | 39 ++++++++++++++++++++++++++
 2 files changed, 115 insertions(+), 9 deletions(-)
 create mode 100644 udp.go

diff --git a/socks.go b/socks.go
index 3d05b8e..191b73a 100644
--- a/socks.go
+++ b/socks.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"crypto/tls"
 	"errors"
 	"github.com/ginuerzh/gosocks5"
@@ -268,24 +269,90 @@ func handleSocks5Request(req *gosocks5.Request, conn net.Conn) {
 			}
 		}
 
-		srvTunnelUDP(conn, uconn)
-	}
-}
+		clientConn, dgram, err := createClientConn(conn, uconn)
+		if err != nil {
+			if glog.V(LWARNING) {
+				glog.Warningln("socks5 udp:", err)
+			}
+			return
+		}
+		if glog.V(LDEBUG) {
+			glog.Infof("[udp] length %d, to %s", len(dgram.Data), dgram.Header.Addr)
+		}
 
-func serveUDP(conn *net.UDPConn) {
-	if len(forwardArgs) > 0 {
-		fconn, _, err = forwardChain(forwardArgs...)
+		serverConn, err := createServerConn(uconn)
 		if err != nil {
 			if glog.V(LWARNING) {
 				glog.Warningln("socks5 udp forward:", err)
 			}
-			if fconn != nil {
-				fconn.Close()
-			}
+		}
+	default:
+		if glog.V(LWARNING) {
+			glog.Warningln("Unrecognized request: ", req)
+		}
+	}
+}
+
+func createClientConn(conn net.Conn, uconn *net.UDPConn) (c *UDPConn, dgram *gosocks5.UDPDatagram, err error) {
+	var raddr *net.UDPAddr
+	dgramChan := make(chan *gosocks5.UDPDatagram, 1)
+	errChan := make(chan error, 1)
+	go func() {
+		b := make([]byte, 64*1024+262)
+
+		n, addr, err := uconn.ReadFromUDP(b)
+		if err != nil {
+			errChan <- err
 			return
 		}
+		raddr = addr
 
+		dgram, err := gosocks5.ReadUDPDatagram(bytes.NewReader(b[:n]))
+		if err != nil {
+			errChan <- err
+			return
+		}
+		dgramChan <- dgram
+	}()
+
+	go func() {
+		dgram, err := gosocks5.ReadUDPDatagram(conn)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		dgramChan <- dgram
+	}()
+
+	select {
+	case dgram = <-dgramChan:
+		if raddr != nil {
+			c = Client(uconn, raddr)
+		} else {
+			c = Client(conn, nil)
+		}
+	case err = <-errChan:
 	}
+
+	return
+}
+
+func createServerConn(uconn *net.UDPConn) (c *UDPConn, err error) {
+	if len(forwardArgs) == 0 {
+		c = Server(uconn)
+		return
+	}
+
+	fconn, _, err := forwardChain(forwardArgs...)
+	if err != nil {
+		if fconn != nil {
+			fconn.Close()
+		}
+		return
+	}
+
+	c = Server(fconn)
+	return
 }
 
 func forwardUDP(req *gosocks5.Request) (conn net.Conn, err error) {
diff --git a/udp.go b/udp.go
new file mode 100644
index 0000000..4301be3
--- /dev/null
+++ b/udp.go
@@ -0,0 +1,39 @@
+package main
+
+import (
+	"github.com/ginuerzh/gosocks5"
+	"github.com/golang/glog"
+	"net"
+)
+
+type UDPConn struct {
+	isClient bool
+	udpConn  *net.UDPConn
+	udpAddr  *net.UDPAddr
+	tcpConn  net.Conn
+}
+
+func Client(conn net.Conn, addr net.Addr) *UDPConn {
+	client := &UDPConn{isClient: true}
+
+	switch conn := conn.(type) {
+	case net.UDPConn:
+		client.udpConn = conn
+		client.udpAddr = addr
+	default:
+		client.tcpConn = conn
+	}
+
+	return client
+}
+
+func Server(conn net.Conn) *UDPConn {
+	server := &UDPConn{}
+	switch conn := conn.(type) {
+	case net.UDPConn:
+		server.udpConn = conn
+	default:
+		server.tcpConn = conn
+	}
+	return server
+}