add remote udp port forwarding

conn.go

@@ -29,7 +29,7 @@ var (
 	// tcp buffer pool
 	tcpPool = sync.Pool{
 		New: func() interface{} {
-			return make([]byte, 16*1024)
+			return make([]byte, 32*1024)
 		},
 	}
 	// udp buffer pool
@@ -59,6 +59,7 @@ func listenAndServe(arg Args) error {
 	case "rtcp": // Remote TCP port forwarding
 		return serveRTcpForward(arg)
 	case "rudp": // Remote UDP port forwarding
+		return serveRUdpForward(arg)
 	default:
 		ln, err = net.Listen("tcp", arg.Addr)
 	}
@@ -148,6 +149,31 @@ func serveRTcpForward(arg Args) error {
 	}
 }
 
+func serveRUdpForward(arg Args) error {
+	if len(forwardArgs) == 0 {
+		return errors.New("rudp: at least one -F must be assigned")
+	}
+
+	retry := 0
+	for {
+		conn, _, err := forwardChain(forwardArgs...)
+		if err != nil {
+			glog.V(LWARNING).Infof("[rudp] %s - %s : %s", arg.Addr, arg.Remote, err)
+			time.Sleep((1 << uint(retry)) * time.Second)
+			if retry < 5 {
+				retry++
+			}
+			continue
+		}
+		retry = 0
+
+		if err := connectRUdpForward(conn, arg); err != nil {
+			conn.Close()
+			time.Sleep(10 * time.Second)
+		}
+	}
+}
+
 func handleConn(conn net.Conn, arg Args) {
 	atomic.AddInt32(&connCounter, 1)
 	glog.V(LDEBUG).Infof("%s connected, connections: %d",

forward.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"errors"
+	"fmt"
 	"github.com/ginuerzh/gosocks5"
 	"github.com/golang/glog"
 	"net"
@@ -72,27 +73,32 @@ func handleUdpForward(conn *net.UDPConn, raddr *net.UDPAddr, data []byte, arg Ar
 		return
 	}
 
-	fconn, _, err := forwardChain(forwardArgs...)
+	tun, _, err := forwardChain(forwardArgs...)
 	if err != nil {
 		glog.V(LWARNING).Infof("[udp-forward] %s -> %s : %s", raddr, arg.Remote, err)
 		return
 	}
-	defer fconn.Close()
+	defer tun.Close()
 
 	glog.V(LINFO).Infof("[udp-forward] %s -> %s ASSOCIATE", raddr, arg.Remote)
 
 	req := gosocks5.NewRequest(CmdUdpTun, nil)
-	if err = req.Write(fconn); err != nil {
+	tun.SetWriteDeadline(time.Now().Add(time.Second * 90))
+	if err = req.Write(tun); err != nil {
 		glog.V(LWARNING).Infof("[udp-forward] %s -> %s ASSOCIATE : %s", raddr, arg.Remote, err)
 		return
 	}
+	tun.SetWriteDeadline(time.Time{})
 	glog.V(LDEBUG).Infof("[udp-forward] %s -> %s\n%s", raddr, arg.Remote, req)
 
-	rep, err := gosocks5.ReadReply(fconn)
+	tun.SetReadDeadline(time.Now().Add(90 * time.Second))
+	rep, err := gosocks5.ReadReply(tun)
 	if err != nil {
 		glog.V(LWARNING).Infof("[udp-forward] %s <- %s ASSOCIATE : %s", raddr, arg.Remote, err)
 		return
 	}
+	tun.SetReadDeadline(time.Time{})
+
 	glog.V(LDEBUG).Infof("[udp-forward] %s <- %s\n%s", raddr, arg.Remote, rep)
 	if rep.Rep != gosocks5.Succeeded {
 		glog.V(LWARNING).Infof("[udp-forward] %s <- %s ASSOCIATE failured", raddr, arg.Remote)
@@ -103,25 +109,28 @@ func handleUdpForward(conn *net.UDPConn, raddr *net.UDPAddr, data []byte, arg Ar
 	dgram := gosocks5.NewUDPDatagram(
 		gosocks5.NewUDPHeader(uint16(len(data)), 0, ToSocksAddr(faddr)), data)
 
-	fconn.SetWriteDeadline(time.Now().Add(time.Second * 90))
-	if err = dgram.Write(fconn); err != nil {
+	tun.SetWriteDeadline(time.Now().Add(time.Second * 90))
+	if err = dgram.Write(tun); err != nil {
 		glog.V(LWARNING).Infof("[udp-forward] %s -> %s : %s", raddr, arg.Remote, err)
 		return
 	}
+	tun.SetWriteDeadline(time.Time{})
 	glog.V(LDEBUG).Infof("[udp-forward] %s >>> %s length %d", raddr, arg.Remote, len(data))
 
-	fconn.SetReadDeadline(time.Now().Add(time.Second * 90))
-	dgram, err = gosocks5.ReadUDPDatagram(fconn)
+	tun.SetReadDeadline(time.Now().Add(time.Second * 90))
+	dgram, err = gosocks5.ReadUDPDatagram(tun)
 	if err != nil {
 		glog.V(LWARNING).Infof("[udp-forward] %s <- %s : %s", raddr, arg.Remote, err)
 		return
 	}
+	tun.SetReadDeadline(time.Time{})
 	glog.V(LDEBUG).Infof("[udp-forward] %s <<< %s length %d", raddr, dgram.Header.Addr, len(dgram.Data))
 
 	if _, err = conn.WriteToUDP(dgram.Data, raddr); err != nil {
 		glog.V(LWARNING).Infof("[udp-forward] %s <- %s : %s", raddr, arg.Remote, err)
 	}
 
+	// NOTE: for now we only get one response from peer
 	glog.V(LINFO).Infof("[udp-forward] %s >-< %s", raddr, arg.Remote)
 }
 
@@ -137,11 +146,13 @@ func connectRTcpForward(conn net.Conn, arg Args) error {
 	}
 
 	// first reply, bind status
+	conn.SetReadDeadline(time.Now().Add(90 * time.Second))
 	rep, err := gosocks5.ReadReply(conn)
 	if err != nil {
 		glog.V(LWARNING).Infof("[rtcp] %s -> %s : %s", bindAddr, arg.Remote, err)
 		return err
 	}
+	conn.SetReadDeadline(time.Time{})
 	if rep.Rep != gosocks5.Succeeded {
 		glog.V(LWARNING).Infof("[rtcp] %s -> %s : bind on %s failure", bindAddr, arg.Remote, arg.Addr)
 		return errors.New("Bind on " + arg.Addr + " failure")
@@ -178,3 +189,79 @@ func connectRTcpForward(conn net.Conn, arg Args) error {
 
 	return nil
 }
+
+func connectRUdpForward(conn net.Conn, arg Args) error {
+	glog.V(LINFO).Infof("[rudp] %s - %s", arg.Addr, arg.Remote)
+
+	addr, _ := net.ResolveUDPAddr("udp", arg.Addr)
+	req := gosocks5.NewRequest(CmdUdpTun, ToSocksAddr(addr))
+	bindAddr := req.Addr
+	conn.SetWriteDeadline(time.Now().Add(time.Second * 90))
+	if err := req.Write(conn); err != nil {
+		glog.V(LWARNING).Infof("[rudp] %s -> %s : %s", bindAddr, arg.Remote, err)
+		return err
+	}
+	conn.SetWriteDeadline(time.Time{})
+
+	conn.SetReadDeadline(time.Now().Add(90 * time.Second))
+	rep, err := gosocks5.ReadReply(conn)
+	if err != nil {
+		glog.V(LWARNING).Infof("[rudp] %s <- %s : %s", bindAddr, arg.Remote, err)
+		return err
+	}
+	conn.SetReadDeadline(time.Time{})
+
+	if rep.Rep != gosocks5.Succeeded {
+		glog.V(LWARNING).Infof("[rudp] %s <- %s : bind on %s failure", bindAddr, arg.Remote, arg.Addr)
+		return errors.New(fmt.Sprintf("Bind on %s failure", bindAddr))
+	}
+
+	glog.V(LINFO).Infof("[rudp] %s - %s BIND ON %s OK", bindAddr, arg.Remote, rep.Addr)
+
+	raddr, err := net.ResolveUDPAddr("udp", arg.Remote)
+	if err != nil {
+		glog.V(LWARNING).Infof("[rudp] %s <- %s : %s", bindAddr, arg.Remote, err)
+		return err
+	}
+
+	for {
+		dgram, err := gosocks5.ReadUDPDatagram(conn)
+		if err != nil {
+			glog.V(LWARNING).Infof("[rudp] %s <- %s : %s", bindAddr, arg.Remote, err)
+			return err
+		}
+
+		go func() {
+			b := udpPool.Get().([]byte)
+			defer udpPool.Put(b)
+
+			relay, err := net.DialUDP("udp", nil, raddr)
+			if err != nil {
+				glog.V(LWARNING).Infof("[rudp] %s -> %s : %s", bindAddr, arg.Remote, err)
+				return
+			}
+			defer relay.Close()
+
+			if _, err := relay.Write(dgram.Data); err != nil {
+				glog.V(LWARNING).Infof("[rudp] %s -> %s : %s", bindAddr, arg.Remote, err)
+				return
+			}
+
+			relay.SetReadDeadline(time.Now().Add(time.Second * 60))
+			n, err := relay.Read(b)
+			if err != nil {
+				glog.V(LWARNING).Infof("[rudp] %s <- %s : %s", bindAddr, arg.Remote, err)
+				return
+			}
+			relay.SetReadDeadline(time.Time{})
+
+			conn.SetWriteDeadline(time.Now().Add(time.Second * 90))
+			if err := gosocks5.NewUDPDatagram(gosocks5.NewUDPHeader(uint16(n), 0, dgram.Header.Addr), b[:n]).Write(conn); err != nil {
+				glog.V(LWARNING).Infof("[rudp] %s <- %s : %s", bindAddr, arg.Remote, err)
+				return
+			}
+			conn.SetWriteDeadline(time.Time{})
+		}()
+	}
+
+}

socks.go

 package main
 
 import (
-	"bytes"
+	//"bytes"
 	"crypto/tls"
 	"errors"
 	"github.com/ginuerzh/gosocks5"
 	"github.com/golang/glog"
 	//"os/exec"
 	//"io"
-	"io/ioutil"
+	//"io/ioutil"
 	"net"
 	"net/url"
 	"strconv"
@@ -218,48 +218,14 @@ func handleSocks5Request(req *gosocks5.Request, conn net.Conn) {
 		Transport(conn, fconn)
 		glog.V(LINFO).Infof("[socks5-bind] %s >-< %s", conn.RemoteAddr(), fconn.RemoteAddr())
 
-		/*
-			case gosocks5.CmdUdp:
-			case CmdUdpTun:
-				glog.V(LINFO).Infof("[socks5-udp] %s - %s ASSOCIATE", conn.RemoteAddr(), req.Addr)
-				if len(forwardArgs) > 0 { // direct forward
-					fconn, _, err := forwardChain(forwardArgs...)
-					if err != nil {
-						glog.V(LWARNING).Infof("[socks5-udp] %s -> %s : %s", conn.RemoteAddr(), req.Addr, err)
-						rep := gosocks5.NewReply(gosocks5.Failure, nil)
-						if err := rep.Write(conn); err != nil {
-							glog.V(LWARNING).Infof("[socks5-udp] %s <- %s : %s", conn.RemoteAddr(), req.Addr, err)
-						} else {
-							glog.V(LDEBUG).Infof("[socks5-udp] %s <- %s\n%s", conn.RemoteAddr(), req.Addr, rep)
-						}
-						return
-					}
-					defer fconn.Close()
-					if err := req.Write(fconn); err != nil {
-						glog.V(LWARNING).Infof("[socks5-udp] %s -> %s : %s", conn.RemoteAddr(), req.Addr, err)
-						rep := gosocks5.NewReply(gosocks5.Failure, nil)
-						if err := rep.Write(conn); err != nil {
-							glog.V(LWARNING).Infof("[socks5-udp] %s <- %s : %s", conn.RemoteAddr(), req.Addr, err)
-						} else {
-							glog.V(LDEBUG).Infof("[socks5-udp] %s <- %s\n%s", conn.RemoteAddr(), req.Addr, rep)
-						}
-						return
-					}
-
-					glog.V(LINFO).Infof("[socks5-udp] %s <-> %s", conn.RemoteAddr(), req.Addr)
-					Transport(conn, fconn)
-					glog.V(LINFO).Infof("[socks5-udp] %s >-< %s", conn.RemoteAddr(), req.Addr)
-				} else {
-
-				}
-		*/
-	case gosocks5.CmdUdp, CmdUdpTun:
-		// TODO: udp tunnel <-> forward chain
-		glog.V(LINFO).Infof("[socks5-udp] %s - %s ASSOCIATE", conn.RemoteAddr(), req.Addr)
-		uconn, err := net.ListenUDP("udp", nil)
-		if err != nil {
-			glog.V(LWARNING).Infof("[socks5-udp] %s -> %s : %s", conn.RemoteAddr(), req.Addr, err)
+	case gosocks5.CmdUdp:
+		glog.V(LINFO).Infof("[socks5-udp] %s - %s", conn.RemoteAddr(), req.Addr)
+		socks5UDP(req, conn)
 
+	case CmdUdpTun:
+		glog.V(LINFO).Infof("[socks5-udp] %s - %s", conn.RemoteAddr(), req.Addr)
+		if err := socks5TunnelUDP(req, conn); err != nil {
+			glog.V(LWARNING).Infof("[socks5-udp] %s - %s : %s", conn.RemoteAddr(), req.Addr, err)
 			rep := gosocks5.NewReply(gosocks5.Failure, nil)
 			if err := rep.Write(conn); err != nil {
 				glog.V(LWARNING).Infof("[socks5-udp] %s <- %s : %s", conn.RemoteAddr(), req.Addr, err)
@@ -268,81 +234,122 @@ func handleSocks5Request(req *gosocks5.Request, conn net.Conn) {
 			}
 			return
 		}
-		defer uconn.Close()
 
-		addr := ToSocksAddr(uconn.LocalAddr())
-		addr.Host, _, _ = net.SplitHostPort(conn.LocalAddr().String()) // BUG: when server has multi-interfaces, this may cause a mistake
+	default:
+		glog.V(LWARNING).Infoln("[socks5] Unrecognized request:", req.Cmd)
+	}
+}
+
+func socks5UDP(req *gosocks5.Request, conn net.Conn) error {
+	bindAddr, _ := net.ResolveUDPAddr("udp", req.Addr.String())
+	relay, err := net.ListenUDP("udp", bindAddr) // udp associate, strict mode: if the port already in use, it will return error
+	if err != nil {
+		glog.V(LWARNING).Infof("[socks5-udp] %s -> %s : %s", conn.RemoteAddr(), req.Addr, err)
 
-		rep := gosocks5.NewReply(gosocks5.Succeeded, addr)
+		rep := gosocks5.NewReply(gosocks5.Failure, nil)
 		if err := rep.Write(conn); err != nil {
 			glog.V(LWARNING).Infof("[socks5-udp] %s <- %s : %s", conn.RemoteAddr(), req.Addr, err)
-			return
 		} else {
 			glog.V(LDEBUG).Infof("[socks5-udp] %s <- %s\n%s", conn.RemoteAddr(), req.Addr, rep)
-			glog.V(LINFO).Infof("[socks5-udp] %s -> %s LISTEN ON %s", conn.RemoteAddr(), req.Addr, addr)
 		}
+		return err
+	}
+	defer relay.Close()
 
-		var cc *UDPConn
-		var dgram *gosocks5.UDPDatagram
-		if req.Cmd == CmdUdpTun {
-			dgram, err = gosocks5.ReadUDPDatagram(conn)
-			if err != nil {
-				glog.V(LWARNING).Infof("[socks5-udp] %s -> %s : %s", conn.RemoteAddr(), req.Addr, err)
-				return
-			}
-			glog.V(LINFO).Infof("[socks5-udp] %s >>> %s, length %d", conn.RemoteAddr(), dgram.Header.Addr, len(dgram.Data))
-			cc = Client(conn, nil)
-		} else {
-			b := udpPool.Get().([]byte)
-			defer udpPool.Put(b)
+	addr := ToSocksAddr(relay.LocalAddr())
+	addr.Host, _, _ = net.SplitHostPort(conn.LocalAddr().String())
+	rep := gosocks5.NewReply(gosocks5.Succeeded, addr)
+	if err := rep.Write(conn); err != nil {
+		return err
+	}
+	glog.V(LDEBUG).Infof("[socks5-udp] %s <- %s\n%s", conn.RemoteAddr(), req.Addr, rep)
 
-			n, raddr, err := uconn.ReadFromUDP(b)
-			if err != nil {
-				glog.V(LWARNING).Infof("[socks5-udp] %s -> %s : %s", conn.RemoteAddr(), req.Addr, err)
-				return
-			}
-			dgram, err = gosocks5.ReadUDPDatagram(bytes.NewReader(b[:n]))
-			if err != nil {
-				glog.V(LWARNING).Infof("[socks5-udp] %s -> %s : %s", conn.RemoteAddr(), req.Addr, err)
-				return
-			}
-			glog.V(LINFO).Infof("[socks5-udp] %s >>> %s, length %d", raddr, dgram.Header.Addr, len(dgram.Data))
-			cc = Client(uconn, raddr)
+	glog.V(LINFO).Infof("[socks5-udp] %s - %s BIND ON %s OK", conn.RemoteAddr(), req.Addr, addr)
+
+	if len(forwardArgs) > 0 { // client -> tunnel, tunnel udp over tcp
+		tun, _, err := forwardChain(forwardArgs...)
+		if err != nil {
+			glog.V(LWARNING).Infof("[socks5-udp] %s -> %s : %s", conn.RemoteAddr(), req.Addr, err)
+			return err
 		}
+		defer tun.Close()
 
-		sc, err := createServerConn(uconn)
+		tun.SetWriteDeadline(time.Now().Add(time.Second * 90))
+		if err := gosocks5.NewRequest(CmdUdpTun, nil).Write(tun); err != nil {
+			glog.V(LWARNING).Infoln("[socks5-udp] %s -> %s : %s", conn.RemoteAddr(), req.Addr, err)
+			return err
+		}
+		tun.SetWriteDeadline(time.Time{})
+
+		tun.SetReadDeadline(time.Now().Add(time.Second * 90))
+		rep, err := gosocks5.ReadReply(tun)
 		if err != nil {
-			glog.V(LWARNING).Infof("[socks5-udp] %s", err)
-			return
+			glog.V(LWARNING).Infoln("[socks5-udp] %s -> %s : %s", conn.RemoteAddr(), req.Addr, err)
+			return err
+		}
+		if rep.Rep != gosocks5.Succeeded {
+			return errors.New("udp associate error")
 		}
-		defer sc.Close()
+		tun.SetReadDeadline(time.Time{})
 
-		if err = sc.WriteUDPTimeout(dgram, time.Second*90); err != nil {
-			glog.V(LWARNING).Infoln("socks5 udp:", err)
-			return
+		glog.V(LINFO).Infof("[socks5-udp] %s <-> %s", conn.RemoteAddr(), req.Addr)
+		go tunnelUDP(relay, tun, true)
+	} else { // standard socks5 udp relay
+		peer, err := net.ListenUDP("udp", nil)
+		if err != nil {
+			glog.V(LWARNING).Infof("[socks5-udp] %s -> %s : %s", conn.RemoteAddr(), req.Addr, err)
+			return err
 		}
-		dgram, err = sc.ReadUDPTimeout(time.Second * 90)
+		defer peer.Close()
+
+		glog.V(LINFO).Infof("[socks5-udp] %s <-> %s", conn.RemoteAddr(), req.Addr)
+		go transportUDP(relay, peer)
+	}
+
+	b := tcpPool.Get().([]byte)
+	defer tcpPool.Put(b)
+	for {
+		_, err := conn.Read(b) // discard any data from tcp connection
 		if err != nil {
-			glog.V(LWARNING).Infoln("socks5 udp:", err)
-			return
+			break // client disconnected
 		}
-		glog.V(LINFO).Infof("[socks5-udp] <<< %s, length %d", dgram.Header.Addr, len(dgram.Data))
+	}
+	glog.V(LINFO).Infof("[socks5-udp] %s >-< %s", conn.RemoteAddr(), req.Addr)
+	return nil
+}
 
-		if err = cc.WriteUDPTimeout(dgram, time.Second*90); err != nil {
-			glog.V(LWARNING).Infoln("socks5 udp:", err)
-			return
+func socks5TunnelUDP(req *gosocks5.Request, conn net.Conn) error {
+	if len(forwardArgs) > 0 { // tunnel -> tunnel, direct forward
+		tun, _, err := forwardChain(forwardArgs...)
+		if err != nil {
+			return err
 		}
+		defer tun.Close()
 
-		if req.Cmd == gosocks5.CmdUdp {
-			go TransportUDP(cc, sc)
-			ioutil.ReadAll(conn) // wait for client exit
-			glog.V(LINFO).Infoln("[udp] transfer done")
-		} else {
-			TransportUDP(cc, sc)
+		if err := req.Write(tun); err != nil {
+			return err
 		}
-	default:
-		glog.V(LWARNING).Infoln("[socks5] Unrecognized request:", req.Cmd)
+
+		glog.V(LINFO).Infof("[socks5-udp] %s <-> %s[tun]", conn.RemoteAddr(), tun.RemoteAddr())
+		Transport(conn, tun)
+		glog.V(LINFO).Infof("[socks5-udp] %s >-< %s[tun]", conn.RemoteAddr(), tun.RemoteAddr())
+	} else { // tunnel -> remote, handle tunnel udp request
+		bindAddr, _ := net.ResolveUDPAddr("udp", req.Addr.String())
+		uconn, err := net.ListenUDP("udp", bindAddr)
+		if err != nil {
+			return err
+		}
+		defer uconn.Close()
+
+		if err := gosocks5.NewReply(gosocks5.Succeeded, ToSocksAddr(uconn.LocalAddr())).Write(conn); err != nil {
+			return nil
+		}
+
+		glog.V(LINFO).Infof("[socks5-udp] %s <-> %s", conn.RemoteAddr(), uconn.LocalAddr())
+		tunnelUDP(uconn, conn, false)
+		glog.V(LINFO).Infof("[socks5-udp] %s >-< %s", conn.RemoteAddr(), uconn.LocalAddr())
 	}
+	return nil
 }
 
 func socks5Bind(req *gosocks5.Request, conn net.Conn) (*gosocks5.Reply, net.Conn, error) {
@@ -361,13 +368,13 @@ func socks5Bind(req *gosocks5.Request, conn net.Conn) (*gosocks5.Reply, net.Conn
 	}
 
 	bindAddr, _ := net.ResolveTCPAddr("tcp", req.Addr.String())
-	ln, err := net.ListenTCP("tcp", bindAddr)
+	ln, err := net.ListenTCP("tcp", bindAddr) // strict mode: if the port already in use, it will return error
 	if err != nil {
 		return gosocks5.NewReply(gosocks5.Failure, nil), nil, err
 	}
 
 	addr := ToSocksAddr(ln.Addr())
-	// Issue: may not reachable when host has two interfaces
+	// Issue: may not reachable when host has multi-interface
 	addr.Host, _, _ = net.SplitHostPort(conn.LocalAddr().String())
 	rep := gosocks5.NewReply(gosocks5.Succeeded, addr)
 	if err := rep.Write(conn); err != nil {
@@ -440,41 +447,6 @@ out:
 	return rep, pconn, nil
 }
 
-func createServerConn(uconn *net.UDPConn) (c *UDPConn, err error) {
-	if len(forwardArgs) == 0 {
-		c = Server(uconn)
-		return
-	}
-
-	fconn, _, err := forwardChain(forwardArgs...)
-	if err != nil {
-		return
-	}
-	glog.V(LINFO).Infoln("[udp] forward associate")
-
-	req := gosocks5.NewRequest(CmdUdpTun, nil)
-	if err = req.Write(fconn); err != nil {
-		fconn.Close()
-		return
-	}
-	glog.V(LDEBUG).Infoln(req)
-
-	rep, err := gosocks5.ReadReply(fconn)
-	if err != nil {
-		fconn.Close()
-		return
-	}
-	glog.V(LDEBUG).Infoln(rep)
-	if rep.Rep != gosocks5.Succeeded {
-		fconn.Close()
-		return nil, errors.New("Failure")
-	}
-	glog.V(LINFO).Infoln("[udp] forward associate on", rep.Addr, "OK")
-
-	c = Server(fconn)
-	return
-}
-
 func ToSocksAddr(addr net.Addr) *gosocks5.Addr {
 	host := "0.0.0.0"
 	port := 0
@@ -489,43 +461,3 @@ func ToSocksAddr(addr net.Addr) *gosocks5.Addr {
 		Port: uint16(port),
 	}
 }
-
-func PipeUDP(src, dst *UDPConn, ch chan<- error) {
-	var err error
-
-	for {
-		var dgram *gosocks5.UDPDatagram
-		dgram, err = src.ReadUDP()
-		if err != nil {
-			break
-		}
-		if src.isClient {
-			glog.V(LDEBUG).Infof("[udp] -> %s, length %d", dgram.Header.Addr, len(dgram.Data))
-		} else {
-			glog.V(LDEBUG).Infof("[udp] <- %s, length %d", dgram.Header.Addr, len(dgram.Data))
-		}
-		if err = dst.WriteUDP(dgram); err != nil {
-			break
-		}
-	}
-
-	ch <- err
-	close(ch)
-}
-
-func TransportUDP(cc, sc *UDPConn) (err error) {
-	rChan := make(chan error, 1)
-	wChan := make(chan error, 1)
-
-	go PipeUDP(cc, sc, wChan)
-	go PipeUDP(sc, cc, rChan)
-
-	select {
-	case err = <-wChan:
-		// glog.V(LDEBUG).Infoln("w exit", err)
-	case err = <-rChan:
-		// glog.V(LDEBUG).Infoln("r exit", err)
-	}
-
-	return
-}

udp.go

@@ -5,174 +5,166 @@ import (
 	"github.com/ginuerzh/gosocks5"
 	"github.com/golang/glog"
 	"net"
-	"time"
+	//"time"
 )
 
-type UDPConn struct {
-	isClient bool
-	udp      *net.UDPConn
-	addr     net.Addr
-	tcp      net.Conn
-}
-
-func Client(conn net.Conn, addr net.Addr) *UDPConn {
-	c := &UDPConn{isClient: true}
-
-	switch conn := conn.(type) {
-	case *net.UDPConn:
-		c.udp = conn
-		c.addr = addr
-	default:
-		c.tcp = conn
-	}
-
-	return c
-}
-
-func Server(conn net.Conn) *UDPConn {
-	c := &UDPConn{}
-
-	switch conn := conn.(type) {
-	case *net.UDPConn:
-		c.udp = conn
-	default:
-		c.tcp = conn
-	}
-
-	return c
-}
-
-func (c *UDPConn) ReadUDP() (*gosocks5.UDPDatagram, error) {
-	if c.isClient {
-		return c.readUDPClient()
-	}
-	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 transportUDP(relay, peer *net.UDPConn) (err error) {
+	rChan := make(chan error, 1)
+	wChan := make(chan error, 1)
 
-func (c *UDPConn) readUDPClient() (*gosocks5.UDPDatagram, error) {
-	if c.udp != nil {
-		return gosocks5.ReadUDPDatagram(c.udp)
-	}
-	return gosocks5.ReadUDPDatagram(c.tcp)
-}
+	var clientAddr *net.UDPAddr
 
-func (c *UDPConn) readUDPServer() (*gosocks5.UDPDatagram, error) {
-	if c.udp != nil {
-		// b := make([]byte, 65535)
+	go func() {
 		b := udpPool.Get().([]byte)
 		defer udpPool.Put(b)
 
-		n, addr, err := c.udp.ReadFrom(b)
-		if err != nil {
-			return nil, err
+		for {
+			n, laddr, err := relay.ReadFromUDP(b)
+			if err != nil {
+				rChan <- err
+				return
+			}
+			if clientAddr == nil {
+				clientAddr = laddr
+			}
+			dgram, err := gosocks5.ReadUDPDatagram(bytes.NewReader(b[:n]))
+			if err != nil {
+				rChan <- err
+				return
+			}
+
+			raddr, err := net.ResolveUDPAddr("udp", dgram.Header.Addr.String())
+			if err != nil {
+				continue // drop silently
+			}
+			if _, err := peer.WriteToUDP(dgram.Data, raddr); err != nil {
+				rChan <- err
+				return
+			}
+			glog.V(LDEBUG).Infof("[socks5-udp] %s >>> %s length: %d", relay.LocalAddr(), raddr, len(dgram.Data))
 		}
-		dgram := gosocks5.NewUDPDatagram(
-			gosocks5.NewUDPHeader(0, 0, ToSocksAddr(addr)), b[:n])
-		return dgram, nil
-	}
-	return gosocks5.ReadUDPDatagram(c.tcp)
-}
+	}()
 
-func (c *UDPConn) WriteUDP(dgram *gosocks5.UDPDatagram) error {
-	if c.isClient {
-		return c.writeUDPClient(dgram)
-	}
-	return c.writeUDPServer(dgram)
-}
+	go func() {
+		b := udpPool.Get().([]byte)
+		defer udpPool.Put(b)
 
-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)
-}
+		for {
+			n, raddr, err := peer.ReadFrom(b)
+			if err != nil {
+				wChan <- err
+				return
+			}
+			if clientAddr == nil {
+				continue
+			}
+			buf := bytes.Buffer{}
+			dgram := gosocks5.NewUDPDatagram(gosocks5.NewUDPHeader(0, 0, ToSocksAddr(raddr)), b[:n])
+			dgram.Write(&buf)
+			if _, err := relay.WriteToUDP(buf.Bytes(), clientAddr); err != nil {
+				wChan <- err
+				return
+			}
+			glog.V(LDEBUG).Infof("[socks5-udp] %s <<< %s length: %d", relay.LocalAddr(), raddr, len(dgram.Data))
+		}
+	}()
 
-func (c *UDPConn) writeUDPClient(dgram *gosocks5.UDPDatagram) error {
-	if c.udp != nil {
-		dgram.Header.Rsv = 0
-		buffer := bytes.Buffer{}
-		dgram.Write(&buffer)
-		_, err := c.udp.WriteTo(buffer.Bytes(), c.addr)
-		return err
+	select {
+	case err = <-wChan:
+		//log.Println("w exit", err)
+	case err = <-rChan:
+		//log.Println("r exit", err)
 	}
 
-	dgram.Header.Rsv = uint16(len(dgram.Data))
-	return dgram.Write(c.tcp)
+	return
 }
 
-func (c *UDPConn) writeUDPServer(dgram *gosocks5.UDPDatagram) error {
-	if c.udp != nil {
-		addr, err := net.ResolveUDPAddr("udp", dgram.Header.Addr.String())
-		if err != nil {
-			glog.V(LWARNING).Infoln(err)
-			return nil // drop silently
-		}
-		_, err = c.udp.WriteTo(dgram.Data, addr)
-		return err
-	}
-	dgram.Header.Rsv = uint16(len(dgram.Data))
-	return dgram.Write(c.tcp)
-}
-
-func (c *UDPConn) Close() error {
-	if c.udp != nil {
-		return c.udp.Close()
-	}
-	return c.tcp.Close()
-}
+func tunnelUDP(conn *net.UDPConn, tun net.Conn, client bool) (err error) {
+	rChan := make(chan error, 1)
+	wChan := make(chan error, 1)
 
-func (c *UDPConn) LocalAddr() net.Addr {
-	if c.udp != nil {
-		return c.udp.LocalAddr()
-	}
-	return c.tcp.LocalAddr()
-}
+	var clientAddr *net.UDPAddr
 
-func (c *UDPConn) RemoteAddr() net.Addr {
-	if c.udp != nil {
-		return c.udp.RemoteAddr()
-	}
-	return c.tcp.RemoteAddr()
-}
+	go func() {
+		b := udpPool.Get().([]byte)
+		defer udpPool.Put(b)
 
-func (c *UDPConn) SetDeadline(t time.Time) error {
-	if c.udp != nil {
-		return c.udp.SetDeadline(t)
-	}
-	return c.tcp.SetDeadline(t)
-}
+		for {
+			n, addr, err := conn.ReadFromUDP(b)
+			if err != nil {
+				rChan <- err
+				return
+			}
+
+			var dgram *gosocks5.UDPDatagram
+			if client { // pipe from relay to tunnel
+				dgram, err = gosocks5.ReadUDPDatagram(bytes.NewReader(b[:n]))
+				if err != nil {
+					rChan <- err
+					return
+				}
+				if clientAddr == nil {
+					clientAddr = addr
+				}
+				dgram.Header.Rsv = uint16(len(dgram.Data))
+				if err := dgram.Write(tun); err != nil {
+					rChan <- err
+					return
+				}
+				glog.V(LDEBUG).Infof("[socks5-udp] %s >>> %s length: %d", conn.LocalAddr(), dgram.Header.Addr, len(dgram.Data))
+			} else { // pipe from peer to tunnel
+				dgram = gosocks5.NewUDPDatagram(
+					gosocks5.NewUDPHeader(uint16(n), 0, ToSocksAddr(addr)), b[:n])
+				if err := dgram.Write(tun); err != nil {
+					rChan <- err
+					return
+				}
+				glog.V(LDEBUG).Infof("[socks5-udp] %s <<< %s length: %d", tun.RemoteAddr(), dgram.Header.Addr, len(dgram.Data))
+			}
+		}
+	}()
+
+	go func() {
+		for {
+			dgram, err := gosocks5.ReadUDPDatagram(tun)
+			if err != nil {
+				wChan <- err
+				return
+			}
+
+			if client { // pipe from tunnel to relay
+				if clientAddr == nil {
+					continue
+				}
+				dgram.Header.Rsv = 0
+
+				buf := bytes.Buffer{}
+				dgram.Write(&buf)
+				if _, err := conn.WriteToUDP(buf.Bytes(), clientAddr); err != nil {
+					wChan <- err
+					return
+				}
+				glog.V(LDEBUG).Infof("[socks5-udp] %s <<< %s length: %d", conn.LocalAddr(), dgram.Header.Addr, len(dgram.Data))
+			} else { // pipe from tunnel to peer
+				addr, err := net.ResolveUDPAddr("udp", dgram.Header.Addr.String())
+				if err != nil {
+					continue // drop silently
+				}
+				if _, err := conn.WriteToUDP(dgram.Data, addr); err != nil {
+					wChan <- err
+					return
+				}
+				glog.V(LDEBUG).Infof("[socks5-udp] %s >>> %s length: %d", tun.RemoteAddr(), addr, len(dgram.Data))
+			}
+		}
+	}()
 
-func (c *UDPConn) SetReadDeadline(t time.Time) error {
-	if c.udp != nil {
-		return c.udp.SetReadDeadline(t)
+	select {
+	case err = <-wChan:
+		//log.Println("w exit", err)
+	case err = <-rChan:
+		//log.Println("r exit", err)
 	}
-	return c.tcp.SetReadDeadline(t)
-}
 
-func (c *UDPConn) SetWriteDeadline(t time.Time) error {
-	if c.udp != nil {
-		return c.udp.SetWriteDeadline(t)
-	}
-	return c.tcp.SetWriteDeadline(t)
+	return
 }

util.go

@@ -91,13 +91,15 @@ func parseArgs(ss []string) (args []Args) {
 
 // Based on io.Copy, but the io.ErrShortWrite is ignored (mainly for websocket)
 func Copy(dst io.Writer, src io.Reader) (written int64, err error) {
-	buf := make([]byte, 32*1024)
+	// b := make([]byte, 32*1024)
+	b := tcpPool.Get().([]byte)
+	defer tcpPool.Put(b)
 
 	for {
-		nr, er := src.Read(buf)
+		nr, er := src.Read(b)
 		//log.Println("cp r", nr, er)
 		if nr > 0 {
-			nw, ew := dst.Write(buf[:nr])
+			nw, ew := dst.Write(b[:nr])
 			//log.Println("cp w", nw, ew)
 			if nw > 0 {
 				written += int64(nw)