From ad123eb093073b72dbe45c8ab1486e29e5db0174 Mon Sep 17 00:00:00 2001
From: netbyte <alex.tsai@protonmail.com>
Date: Wed, 13 Apr 2022 10:11:24 +0800
Subject: [PATCH] fix

---
 tcp/tcpclient.go | 46 +++++++++++++++++++++++++---------------------
 ws/wsclient.go   | 43 +++++++++++++++++++++++--------------------
 2 files changed, 48 insertions(+), 41 deletions(-)

diff --git a/tcp/tcpclient.go b/tcp/tcpclient.go
index 116a297..d63a0a0 100644
--- a/tcp/tcpclient.go
+++ b/tcp/tcpclient.go
@@ -4,9 +4,9 @@ import (
 	"io"
 	"log"
 	"net"
-	"sync"
 	"time"
 
+	"github.com/net-byte/vtun/common/cache"
 	"github.com/net-byte/vtun/common/cipher"
 	"github.com/net-byte/vtun/common/config"
 	"github.com/net-byte/vtun/tun"
@@ -17,39 +17,40 @@ import (
 func StartClient(config config.Config) {
 	log.Printf("vtun tcp client started on %v", config.LocalAddr)
 	iface := tun.CreateTun(config)
+	go tunToTcp(config, iface)
 	for {
 		if conn, err := net.DialTimeout("tcp", config.ServerAddr, time.Duration(config.Timeout)*time.Second); conn != nil && err == nil {
-			var wg sync.WaitGroup
-			wg.Add(2)
-			go tcpToTun(&wg, config, conn, iface)
-			go tunToTcp(&wg, config, conn, iface)
-			wg.Wait()
-			conn.Close()
+			cache.GetCache().Set("tcpconn", conn, 24*time.Hour)
+			tcpToTun(config, conn, iface)
+			cache.GetCache().Delete("tcpconn")
 		}
 	}
 }
 
-func tunToTcp(wg *sync.WaitGroup, config config.Config, tcpconn net.Conn, iface *water.Interface) {
-	defer wg.Done()
+func tunToTcp(config config.Config, iface *water.Interface) {
 	packet := make([]byte, config.MTU)
 	for {
 		n, err := iface.Read(packet)
 		if err != nil || n == 0 {
-			break
+			continue
 		}
-		b := packet[:n]
-		if config.Obfs {
-			b = cipher.XOR(b)
-		}
-		tcpconn.SetWriteDeadline(time.Now().Add(time.Duration(config.Timeout) * time.Second))
-		_, err = tcpconn.Write(b)
-		if err != nil {
-			break
+		if v, ok := cache.GetCache().Get("tcpconn"); ok {
+			b := packet[:n]
+			if config.Obfs {
+				packet = cipher.XOR(packet)
+			}
+			tcpconn := v.(net.Conn)
+			tcpconn.SetWriteDeadline(time.Now().Add(time.Duration(config.Timeout) * time.Second))
+			_, err = tcpconn.Write(b)
+			if err != nil {
+				continue
+			}
 		}
 	}
 }
-func tcpToTun(wg *sync.WaitGroup, config config.Config, tcpconn net.Conn, iface *water.Interface) {
-	defer wg.Done()
+
+func tcpToTun(config config.Config, tcpconn net.Conn, iface *water.Interface) {
+	defer tcpconn.Close()
 	packet := make([]byte, config.MTU)
 	for {
 		tcpconn.SetReadDeadline(time.Now().Add(time.Duration(config.Timeout) * time.Second))
@@ -61,6 +62,9 @@ func tcpToTun(wg *sync.WaitGroup, config config.Config, tcpconn net.Conn, iface
 		if config.Obfs {
 			b = cipher.XOR(b)
 		}
-		iface.Write(b)
+		_, err = iface.Write(b)
+		if err != nil {
+			break
+		}
 	}
 }
diff --git a/ws/wsclient.go b/ws/wsclient.go
index d3c36d9..c3278fe 100644
--- a/ws/wsclient.go
+++ b/ws/wsclient.go
@@ -3,10 +3,10 @@ package ws
 import (
 	"log"
 	"net"
-	"sync"
 	"time"
 
 	"github.com/gobwas/ws/wsutil"
+	"github.com/net-byte/vtun/common/cache"
 	"github.com/net-byte/vtun/common/cipher"
 	"github.com/net-byte/vtun/common/config"
 	"github.com/net-byte/vtun/common/netutil"
@@ -18,20 +18,18 @@ import (
 func StartClient(config config.Config) {
 	log.Printf("vtun websocket client started on %v", config.LocalAddr)
 	iface := tun.CreateTun(config)
+	go tunToWs(config, iface)
 	for {
 		if conn := netutil.ConnectServer(config); conn != nil {
-			var wg sync.WaitGroup
-			wg.Add(2)
-			go wsToTun(&wg, config, conn, iface)
-			go tunToWs(&wg, config, conn, iface)
-			wg.Wait()
-			conn.Close()
+			cache.GetCache().Set("wsconn", conn, 24*time.Hour)
+			wsToTun(config, conn, iface)
+			cache.GetCache().Delete("wsconn")
 		}
 	}
 }
 
-func wsToTun(wg *sync.WaitGroup, config config.Config, wsconn net.Conn, iface *water.Interface) {
-	defer wg.Done()
+func wsToTun(config config.Config, wsconn net.Conn, iface *water.Interface) {
+	defer wsconn.Close()
 	for {
 		wsconn.SetReadDeadline(time.Now().Add(time.Duration(config.Timeout) * time.Second))
 		packet, err := wsutil.ReadServerBinary(wsconn)
@@ -41,25 +39,30 @@ func wsToTun(wg *sync.WaitGroup, config config.Config, wsconn net.Conn, iface *w
 		if config.Obfs {
 			packet = cipher.XOR(packet)
 		}
-		iface.Write(packet)
+		_, err = iface.Write(packet)
+		if err != nil {
+			break
+		}
 	}
 }
 
-func tunToWs(wg *sync.WaitGroup, config config.Config, wsconn net.Conn, iface *water.Interface) {
-	defer wg.Done()
+func tunToWs(config config.Config, iface *water.Interface) {
 	packet := make([]byte, config.MTU)
 	for {
 		n, err := iface.Read(packet)
 		if err != nil || n == 0 {
-			break
+			continue
 		}
-		b := packet[:n]
-		if config.Obfs {
-			packet = cipher.XOR(packet)
-		}
-		wsconn.SetWriteDeadline(time.Now().Add(time.Duration(config.Timeout) * time.Second))
-		if err = wsutil.WriteClientBinary(wsconn, b); err != nil {
-			break
+		if v, ok := cache.GetCache().Get("wsconn"); ok {
+			b := packet[:n]
+			if config.Obfs {
+				packet = cipher.XOR(packet)
+			}
+			wsconn := v.(net.Conn)
+			wsconn.SetWriteDeadline(time.Now().Add(time.Duration(config.Timeout) * time.Second))
+			if err = wsutil.WriteClientBinary(wsconn, b); err != nil {
+				continue
+			}
 		}
 	}
 }