Bugfix: The server need to keep running after the connection closed by client

Author: leviathan0992

client/util/client.go

@@ -1,12 +1,12 @@
 package client
 
 import (
-	"sync/atomic"
 	"encoding/binary"
 	"log"
 	"net"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/leviathan1995/spleen/service"
@@ -37,7 +37,7 @@ func NewClient(clientID int, serverIP string, serverPort int, limitRate []string
 	}
 }
 
-var connections int64 = 0;
+var connections uint64 = 0
 
 func (c *client) Run() {
 	log.Printf("Begin to running the client[%d]", c.clientID)
@@ -46,20 +46,21 @@ func (c *client) Run() {
 	}
 
 	for {
-		if atomic.LoadInt64(&connections) < 10 {
-			for i := atomic.LoadInt64(&connections); i < 10; i++ {
+		if atomic.LoadUint64(&connections) < 10 {
+			for i := atomic.LoadUint64(&connections); i < 10; i++ {
 				srvConn, err := c.DialSrv()
 				if err != nil {
-					log.Printf("Connect to the server %s:%d failed: %s. \n", c.srvAddr.IP.String(), c.srvAddr.Port, err)
+					log.Printf("Connect to the proxy %s:%d failed: %s. \n", c.srvAddr.IP.String(), c.srvAddr.Port, err)
 					continue
 				}
-				log.Printf("Connect to the server %s:%d successful.\n", c.srvAddr.IP.String(), c.srvAddr.Port)
+				log.Printf("Connect to the proxy %s:%d successful.\n", c.srvAddr.IP.String(), c.srvAddr.Port)
 				srvConn.SetKeepAlive(true)
 				srvConn.SetLinger(0)
-				connections = atomic.AddInt64(&connections, 1)
+				atomic.StoreUint64(&connections, atomic.AddUint64(&connections, 1))
 				go c.handleConn(srvConn)
 			}
 		} else {
+			log.Printf("Currently, We still have %d active connections.", atomic.LoadUint64(&connections))
 			time.Sleep(1 * time.Second)
 		}
 	}
@@ -70,55 +71,57 @@ func (c *client) DialSrv() (*net.TCPConn, error) {
 }
 
 func (c *client) handleConn(srvConn *net.TCPConn) {
-	defer srvConn.Close()
-
-	transBuf := make([]byte, 8)
+	transBuf := make([]byte, service.IDBuf)
 	/* Send the ID of client to proxy. */
 	binary.LittleEndian.PutUint64(transBuf, uint64(c.clientID))
 	err := c.TCPWrite(srvConn, transBuf)
 	if err != nil {
-		connections = atomic.AddInt64(&connections, -1)
+		atomic.StoreUint64(&connections, atomic.AddUint64(&connections, ^uint64(1-1)))
+		_ = srvConn.Close()
 		log.Println("Try to send the ID of client to the proxy failed.")
 		return
 	}
 
 	/* Waiting for the transfer port from proxy. */
-	nRead, err := srvConn.Read(transBuf)
-	connections = atomic.AddInt64(&connections, -1)
+	err = c.TCPRead(srvConn, transBuf, service.PortBuf)
+	atomic.StoreUint64(&connections, atomic.AddUint64(&connections, ^uint64(1-1)))
 	if err != nil {
-		log.Println("Try to read the destination port failed.")
+		_ = srvConn.Close()
+		log.Println("Try to read destination port from the proxy failed.")
 		return
 	}
-	port := int64(binary.LittleEndian.Uint64(transBuf[:nRead]))
+	port := int64(binary.LittleEndian.Uint64(transBuf))
 
 	/* Try to direct connect to the destination sever. */
 	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())
+		_ = srvConn.Close()
+		log.Printf("Try to resolve TCPAddr %s failed: %s.\n", "localhost"+":"+strconv.FormatInt(port, 10), err.Error())
 		return
 	}
 
 	dstConn, err := net.DialTCP("tcp", nil, dstAddr)
 	if err != nil {
+		_ = srvConn.Close()
 		log.Printf("Connect to localhost:%d failed.", dstAddr.Port)
 		return
 	} else {
 		log.Printf("Connect to the destination address localhost:%d successful.", dstAddr.Port)
 	}
-	defer dstConn.Close()
 
-	dstConn.SetKeepAlive(true)
+	_ = dstConn.SetKeepAlive(true)
 	_ = dstConn.SetLinger(0)
 
 	var limitRate int64
-
 	if rate, found := c.limitRate[port]; found {
 		limitRate = rate * 1024 /* bytes */
 	}
 
 	go func() {
 		errTransfer := c.TransferToTCP(dstConn, srvConn, limitRate)
 		if errTransfer != nil {
+			_ = srvConn.Close()
+			_ = dstConn.Close()
 			return
 		}
 	}()

server/util/server.go

@@ -70,15 +70,23 @@ func (s *server) ListenForIntranet(tcpAddr *net.TCPAddr) {
 		if err != nil {
 			continue
 		}
-		/* The proxy should get the ID of the client first. */
-		transBuf := make([]byte, 8)
-		nRead, err := conn.Read(transBuf)
+
+		/* The proxy should get the magic number and ID of the client first. */
+		transBuf := make([]byte, service.IDBuf)
+		err = s.TCPRead(conn, transBuf, service.IDBuf)
 		if err != nil {
+			_ = conn.Close()
 			log.Println("Try to read the destination port failed.")
-			return
+			continue
+		}
+
+		id := int64(binary.LittleEndian.Uint64(transBuf))
+		if ConnectionPool.Has(s.connectionPool, id) == false {
+			_ = conn.Close()
+			continue
+		} else {
+			s.connectionPool[id] <- conn
 		}
-		id := int64(binary.LittleEndian.Uint64(transBuf[:nRead]))
-		s.connectionPool[id] <- conn
 	}
 }
 
@@ -100,41 +108,41 @@ func (s *server) Listen() {
 }
 
 func (s *server) handleConn(cliConn *net.TCPConn, clientID int64, transferPort uint64) {
-	defer cliConn.Close()
-
 	select {
 	case intranetConn := <-s.connectionPool[clientID]:
 		_ = intranetConn.SetLinger(0)
 
 		/* Send the transfer port to intranet server . */
-		portBuf := make([]byte, 8)
+		portBuf := make([]byte, service.PortBuf)
 		binary.LittleEndian.PutUint64(portBuf, transferPort)
 		err := s.TCPWrite(intranetConn, portBuf)
 		if err != nil {
-			intranetConn.Close()
+			_ = cliConn.Close()
+			_ = intranetConn.Close()
 			for {
 				/* Close all connections from this client. */
 				select {
 				case intranetConn = <-s.connectionPool[clientID]:
-					intranetConn.Close()
+					_ = intranetConn.SetLinger(0)
+					_ = intranetConn.Close()
 				default:
 					return
 				}
 			}
 		}
 
-		log.Printf("Make a successful connection between the user and the intranet server[Client ID: %d - Port: %d].", clientID, transferPort)
+		log.Printf("Make a successful connection between the user [%s] and the intranet server[Client ID: %d - Port: %d].",
+			cliConn.RemoteAddr().String(), clientID, transferPort)
 		/* Transfer network packets. */
 		go func() {
 			errTransfer := s.TransferToTCP(cliConn, intranetConn, 0)
 			if errTransfer != nil {
-				intranetConn.Close()
+				_ = cliConn.Close()
+				_ = intranetConn.Close()
 				return
 			}
 		}()
 		err = s.TransferToTCP(intranetConn, cliConn, 0)
 		return
-	default:
-		log.Printf("Currently, Do not have any active connection from the intranet server[Client ID: %d - Port: %d].", clientID, transferPort)
 	}
 }

service/service.go

@@ -5,7 +5,9 @@ import (
 	"time"
 )
 
-const BUFFERSIZE = 1024 * 16
+const TransferBuf = 1024 * 16
+const PortBuf = 8
+const IDBuf = 8
 
 type Service struct {
 	IP   string
@@ -25,9 +27,21 @@ func (s *Service) TCPWrite(conn *net.TCPConn, buf []byte) error {
 	return nil
 }
 
+func (s *Service) TCPRead(conn *net.TCPConn, buf []byte, len int) error {
+	nRead := 0
+	for nRead < len {
+		n, errRead := conn.Read(buf[nRead:])
+		if errRead != nil {
+			return errRead
+		}
+		nRead += n
+	}
+	return nil
+}
+
 func (s *Service) TransferToTCP(cliConn net.Conn, dstConn *net.TCPConn, limitRate int64) error {
 	var totalRead, lastTime int64
-	buf := make([]byte, BUFFERSIZE)
+	buf := make([]byte, TransferBuf)
 
 	for {
 		nRead, errRead := cliConn.Read(buf)