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{})
+		}()
+	}
+
+}

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)