Merge pull request #254 from wdvxdr1123/master

feature: _rate_limit

docs/config.md

@@ -27,6 +27,11 @@ go-cqhttp 支持导入CQHTTP的配置文件, 具体步骤为:
 		"relogin_delay": 3,
 		"max_relogin_times": 0
 	},
+    "_rate_limit": {
+		"enabled": false,
+		"frequency": 1,
+		"bucket_size": 1
+    },
 	"post_message_format": "string",
 	"ignore_invalid_cqcode": false,
 	"force_fragmented": true,
@@ -66,6 +71,9 @@ go-cqhttp 支持导入CQHTTP的配置文件, 具体步骤为:
 | relogin              | bool     | 是否自动重新登录                                                     |
 | relogin_delay        | int      | 重登录延时（秒）                                                     |
 | max_relogin_times    | uint     | 最大重登录次数，若0则不设置上限                                        |
+| _rate_limit          | bool     | 是否启用API调用限速                                                  |
+| frequency            | float64  | 1s内能调用API的次数                                                  |
+| bucket_size          | int      | 令牌桶的大小，默认为1，修改此值可允许一定程度内连续调用api                 |
 | post_message_format  | string   | 上报信息类型                                                        |
 | ignore_invalid_cqcode| bool     | 是否忽略错误的CQ码                                                   |
 | force_fragmented     | bool     | 是否强制分片发送群长消息                                              |

global/config.go

@@ -21,6 +21,11 @@ type JsonConfig struct {
 		ReLoginDelay    int  `json:"relogin_delay"`
 		MaxReloginTimes uint `json:"max_relogin_times"`
 	} `json:"relogin"`
+	RateLimit struct {
+		Enabled    bool    `json:"enabled"`
+		Frequency  float64 `json:"frequency"`
+		BucketSize int     `json:"bucket_size"`
+	} `json:"_rate_limit"`
 	IgnoreInvalidCQCode bool                          `json:"ignore_invalid_cqcode"`
 	ForceFragmented     bool                          `json:"force_fragmented"`
 	HeartbeatInterval   time.Duration                 `json:"heartbeat_interval"`
@@ -85,6 +90,15 @@ func DefaultConfig() *JsonConfig {
 			ReLoginDelay:    3,
 			MaxReloginTimes: 0,
 		},
+		RateLimit: struct {
+			Enabled    bool    `json:"enabled"`
+			Frequency  float64 `json:"frequency"`
+			BucketSize int     `json:"bucket_size"`
+		}{
+			Enabled:    false,
+			Frequency:  1,
+			BucketSize: 1,
+		},
 		PostMessageFormat: "string",
 		ForceFragmented:   true,
 		HttpConfig: &GoCQHttpConfig{

global/ratelimit.go

@@ -0,0 +1,20 @@
+package global
+
+import (
+	"context"
+	"golang.org/x/time/rate"
+)
+
+var limiter *rate.Limiter
+var limitEnable = false
+
+func RateLimit(ctx context.Context) {
+	if limitEnable {
+		_ = limiter.Wait(ctx)
+	}
+}
+
+func InitLimiter(r float64, b int) {
+	limitEnable = true
+	limiter = rate.NewLimiter(rate.Limit(r), b)
+}

go.mod

@@ -26,5 +26,6 @@ require (
 	github.com/yinghau76/go-ascii-art v0.0.0-20190517192627-e7f465a30189
 	golang.org/x/net v0.0.0-20200904194848-62affa334b73 // indirect
 	golang.org/x/sys v0.0.0-20200905004654-be1d3432aa8f // indirect
+	golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e
 	gopkg.in/yaml.v2 v2.3.0 // indirect
 )

go.sum

@@ -159,6 +159,8 @@ golang.org/x/sys v0.0.0-20200905004654-be1d3432aa8f h1:Fqb3ao1hUmOR3GkUOg/Y+BadL
 golang.org/x/sys v0.0.0-20200905004654-be1d3432aa8f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
+golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e h1:EHBhcS0mlXEAVwNyO2dLfjToGsyY4j24pTs2ScHnX7s=
+golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY=

main.go

@@ -263,6 +263,9 @@ func main() {
 	} else {
 		coolq.SetMessageFormat(conf.PostMessageFormat)
 	}
+	if conf.RateLimit.Enabled {
+		global.InitLimiter(conf.RateLimit.Frequency, conf.RateLimit.BucketSize)
+	}
 	coolq.IgnoreInvalidCQCode = conf.IgnoreInvalidCQCode
 	coolq.ForceFragmented = conf.ForceFragmented
 	if conf.HttpConfig != nil && conf.HttpConfig.Enabled {

server/http.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"crypto/hmac"
 	"crypto/sha1"
 	"encoding/hex"
@@ -74,93 +75,7 @@ func (s *httpServer) Run(addr, authToken string, bot *coolq.CQBot) {
 		})
 	}
 
-	s.engine.Any("/get_login_info", s.GetLoginInfo)
+	s.engine.Any("/:action", s.HandleActions)
-	s.engine.Any("/get_login_info_async", s.GetLoginInfo)
-
-	s.engine.Any("/get_friend_list", s.GetFriendList)
-	s.engine.Any("/get_friend_list_async", s.GetFriendList)
-
-	s.engine.Any("/get_group_list", s.GetGroupList)
-	s.engine.Any("/get_group_list_async", s.GetGroupList)
-
-	s.engine.Any("/get_group_info", s.GetGroupInfo)
-	s.engine.Any("/get_group_info_async", s.GetGroupInfo)
-
-	s.engine.Any("/get_group_member_list", s.GetGroupMemberList)
-	s.engine.Any("/get_group_member_list_async", s.GetGroupMemberList)
-
-	s.engine.Any("/get_group_member_info", s.GetGroupMemberInfo)
-	s.engine.Any("/get_group_member_info_async", s.GetGroupMemberInfo)
-
-	s.engine.Any("/send_msg", s.SendMessage)
-	s.engine.Any("/send_msg_async", s.SendMessage)
-
-	s.engine.Any("/send_private_msg", s.SendPrivateMessage)
-	s.engine.Any("/send_private_msg_async", s.SendPrivateMessage)
-
-	s.engine.Any("/send_group_msg", s.SendGroupMessage)
-	s.engine.Any("/send_group_msg_async", s.SendGroupMessage)
-
-	s.engine.Any("/send_group_forward_msg", s.SendGroupForwardMessage)
-	s.engine.Any("/send_group_forward_msg_async", s.SendGroupForwardMessage)
-
-	s.engine.Any("/delete_msg", s.DeleteMessage)
-	s.engine.Any("/delete_msg_async", s.DeleteMessage)
-
-	s.engine.Any("/set_friend_add_request", s.ProcessFriendRequest)
-	s.engine.Any("/set_friend_add_request_async", s.ProcessFriendRequest)
-
-	s.engine.Any("/set_group_add_request", s.ProcessGroupRequest)
-	s.engine.Any("/set_group_add_request_async", s.ProcessGroupRequest)
-
-	s.engine.Any("/set_group_card", s.SetGroupCard)
-	s.engine.Any("/set_group_card_async", s.SetGroupCard)
-
-	s.engine.Any("/set_group_special_title", s.SetSpecialTitle)
-	s.engine.Any("/set_group_special_title_async", s.SetSpecialTitle)
-
-	s.engine.Any("/set_group_kick", s.SetGroupKick)
-	s.engine.Any("/set_group_kick_async", s.SetGroupKick)
-
-	s.engine.Any("/set_group_ban", s.SetGroupBan)
-	s.engine.Any("/set_group_ban_async", s.SetGroupBan)
-
-	s.engine.Any("/set_group_whole_ban", s.SetWholeBan)
-	s.engine.Any("/set_group_whole_ban_async", s.SetWholeBan)
-
-	s.engine.Any("/set_group_name", s.SetGroupName)
-	s.engine.Any("/set_group_name_async", s.SetGroupName)
-
-	s.engine.Any("/_send_group_notice", s.SendGroupNotice)
-	s.engine.Any("/_send_group_notice_async", s.SendGroupNotice)
-
-	s.engine.Any("/set_group_leave", s.SetGroupLeave)
-	s.engine.Any("/set_group_leave_async", s.SetGroupLeave)
-
-	s.engine.Any("/get_image", s.GetImage)
-
-	s.engine.Any("/get_forward_msg", s.GetForwardMessage)
-
-	s.engine.Any("/get_group_msg", s.GetGroupMessage)
-
-	s.engine.Any("/get_group_honor_info", s.GetGroupHonorInfo)
-
-	s.engine.Any("/can_send_image", s.CanSendImage)
-	s.engine.Any("/can_send_image_async", s.CanSendImage)
-
-	s.engine.Any("/can_send_record", s.CanSendRecord)
-	s.engine.Any("/can_send_record_async", s.CanSendRecord)
-
-	s.engine.Any("/get_status", s.GetStatus)
-	s.engine.Any("/get_status_async", s.GetStatus)
-
-	s.engine.Any("/get_version_info", s.GetVersionInfo)
-	s.engine.Any("/get_version_info_async", s.GetVersionInfo)
-
-	s.engine.Any("/_get_vip_info", s.GetVipInfo)
-	s.engine.Any("/_get_vip_info_async", s.GetVipInfo)
-
-	s.engine.Any("/.handle_quick_operation", s.HandleQuickOperation)
 
 	go func() {
 		log.Infof("CQ HTTP 服务器已启动: %v", addr)
@@ -213,6 +128,17 @@ func (c *httpClient) onBotPushEvent(m coolq.MSG) {
 	}
 }
 
+func (s *httpServer) HandleActions(c *gin.Context) {
+	global.RateLimit(context.Background())
+	action := strings.ReplaceAll(c.Param("action"), "_async", "")
+	log.Debugf("HTTPServer接收到API调用: %v", action)
+	if f, ok := httpApi[action]; ok {
+		f(s, c)
+	} else {
+		c.JSON(200, coolq.Failed(404))
+	}
+}
+
 func (s *httpServer) GetLoginInfo(c *gin.Context) {
 	c.JSON(200, s.bot.CQGetLoginInfo())
 }
@@ -455,3 +381,99 @@ func getParamWithType(c *gin.Context, k string) (string, gjson.Type) {
 	}
 	return "", gjson.Null
 }
+
+var httpApi = map[string]func(s *httpServer, c *gin.Context){
+	"get_login_info": func(s *httpServer, c *gin.Context) {
+		s.GetLoginInfo(c)
+	},
+	"get_friend_list": func(s *httpServer, c *gin.Context) {
+		s.GetFriendList(c)
+	},
+	"get_group_list": func(s *httpServer, c *gin.Context) {
+		s.GetGroupList(c)
+	},
+	"get_group_info": func(s *httpServer, c *gin.Context) {
+		s.GetGroupInfo(c)
+	},
+	"get_group_member_list": func(s *httpServer, c *gin.Context) {
+		s.GetGroupMemberList(c)
+	},
+	"get_group_member_info": func(s *httpServer, c *gin.Context) {
+		s.GetGroupMemberInfo(c)
+	},
+	"send_msg": func(s *httpServer, c *gin.Context) {
+		s.SendMessage(c)
+	},
+	"send_group_msg": func(s *httpServer, c *gin.Context) {
+		s.SendGroupMessage(c)
+	},
+	"send_group_forward_msg": func(s *httpServer, c *gin.Context) {
+		s.SendGroupForwardMessage(c)
+	},
+	"send_private_msg": func(s *httpServer, c *gin.Context) {
+		s.SendPrivateMessage(c)
+	},
+	"delete_msg": func(s *httpServer, c *gin.Context) {
+		s.DeleteMessage(c)
+	},
+	"set_friend_add_request": func(s *httpServer, c *gin.Context) {
+		s.ProcessFriendRequest(c)
+	},
+	"set_group_add_request": func(s *httpServer, c *gin.Context) {
+		s.ProcessGroupRequest(c)
+	},
+	"set_group_card": func(s *httpServer, c *gin.Context) {
+		s.SetGroupCard(c)
+	},
+	"set_group_special_title": func(s *httpServer, c *gin.Context) {
+		s.SetSpecialTitle(c)
+	},
+	"set_group_kick": func(s *httpServer, c *gin.Context) {
+		s.SetGroupKick(c)
+	},
+	"set_group_ban": func(s *httpServer, c *gin.Context) {
+		s.SetGroupBan(c)
+	},
+	"set_group_whole_ban": func(s *httpServer, c *gin.Context) {
+		s.SetWholeBan(c)
+	},
+	"set_group_name": func(s *httpServer, c *gin.Context) {
+		s.SetGroupName(c)
+	},
+	"_send_group_notice": func(s *httpServer, c *gin.Context) {
+		s.SendGroupNotice(c)
+	},
+	"set_group_leave": func(s *httpServer, c *gin.Context) {
+		s.SetGroupLeave(c)
+	},
+	"get_image": func(s *httpServer, c *gin.Context) {
+		s.GetImage(c)
+	},
+	"get_forward_msg": func(s *httpServer, c *gin.Context) {
+		s.GetForwardMessage(c)
+	},
+	"get_group_msg": func(s *httpServer, c *gin.Context) {
+		s.GetGroupMessage(c)
+	},
+	"get_group_honor_info": func(s *httpServer, c *gin.Context) {
+		s.GetGroupHonorInfo(c)
+	},
+	"can_send_image": func(s *httpServer, c *gin.Context) {
+		s.CanSendImage(c)
+	},
+	"can_send_record": func(s *httpServer, c *gin.Context) {
+		s.CanSendRecord(c)
+	},
+	"get_status": func(s *httpServer, c *gin.Context) {
+		s.GetStatus(c)
+	},
+	"get_version_info": func(s *httpServer, c *gin.Context) {
+		s.GetVersionInfo(c)
+	},
+	"_get_vip_info": func(s *httpServer, c *gin.Context) {
+		s.GetVipInfo(c)
+	},
+	".handle_quick_operation": func(s *httpServer, c *gin.Context) {
+		s.HandleQuickOperation(c)
+	},
+}

server/websocket.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"strconv"
@@ -319,7 +320,7 @@ func (c *websocketConn) handleRequest(bot *coolq.CQBot, payload []byte) {
 			c.Close()
 		}
 	}()
-
+	global.RateLimit(context.Background())
 	j := gjson.ParseBytes(payload)
 	t := strings.ReplaceAll(j.Get("action").Str, "_async", "")
 	log.Debugf("WS接收到API调用: %v 参数: %v", t, j.Get("params").Raw)