Bugfix: The client reads multi-TCP packets causes parsing to fail.

Leviathan1995 · Jan 12, 2022 · 41f6b81 · 41f6b81
1 parent f13e695
commit 41f6b81
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 35 deletions.
diff --git a/README.md b/README.md
@@ -13,11 +13,11 @@
 * 通过 [release]() 下载对应架构的 spleen 包:
 ```shell
 # wget 下载 (请自行替换最新版本)
-> wget https://github.com/Leviathan1995/spleen/releases/download/v0.0.1/spleen_0.0.1_Linux_64-bit.tar.gz
+> wget https://github.com/Leviathan1995/spleen/releases/download/v0.0.2/spleen_0.0.2_Linux_64-bit.tar.gz
 
 # 解压
-> tar -zxvf spleen_0.0.1_Linux_64-bit.tar.gz
-> cd spleen_0.0.1_Linux_64-bit
+> tar -zxvf spleen_0.0.2_Linux_64-bit.tar.gz
+> cd spleen_0.0.2_Linux_64-bit
 
 # 配置端口转发规则
 > vim .server.json
@@ -44,11 +44,11 @@
 * 通过 [release]() 下载对应架构的 spleen 包:
 ```shell
 # wget 下载 (请自行替换最新版本)
-> wget https://github.com/Leviathan1995/spleen/releases/download/v0.0.1/spleen_0.0.1_Linux_64-bit.tar.gz
+> wget https://github.com/Leviathan1995/spleen/releases/download/v0.0.2/spleen_0.0.2_Linux_64-bit.tar.gz
 
 # 解压
-> tar -zxvf spleen_0.0.1_Linux_64-bit.tar.gz
-> cd spleen_0.0.1_Linux_64-bit
+> tar -zxvf spleen_0.0.2_Linux_64-bit.tar.gz
+> cd spleen_0.0.2_Linux_64-bit
 
 # 配置公网服务器地址
 > vim .client.json

diff --git a/client/util/client.go b/client/util/client.go
@@ -1,6 +1,7 @@
 package client
 
 import (
+	"encoding/binary"
 	"log"
 	"net"
 	"strconv"
@@ -31,6 +32,7 @@ func (c *client) Run() {
 				continue
 			}
 			log.Printf("Connect to the server %s:%d successful.\n", c.srvAddr.IP.String(), c.srvAddr.Port)
+			srvConn.SetKeepAlive(true)
 			connectionPool <- srvConn
 			go c.handleConn(srvConn)
 		} else {
@@ -49,22 +51,29 @@ func (c *client) handleConn(srvConn *net.TCPConn) {
 	defer srvConn.Close()
 
 	/* Get the transfer port. */
-	portBuf := make([]byte, 32)
+	portBuf := make([]byte, 8)
 	nRead, err := srvConn.Read(portBuf)
+	_ = <-connectionPool /* Remove a connection from pool. */
 	if err != nil {
+		log.Println("Try to read the destination port failed.")
 		return
 	}
-	_ = <-connectionPool /* Remove a connection from pool. */
+	port := int64(binary.LittleEndian.Uint64(portBuf[:nRead]))
 
 	/* Try to direct connect to the destination sever. */
-	dstAddr, _ := net.ResolveTCPAddr("tcp", "127.0.0.1"+":"+string(portBuf[:nRead]))
-	log.Printf("Try to connect %s:%d.\n", dstAddr.IP.String(), dstAddr.Port)
+	log.Printf("Try to connect %s.\n", "localhost"+":"+strconv.Itoa(int(port)))
+	dstAddr, err := net.ResolveTCPAddr("tcp", ":"+strconv.Itoa(int(port)))
+	if err != nil {
+		log.Printf("Try to resolve TCPAddr %s failed: %s.\n", "localhost"+":"+string(port), err.Error())
+		return
+	}
+
 	dstConn, err := net.DialTCP("tcp", nil, dstAddr)
 	if err != nil {
-		log.Printf("Connect to %s:%d failed.", dstAddr.IP.String(), dstAddr.Port)
+		log.Printf("Connect to localhost:%d failed.", dstAddr.Port)
 		return
 	} else {
-		log.Printf("Connect to the destination address %s:%d successful.", dstAddr.IP, dstAddr.Port)
+		log.Printf("Connect to the destination address localhost:%d successful.", dstAddr.Port)
 	}
 
 	defer dstConn.Close()
@@ -73,7 +82,7 @@ func (c *client) handleConn(srvConn *net.TCPConn) {
 	go func() {
 		errTransfer := c.TransferToTCP(dstConn, srvConn)
 		if errTransfer != nil {
-			log.Println(errTransfer.Error())
+			return
 		}
 	}()
 	err = c.TransferToTCP(srvConn, dstConn)

diff --git a/server/util/server.go b/server/util/server.go
@@ -1,6 +1,7 @@
 package server
 
 import (
+	"encoding/binary"
 	"log"
 	"net"
 	"strconv"
@@ -41,7 +42,8 @@ func (s *server) ListenOnPort(tcpAddr *net.TCPAddr, transferPort string) {
 		if err != nil {
 			continue
 		}
-		go s.handleConn(cliConn, transferPort)
+		port, _ := strconv.Atoi(transferPort)
+		go s.handleConn(cliConn, uint64(port))
 	}
 }
 
@@ -58,7 +60,6 @@ func (s *server) ListenForIntranet(tcpAddr *net.TCPAddr) {
 	for {
 		conn, err := listener.AcceptTCP()
 		if err != nil {
-			log.Println(err.Error())
 			continue
 		}
 		connectionPool <- conn
@@ -76,29 +77,34 @@ func (s *server) Listen() {
 	s.ListenForIntranet(tcpAddr)
 }
 
-func (s *server) handleConn(cliConn *net.TCPConn, transferPort string) {
+func (s *server) handleConn(cliConn *net.TCPConn, transferPort uint64) {
 	defer cliConn.Close()
 
-	intranetConn := <-connectionPool
-	if intranetConn != nil {
-		defer intranetConn.Close()
-		_ = intranetConn.SetLinger(0)
+	for i := 0; i < 20; i++ {
+		intranetConn := <-connectionPool
+		if intranetConn != nil {
+			_ = intranetConn.SetLinger(0)
 
-		/* Send the mapping port to intranet server . */
-		log.Println("Send the mapping port to intranet server.")
-		err := s.TCPWrite(intranetConn, []byte(transferPort))
-		if err != nil {
-			return
-		}
-
-		log.Print("Make a successful connection between client and the intranet server.")
-		/* Transfer network packets. */
-		go func() {
-			errTransfer := s.TransferToTCP(cliConn, intranetConn)
-			if errTransfer != nil {
-				log.Println(errTransfer.Error())
+			/* Send the mapping port to intranet server . */
+			log.Printf("Send the mapping port %d to intranet server.\n", transferPort)
+			portBuf := make([]byte, 8)
+			binary.LittleEndian.PutUint64(portBuf, transferPort)
+			err := s.TCPWrite(intranetConn, portBuf)
+			if err != nil {
+				intranetConn.Close()
+				continue
 			}
-		}()
-		err = s.TransferToTCP(intranetConn, cliConn)
+
+			log.Print("Make a successful connection between client and the intranet server.")
+			/* Transfer network packets. */
+			go func() {
+				errTransfer := s.TransferToTCP(cliConn, intranetConn)
+				if errTransfer != nil {
+					intranetConn.Close()
+					return
+				}
+			}()
+			err = s.TransferToTCP(intranetConn, cliConn)
+		}
 	}
 }