From 9ac6db498e3e165c2b492e969394577d4539bdf4 Mon Sep 17 00:00:00 2001
From: nadoo <287492+nadoo@users.noreply.github.com>
Date: Wed, 17 Jan 2018 22:44:04 +0800
Subject: [PATCH] udptun: optimize code

---
 udptun.go | 55 +++++++++++++++++++++++++++++++++++--------------------
 1 file changed, 35 insertions(+), 20 deletions(-)

diff --git a/udptun.go b/udptun.go
index d6f3f34..aafc1ff 100644
--- a/udptun.go
+++ b/udptun.go
@@ -1,7 +1,9 @@
 package main
 
 import (
+	"encoding/hex"
 	"net"
+	"time"
 )
 
 // UDPTun struct
@@ -35,6 +37,7 @@ func (s *UDPTun) ListenAndServe() {
 	logf("proxy-udptun listening UDP on %s", s.addr)
 
 	buf := make([]byte, udpBufSize)
+
 	for {
 		n, clientAddr, err := c.ReadFrom(buf)
 		if err != nil {
@@ -42,29 +45,41 @@ func (s *UDPTun) ListenAndServe() {
 			continue
 		}
 
-		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
-		}
+		logf("ClientAddr: %s", clientAddr)
+		logf("ReadFrom:\n%s", hex.Dump(buf[:n]))
 
-		n, err = rc.WriteTo(buf[:n], wt)
-		if err != nil {
-			logf("proxy-udptun rc.Write error: %v", err)
-			continue
-		}
+		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
+			}
 
-		buf = make([]byte, udpBufSize)
-		n, _, err = rc.ReadFrom(buf)
-		if err != nil {
-			logf("proxy-udptun rc.Read error: %v", err)
-			continue
-		}
-		rc.Close()
+			logf("WriteTo")
+			n, err = rc.WriteTo(buf[:n], wt)
+			if err != nil {
+				logf("proxy-udptun rc.Write error: %v", err)
+				return
+			}
 
-		// logf("rc resp: \n%s", hex.Dump(buf[:n]))
+			logf("ReadFrom")
+
+			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()
+
+			// logf("rc resp: \n%s", hex.Dump(buf[:n]))
+
+			logf("c.WriteTo")
+			c.WriteTo(rcBuf[:n], clientAddr)
+			logf("proxy-udptun %s <-> %s", clientAddr, s.raddr)
+		}()
 
-		c.WriteTo(buf[:n], clientAddr)
-		logf("proxy-udptun %s <-> %s", clientAddr, s.raddr)
 	}
 }