From 51a37041f1c8ef8c890abffb909a88ac4a54b154 Mon Sep 17 00:00:00 2001
From: Srar <a1020938070@gmail.com>
Date: Wed, 14 Aug 2019 12:51:24 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E5=93=8D=E5=BA=94=E9=A2=91?=
 =?UTF-8?q?=E7=8E=87=E9=99=90=E5=88=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 config.ini                         |  3 ++-
 core/protocol/memcache/memcache.go | 26 +++++++++++++++++++++-----
 utils/setting/setting.go           |  3 ++-
 3 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/config.ini b/config.ini
index 2a20535..3191647 100644
--- a/config.ini
+++ b/config.ini
@@ -53,4 +53,5 @@ addr = 0.0.0.0:21                            # Ftp 服务端地址 注意端口
 
 [memcache]
 status = 0                                   # 是否启动 Memcache 1 启动 0 关闭
-addr = 0.0.0.0:11211                         # Memcache 服务端地址 注意端口冲突
\ No newline at end of file
+addr = 0.0.0.0:11211                         # Memcache 服务端地址 注意端口冲突
+ratelimit = 4                                # 每秒响应次数
\ No newline at end of file
diff --git a/core/protocol/memcache/memcache.go b/core/protocol/memcache/memcache.go
index 9d3a183..4b35843 100644
--- a/core/protocol/memcache/memcache.go
+++ b/core/protocol/memcache/memcache.go
@@ -396,7 +396,7 @@ var commands = map[string]func([]string) ([]byte, int){
 }
 
 // TCP服务端连接
-func tcpServer(address string, exitChan chan int) {
+func tcpServer(address string, rateLimitChan chan int, exitChan chan int) {
 	l, err := net.Listen("tcp", address)
 
 	if err != nil {
@@ -422,6 +422,7 @@ func tcpServer(address string, exitChan chan int) {
 			reader := bufio.NewReader(conn)
 			log.Printf("[Memcache TCP %d] Accepted a client socket from %s\n", trackID, conn.RemoteAddr().String())
 			for {
+				<-rateLimitChan
 				str, err := reader.ReadString('\n')
 				if skip {
 					skip = false
@@ -470,7 +471,7 @@ func tcpServer(address string, exitChan chan int) {
 
 }
 
-func udpServer(address string, exitChan chan int) {
+func udpServer(address string, rateLimitChan chan int, exitChan chan int) {
 	udpAddr, err := net.ResolveUDPAddr("udp", address)
 	if err != nil {
 		fmt.Println(err.Error())
@@ -488,6 +489,7 @@ func udpServer(address string, exitChan chan int) {
 	go func() {
 		buf := make([]byte, 1500)
 		for {
+			<-rateLimitChan
 			plen, addr, _ := l.ReadFromUDP(buf)
 			/* UDP协议需要8个字节的头 */
 			if plen < 8 {
@@ -522,13 +524,27 @@ func udpServer(address string, exitChan chan int) {
 	}()
 }
 
-func Start(addr string) {
+func Start(addr string, rateLimitStr string) {
 	// 创建一个程序结束码的通道
 	exitChan := make(chan int)
 
+	// 响应间隔限制
+	rateLimitChan := make(chan int)
+	rateLimit, err := strconv.Atoi(rateLimitStr)
+	if err != nil {
+		panic(err)
+	}
+	go func() {
+		sleepTime := 1000 / rateLimit
+		for {
+			rateLimitChan <- 1
+			time.Sleep(time.Duration(sleepTime) * time.Millisecond)
+		}
+	}()
+
 	// 将服务器并发运行
-	go tcpServer(addr, exitChan)
-	go udpServer(addr, exitChan)
+	go tcpServer(addr, rateLimitChan, exitChan)
+	go udpServer(addr, rateLimitChan, exitChan)
 
 	// 通道阻塞，等待接受返回值
 	code := <-exitChan
diff --git a/utils/setting/setting.go b/utils/setting/setting.go
index 61a9b8b..da3611b 100644
--- a/utils/setting/setting.go
+++ b/utils/setting/setting.go
@@ -211,11 +211,12 @@ func Run() {
 
 	// 启动 Memcache 蜜罐
 	memcacheStatus := conf.Get("memcache", "status")
+	memcacheRateLimit := conf.Get("memcache", "ratelimit")
 
 	// 判断 暗网 Web 蜜罐 是否开启
 	if memcacheStatus == "1" {
 		memcacheAddr := conf.Get("memcache", "addr")
-		go memcache.Start(memcacheAddr)
+		go memcache.Start(memcacheAddr, memcacheRateLimit)
 	}
 
 	//=========================//