From 912792b4adec2ba4dedbf6e81a828d65c3ea132e Mon Sep 17 00:00:00 2001
From: Alex Tsai <alex.tsai@protonmail.com>
Date: Mon, 24 May 2021 11:25:22 +0800
Subject: [PATCH] refactor

---
 client/udpclient.go     |  7 ++-----
 client/wsclient.go      |  5 ++---
 common/cipher/cipher.go | 16 +++-------------
 main.go                 |  2 +-
 server/udpserver.go     |  7 ++-----
 server/wsserver.go      |  5 ++---
 6 files changed, 12 insertions(+), 30 deletions(-)

diff --git a/client/udpclient.go b/client/udpclient.go
index 2959027..63de382 100644
--- a/client/udpclient.go
+++ b/client/udpclient.go
@@ -12,7 +12,6 @@ import (
 
 // StartUDPClient start udp client
 func StartUDPClient(config config.Config) {
-	config.Init()
 	iface := tun.CreateTun(config.CIDR)
 	serverAddr, err := net.ResolveUDPAddr("udp", config.ServerAddr)
 	if err != nil {
@@ -36,8 +35,7 @@ func StartUDPClient(config config.Config) {
 			if err != nil || n == 0 {
 				continue
 			}
-			// decrypt data
-			b := cipher.Decrypt(buf[:n])
+			b := cipher.XOR(buf[:n])
 			if !waterutil.IsIPv4(b) {
 				continue
 			}
@@ -54,8 +52,7 @@ func StartUDPClient(config config.Config) {
 		if !waterutil.IsIPv4(packet) {
 			continue
 		}
-		// encrypt data
-		b := cipher.Encrypt(packet[:n])
+		b := cipher.XOR(packet[:n])
 		conn.WriteToUDP(b, serverAddr)
 	}
 }
diff --git a/client/wsclient.go b/client/wsclient.go
index fc39c0f..115f99b 100644
--- a/client/wsclient.go
+++ b/client/wsclient.go
@@ -18,7 +18,6 @@ import (
 
 // StartWSClient start ws client
 func StartWSClient(config config.Config) {
-	config.Init()
 	iface := tun.CreateTun(config.CIDR)
 	c := cache.New(30*time.Minute, 10*time.Minute)
 	log.Printf("vtun ws client started,CIDR is %v", config.CIDR)
@@ -50,7 +49,7 @@ func StartWSClient(config config.Config) {
 			c.Set(key, conn, cache.DefaultExpiration)
 			go wsToTun(c, key, conn, iface)
 		}
-		b = cipher.Encrypt(b)
+		b = cipher.XOR(b)
 		conn.WriteMessage(websocket.BinaryMessage, b)
 	}
 }
@@ -63,7 +62,7 @@ func wsToTun(c *cache.Cache, key string, wsConn *websocket.Conn, iface *water.In
 		if err != nil || err == io.EOF {
 			break
 		}
-		b = cipher.Decrypt(b)
+		b = cipher.XOR(b)
 		if !waterutil.IsIPv4(b) {
 			continue
 		}
diff --git a/common/cipher/cipher.go b/common/cipher/cipher.go
index c764fd8..1629abd 100644
--- a/common/cipher/cipher.go
+++ b/common/cipher/cipher.go
@@ -11,22 +11,12 @@ func GenerateKey(key string) {
 	_key = []byte(key)
 }
 
-func Encrypt(data []byte) []byte {
+func XOR(src []byte) []byte {
 	c, err := rc4.NewCipher(_key)
 	if err != nil {
 		log.Fatalln(err)
 	}
-	dst := make([]byte, len(data))
-	c.XORKeyStream(dst, data)
-	return dst
-}
-
-func Decrypt(data []byte) []byte {
-	c, err := rc4.NewCipher(_key)
-	if err != nil {
-		log.Fatalln(err)
-	}
-	dst := make([]byte, len(data))
-	c.XORKeyStream(dst, data)
+	dst := make([]byte, len(src))
+	c.XORKeyStream(dst, src)
 	return dst
 }
diff --git a/main.go b/main.go
index f628fa6..9f3b6ea 100644
--- a/main.go
+++ b/main.go
@@ -18,7 +18,7 @@ func main() {
 	flag.BoolVar(&config.ServerMode, "S", false, "server mode")
 	flag.BoolVar(&config.TLS, "t", false, "enable tls")
 	flag.Parse()
-
+	config.Init()
 	switch config.Protocol {
 	case "udp":
 		if config.ServerMode {
diff --git a/server/udpserver.go b/server/udpserver.go
index 4849c02..54c1f28 100644
--- a/server/udpserver.go
+++ b/server/udpserver.go
@@ -17,7 +17,6 @@ import (
 
 // StartUDPServer start udp server
 func StartUDPServer(config config.Config) {
-	config.Init()
 	iface := tun.CreateTun(config.CIDR)
 	localAddr, err := net.ResolveUDPAddr("udp", config.LocalAddr)
 	if err != nil {
@@ -39,8 +38,7 @@ func StartUDPServer(config config.Config) {
 		if err != nil || n == 0 {
 			continue
 		}
-		// decrypt data
-		b := cipher.Decrypt(buf[:n])
+		b := cipher.XOR(buf[:n])
 		if !waterutil.IsIPv4(b) {
 			continue
 		}
@@ -77,8 +75,7 @@ func (f *Forwarder) forward(iface *water.Interface, conn *net.UDPConn) {
 		key := fmt.Sprintf("%v->%v", dstAddr, srcAddr)
 		v, ok := f.connCache.Get(key)
 		if ok {
-			// encrypt data
-			b = cipher.Encrypt(b)
+			b = cipher.XOR(b)
 			f.localConn.WriteToUDP(b, v.(*net.UDPAddr))
 		}
 	}
diff --git a/server/wsserver.go b/server/wsserver.go
index 6d65cce..a821d28 100644
--- a/server/wsserver.go
+++ b/server/wsserver.go
@@ -30,7 +30,6 @@ var upgrader = websocket.Upgrader{
 
 // StartWSServer start ws server
 func StartWSServer(config config.Config) {
-	config.Init()
 	iface := tun.CreateTun(config.CIDR)
 	c := cache.New(30*time.Minute, 10*time.Minute)
 	go tunToWs(iface, c)
@@ -128,7 +127,7 @@ func tunToWs(iface *water.Interface, c *cache.Cache) {
 		key := fmt.Sprintf("%v->%v", dstAddr, srcAddr)
 		v, ok := c.Get(key)
 		if ok {
-			b = cipher.Encrypt(b)
+			b = cipher.XOR(b)
 			v.(*websocket.Conn).WriteMessage(websocket.BinaryMessage, b)
 		}
 	}
@@ -142,7 +141,7 @@ func wsToTun(wsConn *websocket.Conn, iface *water.Interface, c *cache.Cache) {
 		if err != nil || err == io.EOF {
 			break
 		}
-		b = cipher.Decrypt(b)
+		b = cipher.XOR(b)
 		if !waterutil.IsIPv4(b) {
 			continue
 		}