Minor: Remove unnecessary codes

Leviathan1995 · Feb 13, 2022 · a7ea7f8 · a7ea7f8
1 parent 57ee029
commit a7ea7f8
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 28 deletions.
diff --git a/client/util/client.go b/client/util/client.go
@@ -37,7 +37,7 @@ func NewClient(clientID int, serverIP string, serverPort int, limitRate []string
 	}
 }
 
-var connections uint64 = 0
+var freeConnections int64 = 0
 
 func (c *client) Run() {
 	log.Printf("Begin to running the client[%d]", c.clientID)
@@ -46,8 +46,8 @@ func (c *client) Run() {
 	}
 
 	for {
-		if atomic.LoadUint64(&connections) < 10 {
-			for i := atomic.LoadUint64(&connections); i < 10; i++ {
+		if atomic.LoadInt64(&freeConnections) < 10 {
+			for i := atomic.LoadInt64(&freeConnections); i < 10; i++ {
 				srvConn, err := c.DialSrv()
 				if err != nil {
 					log.Printf("Connect to the proxy %s:%d failed: %s. \n", c.srvAddr.IP.String(), c.srvAddr.Port, err)
@@ -57,12 +57,11 @@ func (c *client) Run() {
 				_ = srvConn.SetLinger(0)
 				_ = srvConn.SetKeepAlive(true)
 				_ = srvConn.SetKeepAlivePeriod(2 * time.Second)
-
-				atomic.StoreUint64(&connections, atomic.AddUint64(&connections, 1))
+				atomic.AddInt64(&freeConnections, 1)
 				go c.handleConn(srvConn)
 			}
 		} else {
-			log.Printf("Currently, We still have %d active connections.", atomic.LoadUint64(&connections))
+			log.Printf("Currently, We still have %d active connections.", atomic.LoadInt64(&freeConnections))
 			time.Sleep(1 * time.Second)
 		}
 	}
@@ -78,7 +77,7 @@ func (c *client) handleConn(srvConn *net.TCPConn) {
 	binary.LittleEndian.PutUint64(transBuf, uint64(c.clientID))
 	err := c.TCPWrite(srvConn, transBuf)
 	if err != nil {
-		atomic.StoreUint64(&connections, atomic.AddUint64(&connections, ^uint64(1-1)))
+		atomic.AddInt64(&freeConnections, -1)
 		_ = srvConn.Close()
 		log.Println("Try to send the ID of client to the proxy failed.")
 		return
@@ -87,7 +86,7 @@ func (c *client) handleConn(srvConn *net.TCPConn) {
 	/* It has to wait 3600 seconds before get the transfer port from the proxy. */
 	_ = srvConn.SetDeadline(time.Now().Add(3600 * time.Second))
 	err = c.TCPRead(srvConn, transBuf, service.PortBuf)
-	atomic.StoreUint64(&connections, atomic.AddUint64(&connections, ^uint64(1-1)))
+	atomic.AddInt64(&freeConnections, -1)
 	if err != nil {
 		_ = srvConn.Close()
 		log.Println("Try to read destination port from the proxy failed or maybe the connection is closed.")
@@ -124,13 +123,12 @@ func (c *client) handleConn(srvConn *net.TCPConn) {
 	}
 
 	go func() {
-		errTransfer := c.TransferToTCP(dstConn, srvConn, limitRate)
-		if errTransfer != nil {
-			_ = srvConn.Close()
-			_ = dstConn.Close()
-			return
-		}
+		_ = c.TransferToTCP(dstConn, srvConn, limitRate)
 	}()
+
 	_ = c.TransferToTCP(srvConn, dstConn, 0)
+	_ = srvConn.Close()
+	_ = dstConn.Close()
+
 	return
 }
diff --git a/server/util/server.go b/server/util/server.go
@@ -6,6 +6,7 @@ import (
 	"log"
 	"net"
 	"strconv"
+	"time"
 )
 
 type ConnectionPool map[uint64]chan *net.TCPConn
@@ -74,7 +75,9 @@ func (s *server) ListenForIntranet(tcpAddr *net.TCPAddr) {
 	} else {
 		log.Printf("The server listening for the intranet server at %s:%d successful.", s.IP, s.Port)
 	}
+
 	defer listener.Close()
+
 	for {
 		conn, err := listener.AcceptTCP()
 		if err != nil {
@@ -116,13 +119,6 @@ func (s *server) Listen() {
 }
 
 func (s *server) handleConn(cliConn *net.TCPConn, clientID uint64, transferPort uint64) {
-	defer func(cliConn *net.TCPConn) {
-		err := cliConn.Close()
-		if err != nil {
-			return
-		}
-	}(cliConn)
-
 	for {
 		select {
 		case intranetConn := <-s.connectionPool[clientID]:
@@ -137,20 +133,21 @@ func (s *server) handleConn(cliConn *net.TCPConn, clientID uint64, transferPort
 				_ = intranetConn.Close()
 				continue
 			} else {
-				log.Printf("Make a successful connection between the user [%s] and the intranet server[Client ID: %d - Port: %d].",
+				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 {
-						_ = cliConn.Close()
-						_ = intranetConn.Close()
-						return
-					}
+					_ = s.TransferToTCP(cliConn, intranetConn, 0)
+					_ = intranetConn.Close()
+					_ = cliConn.Close()
 				}()
+
 				_ = s.TransferToTCP(intranetConn, cliConn, 0)
 				return
 			}
+		case <-time.After(1 * time.Second):
+			log.Printf("Currently, We don't have any active connection from the intranet server[Client ID: %d - Port: %d].",
+				clientID, transferPort)
 		}
 	}
 }