Support the rate-limiting

Author: leviathan0992

.client.json

@@ -1,4 +1,7 @@
 {
-  "server_ip"  : "127.0.0.1",
-  "server_port": 1234
-}
+  "ServerIP"  : "127.0.0.1",
+  "ServerPort": 1234,
+  "LimitRate" : [
+    "5001:512"
+  ]
+}

.server.json

@@ -3,6 +3,6 @@
   "ServerPort"   : 1234,
   "MappingPort" : [
 	"5000:22",
-	"5001:3306"
+	"5002:5001"
   ]
 }

README.md

@@ -4,7 +4,7 @@
 
 ## 介绍
 
-通过在一台具有公网 IP 的小型服务器(阿里云轻量)部署 `spleen`, 可以随时随地访问你的家庭服务器/内网主机(闲置笔记本)的 `TCP` 服务, 例如 `SSH`, `HTTP/S` 等. 
+通过在一台具有公网 IP 的小型服务器(阿里云轻量)部署 `spleen`, 可以随时随地访问你的家庭服务器/内网主机(闲置笔记本)的 `TCP` 服务, 例如 `SSH`, `HTTP/S` 等. `spleen`支持针对内网服务器指定服务进行流速限制, 避免打爆公网小水管.
 
 例如 `SSH` 服务, 在顺利部署 `spleen` 的客户端和服务端后, 通过公网服务器(假定 IP 为`1.1.1.1`), 可以直接通过端口映射来连接你的家庭服务器/内网主机:
 
@@ -64,6 +64,14 @@
   "server_port": 1234 # 公网服务器监听端口
 }
 
+{
+  "ServerIP"  : "127.0.0.1", # 公网服务器 IP
+  "ServerPort": 1234, # 公网服务器监听端口
+  "LimitRate" : [
+    "5001:512" # 指定内网服务的流速限制, 例如指定内网 5001 端口的流速不要超过 512 KB/s, 单位是 KB, 默认不限制, 0 值即为不限制.
+  ]
+}
+
 # 启动
 > ./spleen-client -c .client.json # 默认预留 10 个活跃连接
 2022/01/12 18:55:19 Connect to the server 127.0.0.1:1234 successful.

client/main.go

@@ -9,9 +9,14 @@ import (
 	"github.com/leviathan1995/spleen/client/util"
 )
 
+type Configuration struct {
+	ServerIP   string
+	ServerPort int
+	LimitRate  []string
+}
+
 func main() {
 	var conf string
-	var config map[string]interface{}
 	flag.StringVar(&conf, "c", ".client.json", "The client configuration.")
 	flag.Parse()
 
@@ -20,12 +25,11 @@ func main() {
 		log.Fatalf("Reading %s failed.", conf)
 	}
 
+	var config Configuration
 	if err := json.Unmarshal(bytes, &config); err != nil {
 		log.Fatalf("Parsing %s failed.", conf)
 	}
 
-	serverIP := config["server_ip"].(string)
-	serverPort := int(config["server_port"].(float64))
-	c := client.NewClient(serverIP, serverPort)
+	c := client.NewClient(config.ServerIP, config.ServerPort, config.LimitRate)
 	c.Run()
 }

client/util/client.go

@@ -5,36 +5,53 @@ import (
 	"log"
 	"net"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/leviathan1995/spleen/service"
 )
 
 type client struct {
 	*service.Service
-	srvAddr *net.TCPAddr
+	srvAddr   *net.TCPAddr
+	limitRate map[int64]int64
 }
 
-func NewClient(serverIP string, serverPort int) *client {
+func NewClient(serverIP string, serverPort int, limitRate []string) *client {
 	srvAddr, _ := net.ResolveTCPAddr("tcp", serverIP+":"+strconv.Itoa(serverPort))
+	limits := make(map[int64]int64)
+	for _, rates := range limitRate {
+		rate := strings.Split(rates, ":")
+		port, _ := strconv.Atoi(rate[0])
+		speed, _ := strconv.Atoi(rate[1])
+		limits[int64(port)] = int64(speed)
+	}
+
 	return &client{
 		&service.Service{},
 		srvAddr,
+		limits,
 	}
 }
 
 func (c *client) Run() {
+	for port, rate := range c.limitRate {
+		log.Printf("The limiting rate configurations  Port: %d, Speed: %d KB/s", port, rate)
+	}
+
 	for {
 		if len(connectionPool) < 10 {
-			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)
-				continue
+			for i := len(connectionPool); 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)
+					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)
 			}
-			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 {
 			time.Sleep(1 * time.Second)
 		}
@@ -75,15 +92,22 @@ func (c *client) handleConn(srvConn *net.TCPConn) {
 	} else {
 		log.Printf("Connect to the destination address localhost:%d successful.", dstAddr.Port)
 	}
-
 	defer dstConn.Close()
+
+	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)
+		errTransfer := c.TransferToTCP(dstConn, srvConn, limitRate)
 		if errTransfer != nil {
 			return
 		}
 	}()
-	err = c.TransferToTCP(srvConn, dstConn)
+	err = c.TransferToTCP(srvConn, dstConn, 0)
 }

server/util/server.go

@@ -98,13 +98,15 @@ func (s *server) handleConn(cliConn *net.TCPConn, transferPort uint64) {
 			log.Print("Make a successful connection between the user and the intranet server.")
 			/* Transfer network packets. */
 			go func() {
-				errTransfer := s.TransferToTCP(cliConn, intranetConn)
+				errTransfer := s.TransferToTCP(cliConn, intranetConn, 0)
 				if errTransfer != nil {
 					intranetConn.Close()
 					return
 				}
 			}()
-			err = s.TransferToTCP(intranetConn, cliConn)
+			err = s.TransferToTCP(intranetConn, cliConn, 0)
+			return
 		}
 	}
+	log.Println("Currently, Do not have any active connection from the intranet server.")
 }

service/service.go

@@ -2,6 +2,7 @@ package service
 
 import (
 	"net"
+	"time"
 )
 
 const BUFFERSIZE = 1024 * 16
@@ -24,18 +25,35 @@ func (s *Service) TCPWrite(conn *net.TCPConn, buf []byte) error {
 	return nil
 }
 
-func (s *Service) TransferToTCP(cliConn net.Conn, dstConn *net.TCPConn) error {
+func (s *Service) TransferToTCP(cliConn net.Conn, dstConn *net.TCPConn, limitRate int64) error {
+	var totalRead, lastTime int64
 	buf := make([]byte, BUFFERSIZE)
+
 	for {
-		nRead, err := cliConn.Read(buf)
-		if err != nil {
-			return err
+		nRead, errRead := cliConn.Read(buf)
+		if errRead != nil {
+			return errRead
 		}
 		if nRead > 0 {
 			errWrite := s.TCPWrite(dstConn, buf[0:nRead])
-			if err != nil {
+			if errWrite != nil {
 				return errWrite
 			}
+			if limitRate != 0 {
+				if totalRead > limitRate && ((time.Now().UnixNano()/1e6)-lastTime) >= 1000 {
+					/* Reset the timeout. */
+					totalRead = 0
+					lastTime = time.Now().UnixNano() / 1e6 /* The millisecond */
+				} else if totalRead > limitRate && ((time.Now().UnixNano()/1e6)-lastTime) < 1000 {
+					/* Try to limit the rate of network. */
+					timeout := 1000 - ((time.Now().UnixNano() / 1e6) - lastTime)
+					time.Sleep(time.Duration(timeout) * time.Millisecond)
+					totalRead = 0
+					lastTime = time.Now().UnixNano() / 1e6 /* The millisecond */
+				} else {
+					totalRead += int64(nRead)
+				}
+			}
 		}
 	}
 }