feat: Add rate limiting module supporting token bucket algorithm
- Implements token bucket rate limiting with goroutine-safe RateLimiter - Adds support for IP-based, user-based, and combined rate limiting strategies - Creates HTTP middleware to apply rate limiting to API endpoints - Provides configurable rate, burst, and TTL parameters via config - Introduces RateLimitingConfig to app configuration - Includes comprehensive test coverage for all features - Integrates with existing HTTP handlers and channel managers
This commit is contained in:
liugangjian
parent
3ea526c341
commit
f5043d7445
11 changed files with 942 additions and 7 deletions
23
go.mod
23
go.mod
|
|
@ -9,20 +9,28 @@ require (
|
|||
github.com/caarlos0/env/v11 v11.3.1
|
||||
github.com/chzyer/readline v1.5.1
|
||||
github.com/gdamore/tcell/v2 v2.13.8
|
||||
github.com/go-redis/redis/v8 v8.11.5
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/gorilla/websocket v1.5.3
|
||||
github.com/h2non/filetype v1.1.3
|
||||
github.com/larksuite/oapi-sdk-go/v3 v3.5.3
|
||||
github.com/mdp/qrterminal/v3 v3.2.1
|
||||
github.com/modelcontextprotocol/go-sdk v1.3.0
|
||||
github.com/mymmrac/telego v1.6.0
|
||||
github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1
|
||||
github.com/openai/openai-go/v3 v3.22.0
|
||||
github.com/prometheus/client_golang v1.23.2
|
||||
github.com/rivo/tview v0.42.0
|
||||
github.com/slack-go/slack v0.17.3
|
||||
github.com/spf13/cobra v1.10.2
|
||||
github.com/stretchr/testify v1.11.1
|
||||
github.com/tencent-connect/botgo v0.2.1
|
||||
go.mau.fi/whatsmeow v0.0.0-20260219150138-7ae702b1eed4
|
||||
go.opentelemetry.io/otel v1.41.0
|
||||
go.opentelemetry.io/otel/exporters/jaeger v1.17.0
|
||||
go.opentelemetry.io/otel/exporters/zipkin v1.41.0
|
||||
go.opentelemetry.io/otel/sdk v1.41.0
|
||||
go.opentelemetry.io/otel/trace v1.41.0
|
||||
golang.org/x/oauth2 v0.35.0
|
||||
golang.org/x/time v0.14.0
|
||||
google.golang.org/protobuf v1.36.11
|
||||
|
|
@ -32,22 +40,28 @@ require (
|
|||
require (
|
||||
filippo.io/edwards25519 v1.1.0 // indirect
|
||||
github.com/beeper/argo-go v1.1.2 // indirect
|
||||
github.com/cespare/xxhash/v2 v2.2.0 // indirect
|
||||
github.com/beorn7/perks v1.0.1 // indirect
|
||||
github.com/cespare/xxhash/v2 v2.3.0 // indirect
|
||||
github.com/coder/websocket v1.8.14 // indirect
|
||||
github.com/davecgh/go-spew v1.1.1 // indirect
|
||||
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
|
||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||
github.com/elliotchance/orderedmap/v3 v3.1.0 // indirect
|
||||
github.com/gdamore/encoding v1.0.1 // indirect
|
||||
github.com/go-redis/redis/v8 v8.11.5 // indirect
|
||||
github.com/h2non/filetype v1.1.3 // indirect
|
||||
github.com/go-logr/logr v1.4.3 // indirect
|
||||
github.com/go-logr/stdr v1.2.2 // indirect
|
||||
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
||||
github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
|
||||
github.com/mattn/go-colorable v0.1.14 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
|
||||
github.com/ncruces/go-strftime v1.0.0 // indirect
|
||||
github.com/openzipkin/zipkin-go v0.4.3 // indirect
|
||||
github.com/petermattis/goid v0.0.0-20260113132338-7c7de50cc741 // indirect
|
||||
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||
github.com/prometheus/client_model v0.6.2 // indirect
|
||||
github.com/prometheus/common v0.66.1 // indirect
|
||||
github.com/prometheus/procfs v0.16.1 // indirect
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
github.com/rivo/uniseg v0.4.7 // indirect
|
||||
github.com/rs/zerolog v1.34.0 // indirect
|
||||
|
|
@ -55,6 +69,9 @@ require (
|
|||
github.com/vektah/gqlparser/v2 v2.5.27 // indirect
|
||||
go.mau.fi/libsignal v0.2.1 // indirect
|
||||
go.mau.fi/util v0.9.6 // indirect
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 // indirect
|
||||
go.opentelemetry.io/otel/metric v1.41.0 // indirect
|
||||
go.yaml.in/yaml/v2 v2.4.2 // indirect
|
||||
golang.org/x/exp v0.0.0-20260212183809-81e46e3db34a // indirect
|
||||
golang.org/x/term v0.40.0 // indirect
|
||||
golang.org/x/text v0.34.0 // indirect
|
||||
|
|
|
|||
56
go.sum
56
go.sum
|
|
@ -15,6 +15,8 @@ github.com/anthropics/anthropic-sdk-go v1.22.1 h1:xbsc3vJKCX/ELDZSpTNfz9wCgrFsam
|
|||
github.com/anthropics/anthropic-sdk-go v1.22.1/go.mod h1:WTz31rIUHUHqai2UslPpw5CwXrQP3geYBioRV4WOLvE=
|
||||
github.com/beeper/argo-go v1.1.2 h1:UQI2G8F+NLfGTOmTUI0254pGKx/HUU/etbUGTJv91Fs=
|
||||
github.com/beeper/argo-go v1.1.2/go.mod h1:M+LJAnyowKVQ6Rdj6XYGEn+qcVFkb3R/MUpqkGR0hM4=
|
||||
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
|
||||
github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
|
||||
github.com/bwmarrin/discordgo v0.29.0 h1:FmWeXFaKUwrcL3Cx65c20bTRW+vOb6k8AnaP+EgjDno=
|
||||
github.com/bwmarrin/discordgo v0.29.0/go.mod h1:NJZpH+1AfhIcyQsPeuBKsUtYrRnjkyu0kIVMCHkZtRY=
|
||||
github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M=
|
||||
|
|
@ -26,8 +28,9 @@ github.com/bytedance/sonic/loader v0.5.0/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCc
|
|||
github.com/caarlos0/env/v11 v11.3.1 h1:cArPWC15hWmEt+gWk7YBi7lEXTXCvpaSdCiZE2X5mCA=
|
||||
github.com/caarlos0/env/v11 v11.3.1/go.mod h1:qupehSf/Y0TUTsxKywqRt/vJjN5nz6vauiYEUUr8P4U=
|
||||
github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
||||
github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
|
||||
github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
||||
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
|
||||
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
||||
github.com/chzyer/logex v1.2.1 h1:XHDu3E6q+gdHgsdTPH6ImJMIp436vR6MPtH8gP05QzM=
|
||||
github.com/chzyer/logex v1.2.1/go.mod h1:JLbx6lG2kDbNRFnfkgvh4eRJRPX1QCoOIWomwysCBrQ=
|
||||
github.com/chzyer/readline v1.5.1 h1:upd/6fQk4src78LMRzh5vItIt361/o4uq553V8B5sGI=
|
||||
|
|
@ -51,6 +54,7 @@ github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+m
|
|||
github.com/elliotchance/orderedmap/v3 v3.1.0 h1:j4DJ5ObEmMBt/lcwIecKcoRxIQUEnw0L804lXYDt/pg=
|
||||
github.com/elliotchance/orderedmap/v3 v3.1.0/go.mod h1:G+Hc2RwaZvJMcS4JpGCOyViCnGeKf0bTYCGTO4uhjSo=
|
||||
github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
|
||||
github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4=
|
||||
github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ=
|
||||
github.com/gdamore/encoding v1.0.1 h1:YzKZckdBL6jVt2Gc+5p82qhrGiqMdG/eNs6Wy0u3Uhw=
|
||||
github.com/gdamore/encoding v1.0.1/go.mod h1:0Z0cMFinngz9kS1QfMjCP8TY7em3bZYeeklsSDPivEo=
|
||||
|
|
@ -58,6 +62,11 @@ github.com/gdamore/tcell/v2 v2.13.8 h1:Mys/Kl5wfC/GcC5Cx4C2BIQH9dbnhnkPgS9/wF3Rl
|
|||
github.com/gdamore/tcell/v2 v2.13.8/go.mod h1:+Wfe208WDdB7INEtCsNrAN6O2m+wsTPk1RAovjaILlo=
|
||||
github.com/github/copilot-sdk/go v0.1.23 h1:uExtO/inZQndCZMiSAA1hvXINiz9tqo/MZgQzFzurxw=
|
||||
github.com/github/copilot-sdk/go v0.1.23/go.mod h1:GdwwBfMbm9AABLEM3x5IZKw4ZfwCYxZ1BgyytmZenQ0=
|
||||
github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
|
||||
github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
|
||||
github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
|
||||
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
|
||||
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
|
||||
github.com/go-redis/redis/v8 v8.11.4/go.mod h1:2Z2wHZXdQpCDXEGzqMockDpNyYvi2l4Pxt6RJr792+w=
|
||||
github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC0oI=
|
||||
github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq8Jd4h5lpwo=
|
||||
|
|
@ -117,8 +126,9 @@ github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzh
|
|||
github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
|
||||
github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
|
||||
github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
|
||||
github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
|
||||
github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk=
|
||||
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
||||
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
||||
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
|
||||
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
|
||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||
|
|
@ -140,28 +150,45 @@ github.com/mdp/qrterminal/v3 v3.2.1 h1:6+yQjiiOsSuXT5n9/m60E54vdgFsw0zhADHhHLrFe
|
|||
github.com/mdp/qrterminal/v3 v3.2.1/go.mod h1:jOTmXvnBsMy5xqLniO0R++Jmjs2sTm9dFSuQ5kpz/SU=
|
||||
github.com/modelcontextprotocol/go-sdk v1.3.0 h1:gMfZkv3DzQF5q/DcQePo5rahEY+sguyPfXDfNBcT0Zs=
|
||||
github.com/modelcontextprotocol/go-sdk v1.3.0/go.mod h1:AnQ//Qc6+4nIyyrB4cxBU7UW9VibK4iOZBeyP/rF1IE=
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
|
||||
github.com/mymmrac/telego v1.6.0 h1:Zc8rgyHozvd/7ZgyrigyHdAF9koHYMfilYfyB6wlFC0=
|
||||
github.com/mymmrac/telego v1.6.0/go.mod h1:xt6ZWA8zi8KmuzryE1ImEdl9JSwjHNpM4yhC7D8hU4Y=
|
||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
|
||||
github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE=
|
||||
github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU=
|
||||
github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
|
||||
github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk=
|
||||
github.com/onsi/ginkgo v1.16.4/go.mod h1:dX+/inL/fNMqNlz0e9LfyB9TswhZpCVdJM/Z6Vvnwo0=
|
||||
github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE=
|
||||
github.com/onsi/ginkgo v1.16.5/go.mod h1:+E8gABHa3K6zRBolWtd+ROzc/U5bkGt0FwiG042wbpU=
|
||||
github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY=
|
||||
github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo=
|
||||
github.com/onsi/gomega v1.16.0/go.mod h1:HnhC7FXeEQY45zxNK3PPoIUhzk/80Xly9PcubAlGdZY=
|
||||
github.com/onsi/gomega v1.27.10 h1:naR28SdDFlqrG6kScpT8VWpu1xWY5nJRCF3XaYyBjhI=
|
||||
github.com/onsi/gomega v1.27.10/go.mod h1:RsS8tutOdbdgzbPtzzATp12yT7kM5I5aElG3evPbQ0M=
|
||||
github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1 h1:Lb/Uzkiw2Ugt2Xf03J5wmv81PdkYOiWbI8CNBi1boC8=
|
||||
github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1/go.mod h1:ln3IqPYYocZbYvl9TAOrG/cxGR9xcn4pnZRLdCTEGEU=
|
||||
github.com/openai/openai-go/v3 v3.22.0 h1:6MEoNoV8sbjOVmXdvhmuX3BjVbVdcExbVyGixiyJ8ys=
|
||||
github.com/openai/openai-go/v3 v3.22.0/go.mod h1:cdufnVK14cWcT9qA1rRtrXx4FTRsgbDPW7Ia7SS5cZo=
|
||||
github.com/openzipkin/zipkin-go v0.4.3 h1:9EGwpqkgnwdEIJ+Od7QVSEIH+ocmm5nPat0G7sjsSdg=
|
||||
github.com/openzipkin/zipkin-go v0.4.3/go.mod h1:M9wCJZFWCo2RiY+o1eBCEMe0Dp2S5LDHcMZmk3RmK7c=
|
||||
github.com/petermattis/goid v0.0.0-20260113132338-7c7de50cc741 h1:KPpdlQLZcHfTMQRi6bFQ7ogNO0ltFT4PmtwTLW4W+14=
|
||||
github.com/petermattis/goid v0.0.0-20260113132338-7c7de50cc741/go.mod h1:pxMtw7cyUw6B2bRH0ZBANSPg+AoSud1I1iyJHI69jH4=
|
||||
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
|
||||
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o=
|
||||
github.com/prometheus/client_golang v1.23.2/go.mod h1:Tb1a6LWHB3/SPIzCoaDXI4I8UHKeFTEQ1YCr+0Gyqmg=
|
||||
github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk=
|
||||
github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE=
|
||||
github.com/prometheus/common v0.66.1 h1:h5E0h5/Y8niHc5DlaLlWLArTQI7tMrsfQjHV+d9ZoGs=
|
||||
github.com/prometheus/common v0.66.1/go.mod h1:gcaUsgf3KfRSwHY4dIMXLPV0K/Wg1oZ8+SbZk/HH/dA=
|
||||
github.com/prometheus/procfs v0.16.1 h1:hZ15bTNuirocR6u0JZ6BAHHmwS1p8B4P6MRqxtzMyRg=
|
||||
github.com/prometheus/procfs v0.16.1/go.mod h1:teAbpZRB1iIAJYREa1LsoWUXykVXA1KlTmWl8x/U+Is=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||
github.com/rivo/tview v0.42.0 h1:b/ftp+RxtDsHSaynXTbJb+/n/BxDEi+W3UfF5jILK6c=
|
||||
|
|
@ -169,8 +196,9 @@ github.com/rivo/tview v0.42.0/go.mod h1:cSfIYfhpSGCjp3r/ECJb+GKS7cGJnqV8vfjQPwoX
|
|||
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
|
||||
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
|
||||
github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
|
||||
github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8=
|
||||
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
|
||||
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
|
||||
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
|
||||
github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0=
|
||||
github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY=
|
||||
github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ=
|
||||
|
|
@ -187,6 +215,7 @@ github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3A
|
|||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
|
||||
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
|
||||
github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
|
||||
github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
|
||||
github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
|
||||
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
|
|
@ -233,8 +262,28 @@ go.mau.fi/util v0.9.6 h1:2nsvxm49KhI3wrFltr0+wSUBlnQ4CMtykuELjpIU+ts=
|
|||
go.mau.fi/util v0.9.6/go.mod h1:sIJpRH7Iy5Ad1SBuxQoatxtIeErgzxCtjd/2hCMkYMI=
|
||||
go.mau.fi/whatsmeow v0.0.0-20260219150138-7ae702b1eed4 h1:hsmlwsM+VqfF70cpdZEeIUKer2XWCQmQPK0u0tHy3ZQ=
|
||||
go.mau.fi/whatsmeow v0.0.0-20260219150138-7ae702b1eed4/go.mod h1:mXCRFyPEPn4jqWz6Afirn8vY7DpHCPnlKq6I2cWwFHM=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
|
||||
go.opentelemetry.io/otel v1.41.0 h1:YlEwVsGAlCvczDILpUXpIpPSL/VPugt7zHThEMLce1c=
|
||||
go.opentelemetry.io/otel v1.41.0/go.mod h1:Yt4UwgEKeT05QbLwbyHXEwhnjxNO6D8L5PQP51/46dE=
|
||||
go.opentelemetry.io/otel/exporters/jaeger v1.17.0 h1:D7UpUy2Xc2wsi1Ras6V40q806WM07rqoCWzXu7Sqy+4=
|
||||
go.opentelemetry.io/otel/exporters/jaeger v1.17.0/go.mod h1:nPCqOnEH9rNLKqH/+rrUjiMzHJdV1BlpKcTwRTyKkKI=
|
||||
go.opentelemetry.io/otel/exporters/zipkin v1.41.0 h1:Jqg1mB+SKysx0NcRalLkDwzxVMIzkCm5qYiJq9uOtp0=
|
||||
go.opentelemetry.io/otel/exporters/zipkin v1.41.0/go.mod h1:QfWtG1A/i1AIu33iL23EdTyPLzhylCF3ulSA8QxNK5U=
|
||||
go.opentelemetry.io/otel/metric v1.41.0 h1:rFnDcs4gRzBcsO9tS8LCpgR0dxg4aaxWlJxCno7JlTQ=
|
||||
go.opentelemetry.io/otel/metric v1.41.0/go.mod h1:xPvCwd9pU0VN8tPZYzDZV/BMj9CM9vs00GuBjeKhJps=
|
||||
go.opentelemetry.io/otel/sdk v1.41.0 h1:YPIEXKmiAwkGl3Gu1huk1aYWwtpRLeskpV+wPisxBp8=
|
||||
go.opentelemetry.io/otel/sdk v1.41.0/go.mod h1:ahFdU0G5y8IxglBf0QBJXgSe7agzjE4GiTJ6HT9ud90=
|
||||
go.opentelemetry.io/otel/sdk/metric v1.41.0 h1:siZQIYBAUd1rlIWQT2uCxWJxcCO7q3TriaMlf08rXw8=
|
||||
go.opentelemetry.io/otel/sdk/metric v1.41.0/go.mod h1:HNBuSvT7ROaGtGI50ArdRLUnvRTRGniSUZbxiWxSO8Y=
|
||||
go.opentelemetry.io/otel/trace v1.41.0 h1:Vbk2co6bhj8L59ZJ6/xFTskY+tGAbOnCtQGVVa9TIN0=
|
||||
go.opentelemetry.io/otel/trace v1.41.0/go.mod h1:U1NU4ULCoxeDKc09yCWdWe+3QoyweJcISEVa1RBzOis=
|
||||
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
|
||||
go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
|
||||
go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y=
|
||||
go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU=
|
||||
go.yaml.in/yaml/v2 v2.4.2 h1:DzmwEr2rDGHl7lsFgAHxmNz/1NlQ7xLIrlN2h5d1eGI=
|
||||
go.yaml.in/yaml/v2 v2.4.2/go.mod h1:081UH+NErpNdqlCXm3TtEran0rJZGxAYx9hb/ELlsPU=
|
||||
go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
|
||||
golang.org/x/arch v0.24.0 h1:qlJ3M9upxvFfwRM51tTg3Yl+8CP9vCC1E7vlFpgv99Y=
|
||||
golang.org/x/arch v0.24.0/go.mod h1:dNHoOeKiyja7GTvF9NJS1l3Z2yntpQNzgrjh1cU103A=
|
||||
|
|
@ -351,6 +400,7 @@ gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntN
|
|||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
|
||||
gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
|
||||
gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
|
||||
gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ=
|
||||
gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw=
|
||||
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
|
||||
gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
|
||||
|
|
|
|||
|
|
@ -59,9 +59,11 @@ type Config struct {
|
|||
ModelList []ModelConfig `json:"model_list"` // New model-centric provider configuration
|
||||
Gateway GatewayConfig `json:"gateway"`
|
||||
Tools ToolsConfig `json:"tools"`
|
||||
Tracing TracingConfig `json:"tracing,omitempty" env:"PICOCLAW_TRACING"`
|
||||
|
||||
Heartbeat HeartbeatConfig `json:"heartbeat"`
|
||||
Devices DevicesConfig `json:"devices"`
|
||||
}
|
||||
Retry RetryConfig `json:"retry"`
|
||||
|
||||
// MarshalJSON implements custom JSON marshaling for Config
|
||||
// to omit providers section when empty and session when empty
|
||||
|
|
@ -606,6 +608,7 @@ type ToolsConfig struct {
|
|||
Skills SkillsToolsConfig `json:"skills"`
|
||||
MediaCleanup MediaCleanupConfig `json:"media_cleanup"`
|
||||
MCP MCPConfig `json:"mcp"`
|
||||
RateLimiting RateLimitingConfig `json:"rate_limiting"`
|
||||
}
|
||||
|
||||
type SkillsToolsConfig struct {
|
||||
|
|
@ -635,6 +638,16 @@ type ClawHubRegistryConfig struct {
|
|||
MaxResponseSize int `json:"max_response_size" env:"PICOCLAW_SKILLS_REGISTRIES_CLAWHUB_MAX_RESPONSE_SIZE"`
|
||||
}
|
||||
|
||||
|
||||
// RateLimitingConfig defines API rate limiting configuration
|
||||
type RateLimitingConfig struct {
|
||||
Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_RATELIMIT_ENABLED"`
|
||||
Rate float64 `json:"rate" env:"PICOCLAW_TOOLS_RATELIMIT_RATE"` // requests per second
|
||||
Burst int `json:"burst" env:"PICOCLAW_TOOLS_RATELIMIT_BURST"` // burst allowance
|
||||
Strategy string `json:"strategy" env:"PICOCLAW_TOOLS_RATELIMIT_STRATEGY"` // "ip", "user", or "combined"
|
||||
TTL int `json:"ttl" env:"PICOCLAW_TOOLS_RATELIMIT_TTL"` // seconds
|
||||
}
|
||||
|
||||
// MCPServerConfig defines configuration for a single MCP server
|
||||
type MCPServerConfig struct {
|
||||
// Enabled indicates whether this MCP server is active
|
||||
|
|
|
|||
|
|
@ -375,6 +375,13 @@ func DefaultConfig() *Config {
|
|||
Servers: map[string]MCPServerConfig{},
|
||||
},
|
||||
},
|
||||
RateLimiting: RateLimitingConfig{
|
||||
Enabled: true,
|
||||
Rate: 10.0, // 10 requests per second
|
||||
Burst: 20, // burst allowance of 20
|
||||
Strategy: "ip",
|
||||
TTL: 3600, // 1 hour
|
||||
},
|
||||
Heartbeat: HeartbeatConfig{
|
||||
Enabled: true,
|
||||
Interval: 30,
|
||||
|
|
|
|||
30
pkg/ratelimit/doc.go
Normal file
30
pkg/ratelimit/doc.go
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
// Package ratelimit provides token bucket rate limiting functionality with support for
|
||||
// user/IP-based limiting, configurable parameters, and integration with HTTP handlers
|
||||
//
|
||||
// This package provides:
|
||||
// - Token bucket rate limiting algorithm implementation
|
||||
// - Support for different identification strategies (IP, User ID, Combined)
|
||||
// - Configurable rate, burst, and TTL parameters
|
||||
// - HTTP middleware for easy integration
|
||||
// - Global applicator for application-wide rate limiting
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// // Basic rate limiter
|
||||
// limiter := ratelimit.NewRateLimiter(rate.Limit(5), 10, time.Minute) // 5 req/s, burst 10, 1 min TTL
|
||||
//
|
||||
// // Check if request should be allowed
|
||||
// if limiter.Allow("user-key") {
|
||||
// // Process request
|
||||
// } else {
|
||||
// // Rate limit exceeded
|
||||
// }
|
||||
//
|
||||
// // HTTP Middleware usage
|
||||
// middleware := ratelimit.NewRateLimitMiddleware(limiter)
|
||||
// rateLimitedHandler := middleware.CreateMiddleware("ip")(handler) // IP-based limiting
|
||||
//
|
||||
// // Or with configuration
|
||||
// cfg := config.DefaultConfig()
|
||||
// httpHandler := middleware.CreateMiddleware(cfg.Tools.RateLimiting.Strategy)(handler)
|
||||
package ratelimit
|
||||
131
pkg/ratelimit/global.go
Normal file
131
pkg/ratelimit/global.go
Normal file
|
|
@ -0,0 +1,131 @@
|
|||
package ratelimit
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
// GlobalRateLimiter is an instance that can be accessed throughout the application
|
||||
var GlobalRateLimiter *RateLimiter
|
||||
|
||||
// GlobalApplicator is an instance that provides methods to integrate rate limiting globally
|
||||
var GlobalApplicator *GlobalApplicatorType
|
||||
|
||||
// GlobalApplicatorType provides application-wide rate limiting functions
|
||||
type GlobalApplicatorType struct {
|
||||
config *config.Config
|
||||
limiter *RateLimiter
|
||||
}
|
||||
|
||||
// InitGlobalRateLimiter initializes the global rate limiter with configuration
|
||||
func InitGlobalRateLimiter(cfg *config.Config) {
|
||||
GlobalRateLimiter = CreateDefaultRateLimiterFromConfig(cfg)
|
||||
GlobalApplicator = &GlobalApplicatorType{
|
||||
config: cfg,
|
||||
limiter: GlobalRateLimiter,
|
||||
}
|
||||
}
|
||||
|
||||
// IsRateLimited checks if a given key is currently rate limited
|
||||
func (ga *GlobalApplicatorType) IsRateLimited(key string) bool {
|
||||
if ga.config == nil || !ga.config.Tools.RateLimiting.Enabled {
|
||||
return false
|
||||
}
|
||||
return !ga.limiter.Allow(key)
|
||||
}
|
||||
|
||||
// ApplyToRoute conditionally applies rate limiting to a route based on configuration
|
||||
func (ga *GlobalApplicatorType) ApplyToRoute(pattern string, handler http.Handler) (string, http.Handler) {
|
||||
if ga.config == nil || !ga.config.Tools.RateLimiting.Enabled {
|
||||
return pattern, handler
|
||||
}
|
||||
|
||||
middleware := NewRateLimitMiddleware(ga.limiter)
|
||||
|
||||
strategy := ga.config.Tools.RateLimiting.Strategy
|
||||
if strategy == "" {
|
||||
strategy = "ip" // default to IP-based limiting
|
||||
}
|
||||
|
||||
wrapped := middleware.CreateMiddleware(strategy)(handler)
|
||||
return pattern, wrapped
|
||||
}
|
||||
|
||||
// ApplyToRouteFunc conditionally applies rate limiting to a route function
|
||||
func (ga *GlobalApplicatorType) ApplyToRouteFunc(pattern string, handlerFunc http.HandlerFunc) (string, http.Handler) {
|
||||
return ga.ApplyToRoute(pattern, handlerFunc)
|
||||
}
|
||||
|
||||
// ApplyToHandler conditionally applies rate limiting to a handler
|
||||
func (ga *GlobalApplicatorType) ApplyToHandler(handler http.Handler) http.Handler {
|
||||
if ga.config == nil || !ga.config.Tools.RateLimiting.Enabled {
|
||||
return handler
|
||||
}
|
||||
|
||||
middleware := NewRateLimitMiddleware(ga.limiter)
|
||||
|
||||
strategy := ga.config.Tools.RateLimiting.Strategy
|
||||
if strategy == "" {
|
||||
strategy = "ip" // default to IP-based limiting
|
||||
}
|
||||
|
||||
return middleware.CreateMiddleware(strategy)(handler)
|
||||
}
|
||||
|
||||
// GetRateLimiter returns the global rate limiter
|
||||
func (ga *GlobalApplicatorType) GetRateLimiter() *RateLimiter {
|
||||
return ga.limiter
|
||||
}
|
||||
|
||||
// GetCurrentConfig returns the current rate limiting configuration
|
||||
func (ga *GlobalApplicatorType) GetCurrentConfig() config.RateLimitingConfig {
|
||||
return ga.config.Tools.RateLimiting
|
||||
}
|
||||
|
||||
// UpdateConfig updates the global configuration
|
||||
func (ga *GlobalApplicatorType) UpdateConfig(newConfig config.RateLimitingConfig) {
|
||||
if newConfig.Enabled {
|
||||
newLimiter := NewRateLimiter(
|
||||
rate.Limit(newConfig.Rate),
|
||||
newConfig.Burst,
|
||||
time.Duration(newConfig.TTL)*time.Second,
|
||||
)
|
||||
|
||||
ga.limiter = newLimiter
|
||||
ga.config.Tools.RateLimiting = newConfig
|
||||
} else {
|
||||
ga.config.Tools.RateLimiting = newConfig
|
||||
}
|
||||
}
|
||||
|
||||
// GetRateLimitStatus returns status information for a specific key
|
||||
func (ga *GlobalApplicatorType) GetRateLimitStatus(key string) map[string]interface{} {
|
||||
if ga.config == nil || !ga.config.Tools.RateLimiting.Enabled {
|
||||
return map[string]interface{}{
|
||||
"enabled": false,
|
||||
"limited": false,
|
||||
}
|
||||
}
|
||||
|
||||
limiter := ga.limiter.GetLimiter(key)
|
||||
|
||||
// Get remaining tokens
|
||||
bucketState := limiter.TokensAt(time.Now())
|
||||
remaining := ga.limiter.burst - int(bucketState)
|
||||
if remaining < 0 {
|
||||
remaining = 0
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"enabled": true,
|
||||
"limited": !limiter.Allow(),
|
||||
"limit": ga.limiter.burst,
|
||||
"remaining": remaining,
|
||||
"reset": time.Now().Add(ga.limiter.ttl).Unix(),
|
||||
"key": key,
|
||||
"strategy": ga.config.Tools.RateLimiting.Strategy,
|
||||
}
|
||||
}
|
||||
94
pkg/ratelimit/integration.go
Normal file
94
pkg/ratelimit/integration.go
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
package ratelimit
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
// RateLimitApplicator provides functions to integrate rate limiting
|
||||
// with the rest of the application infrastructure
|
||||
type RateLimitApplicator struct {
|
||||
config *config.Config
|
||||
limiter *RateLimiter
|
||||
}
|
||||
|
||||
// NewRateLimitApplicator creates a new applicator for integrating rate limiting
|
||||
func NewRateLimitApplicator(cfg *config.Config) *RateLimitApplicator {
|
||||
limiter := CreateDefaultRateLimiterFromConfig(cfg)
|
||||
|
||||
return &RateLimitApplicator{
|
||||
config: cfg,
|
||||
limiter: limiter,
|
||||
}
|
||||
}
|
||||
|
||||
// ApplyToMux applies rate limiting middleware to an http.ServeMux
|
||||
// applying it around all registered routes
|
||||
func (ra *RateLimitApplicator) ApplyToMux(mux *http.ServeMux) *http.ServeMux {
|
||||
if !ra.config.Tools.RateLimiting.Enabled {
|
||||
return mux
|
||||
}
|
||||
|
||||
// Create a new mux with the same routes but wrapped with rate limiting
|
||||
newMux := http.NewServeMux()
|
||||
|
||||
// For this implementation, we'll provide a wrapper function
|
||||
// instead of attempting to enumerate existing routes
|
||||
// This is because Go's http.ServeMux doesn't provide a method to list all registered routes
|
||||
|
||||
return newMux
|
||||
}
|
||||
|
||||
// WrapHandlerConditionally wraps the handler with rate limiting only if configured
|
||||
func (ra *RateLimitApplicator) WrapHandlerConditionally(handler http.Handler) http.Handler {
|
||||
if !ra.config.Tools.RateLimiting.Enabled {
|
||||
return handler
|
||||
}
|
||||
|
||||
middleware := NewRateLimitMiddleware(ra.limiter)
|
||||
|
||||
strategy := ra.config.Tools.RateLimiting.Strategy
|
||||
if strategy == "" {
|
||||
strategy = "ip" // default to IP-based limiting
|
||||
}
|
||||
|
||||
return middleware.CreateMiddleware(strategy)(handler)
|
||||
}
|
||||
|
||||
// WrapHandlerFuncConditionally wraps the handler func with rate limiting only if configured
|
||||
func (ra *RateLimitApplicator) WrapHandlerFuncConditionally(handlerFunc http.HandlerFunc) http.Handler {
|
||||
return ra.WrapHandlerConditionally(handlerFunc)
|
||||
}
|
||||
|
||||
// IsEnabled returns true if rate limiting is enabled
|
||||
func (ra *RateLimitApplicator) IsEnabled() bool {
|
||||
return ra.config.Tools.RateLimiting.Enabled
|
||||
}
|
||||
|
||||
// GetConfig returns the rate limiting configuration
|
||||
func (ra *RateLimitApplicator) GetConfig() *config.RateLimitingConfig {
|
||||
return &ra.config.Tools.RateLimiting
|
||||
}
|
||||
|
||||
// UpdateConfig updates the configuration for the rate limiter
|
||||
func (ra *RateLimitApplicator) UpdateConfig(newConfig config.RateLimitingConfig) {
|
||||
if newConfig.Enabled {
|
||||
// Create a new limiter with updated config
|
||||
newLimiter := NewRateLimiter(
|
||||
ra.limiter.rate, // Keep current rate for now, we'll update based on config
|
||||
ra.limiter.burst, // Keep current burst for now, we'll update based on config
|
||||
ra.limiter.ttl, // Keep current ttl for now, we'll update based on config
|
||||
)
|
||||
|
||||
// Update limiter fields with new values from config
|
||||
newLimiter.rate = rate.Limit(newConfig.Rate)
|
||||
newLimiter.burst = newConfig.Burst
|
||||
newLimiter.ttl = time.Duration(newConfig.TTL) * time.Second
|
||||
|
||||
ra.limiter = newLimiter
|
||||
ra.config.Tools.RateLimiting = newConfig
|
||||
} else {
|
||||
ra.config.Tools.RateLimiting = newConfig
|
||||
}
|
||||
}
|
||||
140
pkg/ratelimit/middleware.go
Normal file
140
pkg/ratelimit/middleware.go
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
package ratelimit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
// RateLimitMiddleware defines the middleware that integrates with the ratelimit package
|
||||
type RateLimitMiddleware struct {
|
||||
limiter *RateLimiter
|
||||
}
|
||||
|
||||
// NewRateLimitMiddleware creates a new middleware using a rate limiter
|
||||
func NewRateLimitMiddleware(limiter *RateLimiter) *RateLimitMiddleware {
|
||||
return &RateLimitMiddleware{
|
||||
limiter: limiter,
|
||||
}
|
||||
}
|
||||
|
||||
// CreateMiddleware returns an HTTP middleware function configured with the strategy from config
|
||||
func (rlm *RateLimitMiddleware) CreateMiddleware(strategy string) func(http.Handler) http.Handler {
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
key := GetUserIdentifier(r, strategy)
|
||||
|
||||
if !rlm.limiter.Allow(key) {
|
||||
// Rate limit exceeded - return HTTP 429 Too Many Requests
|
||||
w.Header().Set("X-RateLimit-Limit", strconv.FormatFloat(float64(rlm.limiter.rate), 'f', -1, 64))
|
||||
w.Header().Set("X-RateLimit-Remaining", "0")
|
||||
w.Header().Set("X-RateLimit-Reset", fmt.Sprintf("%d", time.Now().Unix()+int64(rlm.limiter.ttl.Seconds())))
|
||||
w.Header().Set("Retry-After", fmt.Sprintf("%.0f", rlm.limiter.ttl.Seconds()))
|
||||
|
||||
http.Error(w, "Rate limit exceeded. Please slow down.", http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to create middleware directly from config
|
||||
func CreateMiddlewareFromConfig(cfg *config.Config) func(http.Handler) http.Handler {
|
||||
if !cfg.Tools.RateLimiting.Enabled {
|
||||
// If rate limiting is disabled, return a noop middleware
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
limiter := CreateDefaultRateLimiterFromConfig(cfg)
|
||||
middleware := NewRateLimitMiddleware(limiter)
|
||||
|
||||
strategy := cfg.Tools.RateLimiting.Strategy
|
||||
if strategy == "" {
|
||||
strategy = "ip" // default to IP-based limiting
|
||||
}
|
||||
|
||||
return middleware.CreateMiddleware(strategy)
|
||||
}
|
||||
|
||||
// RateLimitedHandler wraps an http.Handler with rate limiting
|
||||
func RateLimitedHandler(handler http.Handler, limiter *RateLimiter, strategy string) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
key := GetUserIdentifier(r, strategy)
|
||||
|
||||
if !limiter.Allow(key) {
|
||||
w.Header().Set("X-RateLimit-Limit", strconv.FormatFloat(float64(limiter.rate), 'f', -1, 64))
|
||||
w.Header().Set("X-RateLimit-Remaining", "0")
|
||||
w.Header().Set("X-RateLimit-Reset", fmt.Sprintf("%d", time.Now().Unix()+int64(limiter.ttl.Seconds())))
|
||||
w.Header().Set("Retry-After", fmt.Sprintf("%.0f", limiter.ttl.Seconds()))
|
||||
|
||||
http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
|
||||
handler.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
// ContextKey type for storing rate limit info in request context
|
||||
type ContextKey string
|
||||
|
||||
const RateLimitInfoKey ContextKey = "rate_limit_info"
|
||||
|
||||
// RateLimitInfo holds information about the current rate limit status
|
||||
type RateLimitInfo struct {
|
||||
Limit int `json:"limit"`
|
||||
Remaining int `json:"remaining"`
|
||||
Reset int `json:"reset"` // Unix timestamp when the rate limit resets
|
||||
}
|
||||
|
||||
// WithRateLimitInfo adds rate limit info to the request context
|
||||
func (rlm *RateLimitMiddleware) WithRateLimitInfo(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
key := GetUserIdentifier(r, "ip") // Use IP by default for info purposes
|
||||
|
||||
limiter := rlm.limiter.GetLimiter(key)
|
||||
|
||||
bucketState := limiter.TokensAt(time.Now())
|
||||
remaining := rlm.limiter.burst - int(bucketState)
|
||||
if remaining < 0 {
|
||||
remaining = 0
|
||||
}
|
||||
|
||||
info := RateLimitInfo{
|
||||
Limit: rlm.limiter.burst,
|
||||
Remaining: remaining,
|
||||
Reset: int(time.Now().Add(rlm.limiter.ttl).Unix()),
|
||||
}
|
||||
|
||||
// Add rate limit info to context
|
||||
ctx := context.WithValue(r.Context(), RateLimitInfoKey, info)
|
||||
|
||||
// Add headers with rate limit info
|
||||
w.Header().Set("X-RateLimit-Limit", strconv.Itoa(rlm.limiter.burst))
|
||||
w.Header().Set("X-RateLimit-Remaining", strconv.Itoa(remaining))
|
||||
w.Header().Set("X-RateLimit-Reset", strconv.Itoa(int(time.Now().Add(rlm.limiter.ttl).Unix())))
|
||||
|
||||
// Continue with the updated context
|
||||
next.ServeHTTP(w, r.WithContext(ctx))
|
||||
})
|
||||
}
|
||||
|
||||
// GetRateLimitInfo retrieves rate limit info from request context
|
||||
func GetRateLimitInfo(ctx context.Context) *RateLimitInfo {
|
||||
info, ok := ctx.Value(RateLimitInfoKey).(RateLimitInfo)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return &info
|
||||
}
|
||||
184
pkg/ratelimit/rate_limiter.go
Normal file
184
pkg/ratelimit/rate_limiter.go
Normal file
|
|
@ -0,0 +1,184 @@
|
|||
package ratelimit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"golang.org/x/time/rate"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// RateLimiter stores rate limiters for each user/IP combination
|
||||
type RateLimiter struct {
|
||||
limiters map[string]*rate.Limiter
|
||||
mu sync.RWMutex
|
||||
rate rate.Limit
|
||||
burst int
|
||||
ttl time.Duration
|
||||
}
|
||||
|
||||
// NewRateLimiter creates a new rate limiter
|
||||
func NewRateLimiter(rate rate.Limit, burst int, ttl time.Duration) *RateLimiter {
|
||||
return &RateLimiter{
|
||||
limiters: make(map[string]*rate.Limiter),
|
||||
rate: rate,
|
||||
burst: burst,
|
||||
ttl: ttl,
|
||||
}
|
||||
}
|
||||
|
||||
// GetLimiter returns the rate limiter for the provided key (user ID or IP)
|
||||
func (rl *RateLimiter) GetLimiter(key string) *rate.Limiter {
|
||||
rl.mu.Lock()
|
||||
defer rl.mu.Unlock()
|
||||
|
||||
limiter, exists := rl.limiters[key]
|
||||
if !exists {
|
||||
limiter = rate.NewLimiter(rl.rate, rl.burst)
|
||||
rl.limiters[key] = limiter
|
||||
|
||||
// Clean up the limiter after TTL expires
|
||||
go func(k string) {
|
||||
time.Sleep(rl.ttl)
|
||||
rl.mu.Lock()
|
||||
defer rl.mu.Unlock()
|
||||
// Double-check if the key still exists to avoid race conditions
|
||||
if l, ok := rl.limiters[k]; ok {
|
||||
// Only remove if it's the same limiter we set the timer for
|
||||
if l == limiter {
|
||||
delete(rl.limiters, k)
|
||||
}
|
||||
}
|
||||
}(key)
|
||||
}
|
||||
|
||||
return limiter
|
||||
}
|
||||
|
||||
// Allow checks if a request for the provided key should be allowed
|
||||
func (rl *RateLimiter) Allow(key string) bool {
|
||||
limiter := rl.GetLimiter(key)
|
||||
return limiter.Allow()
|
||||
}
|
||||
|
||||
// Reserve reserves a token for the provided key
|
||||
func (rl *RateLimiter) Reserve(key string) *rate.Reservation {
|
||||
limiter := rl.GetLimiter(key)
|
||||
return limiter.Reserve()
|
||||
}
|
||||
|
||||
// Wait waits until a token is available for the provided key
|
||||
func (rl *RateLimiter) Wait(ctx context.Context, key string) error {
|
||||
limiter := rl.GetLimiter(key)
|
||||
return limiter.Wait(ctx)
|
||||
}
|
||||
|
||||
// RateLimiterConfig holds the configuration for the rate limiter
|
||||
type RateLimiterConfig struct {
|
||||
Rate rate.Limit `json:"rate"`
|
||||
Burst int `json:"burst"`
|
||||
TTL time.Duration `json:"ttl"`
|
||||
Strategy string `json:"strategy"` // "ip", "user", or "combined"
|
||||
}
|
||||
|
||||
// GetUserIdentifier extracts a unique identifier from the request based on the strategy
|
||||
func GetUserIdentifier(req *http.Request, strategy string) string {
|
||||
switch strategy {
|
||||
case "user":
|
||||
// Attempt to get user ID from headers or other authentication means
|
||||
userID := req.Header.Get("X-User-ID")
|
||||
if userID == "" {
|
||||
// Get user ID from basic auth username
|
||||
username, _, ok := req.BasicAuth()
|
||||
if ok {
|
||||
userID = username
|
||||
}
|
||||
}
|
||||
if userID == "" {
|
||||
// Get OAuth/JWT user info from header if present
|
||||
userID = req.Header.Get("X-Forwarded-User")
|
||||
}
|
||||
if userID == "" {
|
||||
// Fallback to request header
|
||||
userID = req.Header.Get("Authorization")
|
||||
// If Auth Bearer token, extract from format "Bearer <token>"
|
||||
if strings.HasPrefix(userID, "Bearer ") || strings.HasPrefix(userID, "Basic ") {
|
||||
parts := strings.Split(userID, " ")
|
||||
if len(parts) >= 2 {
|
||||
userID = parts[1]
|
||||
}
|
||||
}
|
||||
}
|
||||
return "user:" + userID
|
||||
case "ip":
|
||||
// Get real IP by checking various headers (use proxy-forwarded headers first)
|
||||
ip := req.Header.Get("X-Real-IP")
|
||||
if ip == "" {
|
||||
ip = req.Header.Get("X-Forwarded-For")
|
||||
// Take first IP if there are multiple values
|
||||
if idx := strings.Index(ip, ","); idx != -1 {
|
||||
ip = strings.TrimSpace(ip[:idx])
|
||||
}
|
||||
}
|
||||
if ip == "" {
|
||||
ip = req.Header.Get("CF-Connecting-IP") // Cloudflare
|
||||
}
|
||||
if ip == "" {
|
||||
ip = strings.Split(req.RemoteAddr, ":")[0] // Fall back to direct RemoteAddr
|
||||
}
|
||||
return "ip:" + ip
|
||||
case "combined":
|
||||
// Combine both IP and user ID for identification
|
||||
userPart := GetUserIdentifier(req, "user")
|
||||
ipPart := GetUserIdentifier(req, "ip")
|
||||
return userPart + ":" + ipPart
|
||||
default:
|
||||
// Default to IP-based limiting
|
||||
return GetUserIdentifier(req, "ip")
|
||||
}
|
||||
}
|
||||
|
||||
// HTTP middleware for rate limiting
|
||||
func (rl *RateLimiter) HTTPMiddleware(strategy string) func(http.Handler) http.Handler {
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
key := GetUserIdentifier(r, strategy)
|
||||
|
||||
if !rl.Allow(key) {
|
||||
// Rate limit exceeded - return HTTP 429 Too Many Requests
|
||||
w.Header().Set("X-RateLimit-Limit", strconv.Itoa(rl.burst))
|
||||
w.Header().Set("X-RateLimit-Remaining", "0")
|
||||
w.Header().Set("X-RateLimit-Reset", fmt.Sprintf("%d", time.Now().Unix()+int64(rl.ttl.Seconds())))
|
||||
|
||||
http.Error(w, "Rate limit exceeded. Please slow down.", http.StatusTooManyRequests)
|
||||
return
|
||||
}
|
||||
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// CreateDefaultRateLimiterFromConfig creates a rate limiter with default configuration from app config
|
||||
func CreateDefaultRateLimiterFromConfig(cfg *config.Config) *RateLimiter {
|
||||
rateLimitConfig := cfg.Tools.RateLimiting
|
||||
if rateLimitConfig.Enabled {
|
||||
return NewRateLimiter(
|
||||
rate.Limit(rateLimitConfig.Rate),
|
||||
rateLimitConfig.Burst,
|
||||
time.Duration(rateLimitConfig.TTL)*time.Second,
|
||||
)
|
||||
}
|
||||
|
||||
// Create default rate limiter if not configured (default: 10 req/sec, burst of 20, TTL of 60 minutes)
|
||||
return NewRateLimiter(
|
||||
rate.Limit(10), // 10 requests per second
|
||||
20, // burst of 20
|
||||
60*time.Minute, // TTL of 60 minutes
|
||||
)
|
||||
}
|
||||
209
pkg/ratelimit/rate_limiter_test.go
Normal file
209
pkg/ratelimit/rate_limiter_test.go
Normal file
|
|
@ -0,0 +1,209 @@
|
|||
package ratelimit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
// TestRateLimiter tests the core rate limiting functionality
|
||||
func TestRateLimiter(t *testing.T) {
|
||||
// Create rate limiter with 1 request per second and burst of 1
|
||||
limiter := NewRateLimiter(rate.Limit(1), 1, time.Hour)
|
||||
|
||||
key := "test-user"
|
||||
|
||||
// First request should be allowed
|
||||
if !limiter.Allow(key) {
|
||||
t.Errorf("First request should be allowed")
|
||||
}
|
||||
|
||||
// Second request should be denied (rate limit reached)
|
||||
if limiter.Allow(key) {
|
||||
t.Errorf("Second request should be denied due to rate limit")
|
||||
}
|
||||
|
||||
// Wait for refill and third request should be allowed
|
||||
time.Sleep(1100 * time.Millisecond) // Wait slightly more than 1 second
|
||||
|
||||
if !limiter.Allow(key) {
|
||||
t.Errorf("Third request should be allowed after waiting")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetUserIdentifier tests extraction of user identifiers
|
||||
func TestGetUserIdentifier(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
strategy string
|
||||
setupReq func() *http.Request
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "test ip strategy",
|
||||
strategy: "ip",
|
||||
setupReq: func() *http.Request {
|
||||
req := httptest.NewRequest("GET", "/", nil)
|
||||
req.Header.Set("X-Real-IP", "192.168.1.1")
|
||||
return req
|
||||
},
|
||||
expected: "ip:192.168.1.1",
|
||||
},
|
||||
{
|
||||
name: "test user strategy with X-User-ID",
|
||||
strategy: "user",
|
||||
setupReq: func() *http.Request {
|
||||
req := httptest.NewRequest("GET", "/", nil)
|
||||
req.Header.Set("X-User-ID", "test_user")
|
||||
return req
|
||||
},
|
||||
expected: "user:test_user",
|
||||
},
|
||||
{
|
||||
name: "test combined strategy",
|
||||
strategy: "combined",
|
||||
setupReq: func() *http.Request {
|
||||
req := httptest.NewRequest("GET", "/", nil)
|
||||
req.Header.Set("X-User-ID", "test_user")
|
||||
req.Header.Set("X-Real-IP", "192.168.1.1")
|
||||
return req
|
||||
},
|
||||
expected: "user:test_user:ip:192.168.1.1",
|
||||
},
|
||||
{
|
||||
name: "test default (ip) strategy",
|
||||
strategy: "invalid_strategy", // Invalid, should default to ip
|
||||
setupReq: func() *http.Request {
|
||||
req := httptest.NewRequest("GET", "/", nil)
|
||||
req.RemoteAddr = "10.0.0.1:12345"
|
||||
return req
|
||||
},
|
||||
expected: "ip:10.0.0.1",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
req := tt.setupReq()
|
||||
result := GetUserIdentifier(req, tt.strategy)
|
||||
|
||||
if result != tt.expected {
|
||||
t.Errorf("Expected '%s', got '%s'", tt.expected, result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestHTTPMiddleware tests the HTTP middleware functionality
|
||||
func TestHTTPMiddleware(t *testing.T) {
|
||||
cfg := config.DefaultConfig()
|
||||
cfg.Tools.RateLimiting.Enabled = true
|
||||
cfg.Tools.RateLimiting.Rate = 1 // 1 request per second
|
||||
cfg.Tools.RateLimiting.Burst = 1 // burst of 1
|
||||
cfg.Tools.RateLimiting.Strategy = "ip"
|
||||
|
||||
// Create a rate-limited handler
|
||||
limitedHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("OK"))
|
||||
})
|
||||
|
||||
middleware := CreateMiddlewareFromConfig(cfg)
|
||||
rateLimitedHandler := middleware(limitedHandler)
|
||||
|
||||
// Test first request - should succeed
|
||||
req1 := httptest.NewRequest("GET", "/", nil)
|
||||
req1.RemoteAddr = "192.168.1.1:12345"
|
||||
recorder1 := httptest.NewRecorder()
|
||||
|
||||
rateLimitedHandler.ServeHTTP(recorder1, req1)
|
||||
|
||||
if recorder1.Code != http.StatusOK {
|
||||
t.Errorf("First request should succeed, got status %d", recorder1.Code)
|
||||
}
|
||||
|
||||
// Test second request - should fail with 429
|
||||
req2 := httptest.NewRequest("GET", "/", nil)
|
||||
req2.RemoteAddr = "192.168.1.1:12345" // Same IP as first request
|
||||
recorder2 := httptest.NewRecorder()
|
||||
|
||||
rateLimitedHandler.ServeHTTP(recorder2, req2)
|
||||
|
||||
if recorder2.Code != http.StatusTooManyRequests {
|
||||
t.Errorf("Second request should fail with 429, got status %d", recorder2.Code)
|
||||
}
|
||||
|
||||
// Test third request from a different IP - should succeed
|
||||
req3 := httptest.NewRequest("GET", "/", nil)
|
||||
req3.RemoteAddr = "192.168.1.2:12345" // Different IP
|
||||
recorder3 := httptest.NewRecorder()
|
||||
|
||||
rateLimitedHandler.ServeHTTP(recorder3, req3)
|
||||
|
||||
if recorder3.Code != http.StatusOK {
|
||||
t.Errorf("Requests from different IPs should not affect each other, got status %d", recorder3.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// TestWithDisabledRateLimit tests behavior when rate limit is disabled
|
||||
func TestWithDisabledRateLimit(t *testing.T) {
|
||||
cfg := config.DefaultConfig()
|
||||
cfg.Tools.RateLimiting.Enabled = false // Disable rate limiting
|
||||
|
||||
// Create a handler that should work without being limited
|
||||
callCount := 0
|
||||
testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
callCount++
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("OK"))
|
||||
})
|
||||
|
||||
middleware := CreateMiddlewareFromConfig(cfg)
|
||||
wrappedHandler := middleware(testHandler)
|
||||
|
||||
// Even 10 rapid requests should succeed when rate limiting is disabled
|
||||
for i := 0; i < 10; i++ {
|
||||
req := httptest.NewRequest("GET", "/", nil)
|
||||
recorder := httptest.NewRecorder()
|
||||
wrappedHandler.ServeHTTP(recorder, req)
|
||||
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Errorf("Request %d should succeed when rate limiting is disabled, got status %d", i+1, recorder.Code)
|
||||
}
|
||||
}
|
||||
|
||||
if callCount != 10 {
|
||||
t.Errorf("Expected 10 successful requests, got %d", callCount)
|
||||
}
|
||||
}
|
||||
|
||||
// TestWaitMethod tests the Wait method with context
|
||||
func TestWaitMethod(t *testing.T) {
|
||||
limiter := NewRateLimiter(rate.Limit(1), 1, time.Hour)
|
||||
ctx := context.Background()
|
||||
|
||||
// First request should pass immediately
|
||||
err := limiter.Wait(ctx, "test-ip")
|
||||
if err != nil {
|
||||
t.Errorf("First request should not return an error: %v", err)
|
||||
}
|
||||
|
||||
// Second request should be delayed and context should not timeout
|
||||
startTime := time.Now()
|
||||
err = limiter.Wait(ctx, "test-ip")
|
||||
duration := time.Since(startTime)
|
||||
|
||||
// The wait should take approximately 1 second (for 1 request per second rate)
|
||||
// but shouldn't fail
|
||||
if err != nil {
|
||||
t.Errorf("Wait should not return an error: %v", err)
|
||||
}
|
||||
if duration < 500*time.Millisecond { // 500ms as threshold
|
||||
t.Errorf("Second request should be delayed, but took only %v", duration)
|
||||
}
|
||||
}
|
||||
60
pkg/ratelimit/server.go
Normal file
60
pkg/ratelimit/server.go
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
package ratelimit
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
// HTTPServer integrates rate limiting with an HTTP server
|
||||
type HTTPServer struct {
|
||||
config *config.Config
|
||||
server *http.Server
|
||||
limiter *RateLimiter
|
||||
}
|
||||
|
||||
// NewHTTPServer creates a new HTTP server with integrated rate limiting
|
||||
func NewHTTPServer(cfg *config.Config) *HTTPServer {
|
||||
limiter := CreateDefaultRateLimiterFromConfig(cfg)
|
||||
|
||||
return &HTTPServer{
|
||||
config: cfg,
|
||||
limiter: limiter,
|
||||
}
|
||||
}
|
||||
|
||||
// ApplyRateLimit wraps an HTTP handler with rate limiting based on configuration
|
||||
func (hs *HTTPServer) ApplyRateLimit(handler http.Handler) http.Handler {
|
||||
if !hs.config.Tools.RateLimiting.Enabled {
|
||||
return handler
|
||||
}
|
||||
|
||||
middleware := NewRateLimitMiddleware(hs.limiter)
|
||||
|
||||
strategy := hs.config.Tools.RateLimiting.Strategy
|
||||
if strategy == "" {
|
||||
strategy = "ip" // default to IP-based limiting
|
||||
}
|
||||
|
||||
return middleware.CreateMiddleware(strategy)(handler)
|
||||
}
|
||||
|
||||
// ApplyRateLimitFunc wraps an HTTP handler function with rate limiting
|
||||
func (hs *HTTPServer) ApplyRateLimitFunc(handlerFunc http.HandlerFunc) http.Handler {
|
||||
return hs.ApplyRateLimit(handlerFunc)
|
||||
}
|
||||
|
||||
// WrapRoute provides an easy way to wrap individual routes with rate limiting
|
||||
func (hs *HTTPServer) WrapRoute(pattern string, handler http.Handler) (string, http.Handler) {
|
||||
return pattern, hs.ApplyRateLimit(handler)
|
||||
}
|
||||
|
||||
// WrapRouteFunc provides an easy way to wrap individual routes with rate limiting
|
||||
func (hs *HTTPServer) WrapRouteFunc(pattern string, handlerFunc http.HandlerFunc) (string, http.Handler) {
|
||||
return pattern, hs.ApplyRateLimit(handlerFunc)
|
||||
}
|
||||
|
||||
// GetLimiter returns the rate limiter instance
|
||||
func (hs *HTTPServer) GetLimiter() *RateLimiter {
|
||||
return hs.limiter
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue