udptun: add nat map

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)
 	}
 }
 

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
 		}
 
+		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 remote dial error: %v", err)
+				continue
+			}
+
+			nm.Store(raddr.String(), pc)
 			go func() {
-			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)
-				return
-			}
-
-			n, err = rc.WriteTo(buf[:n], wt)
-			if err != nil {
-				logf("proxy-udptun rc.Write error: %v", err)
-				return
-			}
-
-			rcBuf := make([]byte, udpBufSize)
-			rc.SetReadDeadline(time.Now().Add(time.Minute))
-
-			n, _, err = rc.ReadFrom(rcBuf)
-			if err != nil {
-				logf("proxy-udptun rc.Read error: %v", err)
-				return
-			}
-			rc.Close()
-
-			c.WriteTo(rcBuf[:n], clientAddr)
-			logf("proxy-udptun %s <-> %s", clientAddr, s.raddr)
+				timedCopy(c, raddr, pc, 2*time.Minute)
+				pc.Close()
+				nm.Delete(raddr.String())
 			}()
 
+		} else {
+			pc = v.(net.PacketConn)
+		}
+
+		_, err = pc.WriteTo(buf[:n], writeAddr)
+		if err != nil {
+			logf("proxy-udptun remote write error: %v", err)
+			continue
+		}
+
+		logf("proxy-udptun %s <-> %s", raddr, s.raddr)
+
 	}
 }

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)
 	}
 }