支持限流

README.md

@@ -74,3 +74,4 @@ go get -u github.com/CuteReimu/mirai-sdk-http
   - [x] 连接与认证
   - [ ] 断线重连
   - [ ] MiraiCode解析
+  - [x] 请求限流

examples/echo.go

@@ -2,11 +2,17 @@ package main
 
 import (
 	. "github.com/CuteReimu/mirai-sdk-http"
+	"golang.org/x/time/rate"
 	"log/slog"
 )
 
 func main() {
-	b, _ := Connect("localhost", 8080, WsChannelAll, "ABCDEFGHIJK", 123456789, false)
+	b, err := Connect("localhost", 8080, WsChannelAll, "ABCDEFGHIJK", 123456789, false)
+	if err != nil {
+		panic(err)
+	}
+	// 设置限流策略为：令牌桶容量为10，每秒放入一个令牌，超过的消息直接丢弃
+	b.SetLimiter("drop", rate.NewLimiter(1, 10))
 	b.ListenGroupMessage(func(message *GroupMessage) bool {
 		var ret MessageChain
 		ret = append(ret, &Plain{Text: "你说了：\n"})

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/gorilla/websocket v1.5.1
 	github.com/stretchr/testify v1.9.0
 	github.com/tidwall/gjson v1.17.1
+	golang.org/x/time v0.5.0
 )
 
 require (

go.sum

@@ -1,5 +1,3 @@
-github.com/CuteReimu/goutil v0.0.0-20240312041931-d57e0398e8d8 h1:9lZfJGQNQkJGc5mk7nYUu5j4l3ELS8WexeFxv1FJOQs=
-github.com/CuteReimu/goutil v0.0.0-20240312041931-d57e0398e8d8/go.mod h1:+ic3kSXScM0/rxJdO/gGydQKuOO9Y1f/t0ajl1xRfeY=
 github.com/CuteReimu/goutil v0.0.0-20240422034935-940ce604f878 h1:DdooD0YHFR3Aoay37mJpUUcKSQ3WUytbfVYYTmocvUI=
 github.com/CuteReimu/goutil v0.0.0-20240422034935-940ce604f878/go.mod h1:+ic3kSXScM0/rxJdO/gGydQKuOO9Y1f/t0ajl1xRfeY=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
@@ -17,10 +15,10 @@ github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JT
 github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
 github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4=
 github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
-golang.org/x/net v0.23.0 h1:7EYJ93RZ9vYSZAIb2x3lnuvqO5zneoD6IvWjuhfxjTs=
-golang.org/x/net v0.23.0/go.mod h1:JKghWKKOSdJwpW2GEx0Ja7fmaKnMsbu+MWVZTokSYmg=
 golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac=
 golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
+golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
+golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

miraihttp.go

@@ -1,12 +1,14 @@
 package miraihttp
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"github.com/CuteReimu/goutil"
 	"github.com/gorilla/websocket"
 	"github.com/tidwall/gjson"
+	"golang.org/x/time/rate"
 	"log/slog"
 	"runtime/debug"
 	"strconv"
@@ -123,10 +125,33 @@ type Bot struct {
 	handler     map[string][]listenHandler
 	syncIdMap   sync.Map
 	eventChan   *goutil.BlockingQueue[func()]
+	limiter     atomic.Pointer[limiter]
+}
+
+type limiter struct {
+	limiterType string
+	limiter     *rate.Limiter
+}
+
+func (l *limiter) check() bool {
+	if l.limiterType == "wait" {
+		return l.limiter.Wait(context.Background()) == nil
+	} else {
+		return l.limiter.Allow()
+	}
+}
+
+// SetLimiter 设置限流器，limiterType为"wait"表示等待，为"drop"表示丢弃
+func (b *Bot) SetLimiter(limiterType string, l *rate.Limiter) {
+	b.limiter.Store(&limiter{limiterType: limiterType, limiter: l})
 }
 
 // request 发送请求
 func (b *Bot) request(command, subCommand string, m any) (gjson.Result, error) {
+	limiter := b.limiter.Load()
+	if limiter != nil && !limiter.check() {
+		return gjson.Result{}, errors.New("rate limit exceeded")
+	}
 	msg := &requestMessage{
 		SyncId:     b.syncId.Add(1),
 		Command:    command,