From 601e0ddcdaef389c96de88cab0a67ce3551ad4dd Mon Sep 17 00:00:00 2001
From: nadoo <287492+nadoo@users.noreply.github.com>
Date: Sat, 20 Jan 2018 15:36:53 +0800
Subject: [PATCH] udptun: add nat map

---
 ss.go     |  6 +++---
 udptun.go | 48 +++++++++++++++++++++++++++---------------------
 uottun.go |  3 +--
 3 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/ss.go b/ss.go
index c045d3b..d85bd33 100644
--- a/ss.go
+++ b/ss.go
@@ -163,14 +163,12 @@ func (s *SS) ListenAndServeUDP() {
 			continue
 		}
 
-		logf("proxy-ss-udp %s <-> %s", raddr, c.tgtAddr)
-
 		var pc *PktConn
 		v, ok := nm.Load(raddr.String())
 		if !ok && v == nil {
 			lpc, nextHop, err := s.sDialer.DialUDP("udp", c.tgtAddr.String())
 			if err != nil {
-				logf("proxy-ss-udp remote listen error: %v", err)
+				logf("proxy-ss-udp remote dial error: %v", err)
 				continue
 			}
 
@@ -192,6 +190,8 @@ func (s *SS) ListenAndServeUDP() {
 			logf("proxy-ss-udp remote write error: %v", err)
 			continue
 		}
+
+		logf("proxy-ss-udp %s <-> %s", raddr, c.tgtAddr)
 	}
 }
 
diff --git a/udptun.go b/udptun.go
index 3ffcb1c..7c14cbf 100644
--- a/udptun.go
+++ b/udptun.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"net"
+	"sync"
 	"time"
 )
 
@@ -35,41 +36,46 @@ func (s *UDPTun) ListenAndServe() {
 
 	logf("proxy-udptun listening UDP on %s", s.addr)
 
+	var nm sync.Map
 	buf := make([]byte, udpBufSize)
 
 	for {
-		n, clientAddr, err := c.ReadFrom(buf)
+		n, raddr, err := c.ReadFrom(buf)
 		if err != nil {
 			logf("proxy-udptun read error: %v", err)
 			continue
 		}
 
-		go func() {
-			rc, wt, err := s.sDialer.DialUDP("udp", s.raddr)
+		var pc net.PacketConn
+		var writeAddr net.Addr
+
+		v, ok := nm.Load(raddr.String())
+		if !ok && v == nil {
+
+			pc, writeAddr, err = s.sDialer.DialUDP("udp", s.raddr)
 			if err != nil {
-				logf("proxy-udptun failed to connect to server %v: %v", s.raddr, err)
-				return
+				logf("proxy-udptun remote dial error: %v", err)
+				continue
 			}
 
-			n, err = rc.WriteTo(buf[:n], wt)
-			if err != nil {
-				logf("proxy-udptun rc.Write error: %v", err)
-				return
-			}
+			nm.Store(raddr.String(), pc)
+			go func() {
+				timedCopy(c, raddr, pc, 2*time.Minute)
+				pc.Close()
+				nm.Delete(raddr.String())
+			}()
 
-			rcBuf := make([]byte, udpBufSize)
-			rc.SetReadDeadline(time.Now().Add(time.Minute))
+		} else {
+			pc = v.(net.PacketConn)
+		}
 
-			n, _, err = rc.ReadFrom(rcBuf)
-			if err != nil {
-				logf("proxy-udptun rc.Read error: %v", err)
-				return
-			}
-			rc.Close()
+		_, err = pc.WriteTo(buf[:n], writeAddr)
+		if err != nil {
+			logf("proxy-udptun remote write error: %v", err)
+			continue
+		}
 
-			c.WriteTo(rcBuf[:n], clientAddr)
-			logf("proxy-udptun %s <-> %s", clientAddr, s.raddr)
-		}()
+		logf("proxy-udptun %s <-> %s", raddr, s.raddr)
 
 	}
 }
diff --git a/uottun.go b/uottun.go
index 5e910d6..7b99c9a 100644
--- a/uottun.go
+++ b/uottun.go
@@ -51,8 +51,6 @@ func (s *UoTTun) ListenAndServe() {
 			continue
 		}
 
-		rc.Write(buf[:n])
-
 		go func() {
 			// no remote forwarder, just a local udp forwarder
 			if urc, ok := rc.(*net.UDPConn); ok {
@@ -71,6 +69,7 @@ func (s *UoTTun) ListenAndServe() {
 			c.WriteTo(resp, clientAddr)
 		}()
 
+		rc.Write(buf[:n])
 		logf("proxy-uottun %s <-> %s", clientAddr, s.raddr)
 	}
 }