From 31930f2955cb1a054032f6ec3dfdb38e8722110b Mon Sep 17 00:00:00 2001
From: nadoo <287492+nadoo@users.noreply.github.com>
Date: Wed, 17 Jan 2018 00:26:38 +0800
Subject: [PATCH] udptun: add DialUDP to forwarders

---
 dialer.go               |  3 ++
 direct.go               | 21 +++++++++-
 http.go                 |  7 +++-
 rule.go                 |  4 ++
 socks5.go               |  5 +++
 ss.go                   | 85 +++++++++++++++++++++++++++--------------
 strategy.go             | 15 ++++++++
 systemd/glider@.service |  1 -
 udptun.go               |  6 +--
 9 files changed, 113 insertions(+), 34 deletions(-)

diff --git a/dialer.go b/dialer.go
index 50733b3..c20e7da 100644
--- a/dialer.go
+++ b/dialer.go
@@ -14,6 +14,9 @@ type Dialer interface {
 	// Dial connects to the given address via the proxy.
 	Dial(network, addr string) (c net.Conn, err error)
 
+	// DialUDP connects to the given address via the proxy.
+	DialUDP(network, addr string) (pc net.PacketConn, writeTo net.Addr, err error)
+
 	// Get the dialer by dstAddr
 	NextDialer(dstAddr string) Dialer
 }
diff --git a/direct.go b/direct.go
index 33013e7..084a3e7 100644
--- a/direct.go
+++ b/direct.go
@@ -1,6 +1,8 @@
 package main
 
-import "net"
+import (
+	"net"
+)
 
 // direct proxy
 type direct struct{}
@@ -16,10 +18,27 @@ func (d *direct) Dial(network, addr string) (net.Conn, error) {
 	}
 
 	c, err := net.Dial(network, addr)
+	if err != nil {
+		return nil, err
+	}
+
 	if c, ok := c.(*net.TCPConn); ok {
 		c.SetKeepAlive(true)
 	}
+
 	return c, err
 }
 
+// DialUDP connects to the given address via the proxy.
+func (d *direct) DialUDP(network, addr string) (pc net.PacketConn, writeTo net.Addr, err error) {
+	pc, err = net.ListenPacket(network, "")
+	if err != nil {
+		logf("ListenPacket error: %s", err)
+		return nil, nil, err
+	}
+
+	uAddr, err := net.ResolveUDPAddr("udp", addr)
+	return pc, uAddr, nil
+}
+
 func (d *direct) NextDialer(dstAddr string) Dialer { return d }
diff --git a/http.go b/http.go
index 10d259a..218efbb 100644
--- a/http.go
+++ b/http.go
@@ -212,7 +212,7 @@ func (s *HTTP) servHTTPS(method, requestURI, proto string, c net.Conn) {
 
 // Dial connects to the address addr on the network net via the proxy.
 func (s *HTTP) Dial(network, addr string) (net.Conn, error) {
-	rc, err := s.cDialer.Dial("tcp", s.addr)
+	rc, err := s.cDialer.Dial(network, s.addr)
 	if err != nil {
 		logf("dial to %s error: %s", s.addr, err)
 		return nil, err
@@ -247,6 +247,11 @@ func (s *HTTP) Dial(network, addr string) (net.Conn, error) {
 	return nil, errors.New("proxy-http cound not connect remote address: " + addr + ". error code: " + code)
 }
 
+// DialUDP connects to the given address via the proxy.
+func (s *HTTP) DialUDP(network, addr string) (pc net.PacketConn, writeTo net.Addr, err error) {
+	return nil, nil, errors.New("DialUDP not supported")
+}
+
 // parseFirstLine parses "GET /foo HTTP/1.1" OR "HTTP/1.1 200 OK" into its three parts.
 func parseFirstLine(tp *textproto.Reader) (r1, r2, r3 string, ok bool) {
 	line, err := tp.ReadLine()
diff --git a/rule.go b/rule.go
index 3102909..5e04d8a 100644
--- a/rule.go
+++ b/rule.go
@@ -113,6 +113,10 @@ func (rd *RuleDialer) Dial(network, addr string) (net.Conn, error) {
 	return rd.NextDialer(addr).Dial(network, addr)
 }
 
+func (rd *RuleDialer) DialUDP(network, addr string) (pc net.PacketConn, writeTo net.Addr, err error) {
+	return rd.NextDialer(addr).DialUDP(network, addr)
+}
+
 // AddDomainIP used to update ipMap rules according to domainMap rule
 func (rd *RuleDialer) AddDomainIP(domain, ip string) error {
 	if ip != "" {
diff --git a/socks5.go b/socks5.go
index 3e686fd..e025144 100644
--- a/socks5.go
+++ b/socks5.go
@@ -156,6 +156,11 @@ func (s *SOCKS5) Dial(network, addr string) (net.Conn, error) {
 	return c, nil
 }
 
+// DialUDP connects to the given address via the proxy.
+func (s *SOCKS5) DialUDP(network, addr string) (pc net.PacketConn, writeTo net.Addr, err error) {
+	return nil, nil, errors.New("DialUDP not supported")
+}
+
 // connect takes an existing connection to a socks5 proxy server,
 // and commands the server to extend that connection to target,
 // which must be a canonical address with a host and port.
diff --git a/ss.go b/ss.go
index 3674b87..afe9b03 100644
--- a/ss.go
+++ b/ss.go
@@ -225,8 +225,8 @@ func (s *SS) Dial(network, addr string) (net.Conn, error) {
 	case "uot":
 		target[0] = target[0] | 0x8
 		return s.dialTCP(target)
-	case "udp":
-		return s.dialUDP(target)
+	// case "udp":
+	// 	return s.dialUDP(target)
 	default:
 		return nil, errors.New("Unknown schema: " + network)
 	}
@@ -254,34 +254,51 @@ func (s *SS) dialTCP(target Addr) (net.Conn, error) {
 	return c, err
 }
 
-// TODO: support forwarder chain
-func (s *SS) dialUDP(target Addr) (net.Conn, error) {
-	c, err := net.ListenPacket("udp", "")
+// DialUDP connects to the given address via the proxy.
+func (s *SS) DialUDP(network, addr string) (pc net.PacketConn, writeTo net.Addr, err error) {
+	// TODO: check forward chain
+	pc, nextHop, err := s.cDialer.DialUDP(network, addr)
 	if err != nil {
-		logf("proxy-ss dialudp failed to listen packet: %v", err)
-		return nil, err
+		logf("proxy-ss dialudp to %s error: %s", addr, err)
+		return nil, nil, err
 	}
 
-	sUDPAddr, err := net.ResolveUDPAddr("udp", s.Addr())
-	suc := &ssUDPConn{
-		PacketConn: s.PacketConn(c),
-		addr:       sUDPAddr,
-		target:     target,
-	}
+	nextHopAddr := ParseAddr(nextHop.String())
+	writeTo, err = net.ResolveUDPAddr("udp", s.Addr())
 
-	return suc, err
+	pkc := NewPktConn(s.PacketConn(pc), writeTo, nextHopAddr, true)
+
+	return pkc, writeTo, err
 }
 
-type ssUDPConn struct {
+// PktConn wraps a net.PacketConn and support Write method like net.Conn
+type PktConn struct {
 	net.PacketConn
 
-	addr   net.Addr
-	target Addr
+	addr      net.Addr // write to and read from addr
+	target    Addr
+	tgtHeader bool
 }
 
-func (uc *ssUDPConn) Read(b []byte) (int, error) {
+// NewPktConn returns a PktConn
+func NewPktConn(c net.PacketConn, addr net.Addr, target Addr, tgtHeader bool) *PktConn {
+	pc := &PktConn{
+		PacketConn: c,
+		addr:       addr,
+		target:     target,
+		tgtHeader:  tgtHeader}
+	return pc
+}
+
+func (pc *PktConn) Read(b []byte) (int, error) {
+
+	if !pc.tgtHeader {
+		n, _, err := pc.PacketConn.ReadFrom(b)
+		return n, err
+	}
+
 	buf := make([]byte, len(b))
-	n, raddr, err := uc.PacketConn.ReadFrom(buf)
+	n, raddr, err := pc.PacketConn.ReadFrom(buf)
 	if err != nil {
 		return 0, err
 	}
@@ -292,20 +309,32 @@ func (uc *ssUDPConn) Read(b []byte) (int, error) {
 	return n - len(srcAddr), err
 }
 
-func (uc *ssUDPConn) Write(b []byte) (int, error) {
-	buf := make([]byte, len(uc.target)+len(b))
-	copy(buf, uc.target)
-	copy(buf[len(uc.target):], b)
+func (pc *PktConn) ReadFrom(b []byte) (int, net.Addr, error) {
+
+	n, err := pc.Read(b)
+
+	// TODO: Addr
+	return n, nil, err
+}
+
+func (pc *PktConn) Write(b []byte) (int, error) {
+	if !pc.tgtHeader {
+		return pc.PacketConn.WriteTo(b, pc.addr)
+	}
+
+	buf := make([]byte, len(pc.target)+len(b))
+	copy(buf, pc.target)
+	copy(buf[len(pc.target):], b)
 
 	// logf("Write: \n%s", hex.Dump(buf))
 
-	return uc.PacketConn.WriteTo(buf, uc.addr)
+	return pc.PacketConn.WriteTo(buf, pc.addr)
 }
 
-func (uc *ssUDPConn) WriteTo(b []byte, addr net.Addr) (int, error) {
-	return 0, errors.New("not available")
+func (pc *PktConn) WriteTo(b []byte, addr net.Addr) (int, error) {
+	return pc.Write(b)
 }
 
-func (uc *ssUDPConn) RemoteAddr() net.Addr {
-	return uc.addr
+func (pc *PktConn) RemoteAddr() net.Addr {
+	return pc.addr
 }
diff --git a/strategy.go b/strategy.go
index 1af69d0..d87e841 100644
--- a/strategy.go
+++ b/strategy.go
@@ -67,6 +67,10 @@ func (rr *rrDialer) Dial(network, addr string) (net.Conn, error) {
 	return rr.NextDialer(addr).Dial(network, addr)
 }
 
+func (rr *rrDialer) DialUDP(network, addr string) (pc net.PacketConn, writeTo net.Addr, err error) {
+	return rr.NextDialer(addr).DialUDP(network, addr)
+}
+
 func (rr *rrDialer) NextDialer(dstAddr string) Dialer {
 	n := len(rr.dialers)
 	if n == 1 {
@@ -158,3 +162,14 @@ func (ha *haDialer) Dial(network, addr string) (net.Conn, error) {
 
 	return d.Dial(network, addr)
 }
+
+func (ha *haDialer) DialUDP(network, addr string) (pc net.PacketConn, writeTo net.Addr, err error) {
+	d := ha.dialers[ha.idx]
+
+	result, ok := ha.status.Load(ha.idx)
+	if ok && !result.(bool) {
+		d = ha.NextDialer(addr)
+	}
+
+	return d.DialUDP(network, addr)
+}
diff --git a/systemd/glider@.service b/systemd/glider@.service
index 2b7d464..2891ffe 100644
--- a/systemd/glider@.service
+++ b/systemd/glider@.service
@@ -1,6 +1,5 @@
 [Unit]
 Description=Glider Service (%i)
-After=network.target
 Before=iptables.service ip6tables.service
 
 [Service]
diff --git a/udptun.go b/udptun.go
index 46990b8..d6f3f34 100644
--- a/udptun.go
+++ b/udptun.go
@@ -42,20 +42,20 @@ func (s *UDPTun) ListenAndServe() {
 			continue
 		}
 
-		rc, err := s.sDialer.Dial("udp", s.raddr)
+		rc, wt, err := s.sDialer.DialUDP("udp", s.raddr)
 		if err != nil {
 			logf("proxy-udptun failed to connect to server %v: %v", s.raddr, err)
 			continue
 		}
 
-		n, err = rc.Write(buf[:n])
+		n, err = rc.WriteTo(buf[:n], wt)
 		if err != nil {
 			logf("proxy-udptun rc.Write error: %v", err)
 			continue
 		}
 
 		buf = make([]byte, udpBufSize)
-		n, err = rc.Read(buf)
+		n, _, err = rc.ReadFrom(buf)
 		if err != nil {
 			logf("proxy-udptun rc.Read error: %v", err)
 			continue