diff --git a/go.mod b/go.mod index 7c3cd077f..6c5965db4 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index a3efd8b4c..58e36b561 100644 --- a/go.sum +++ b/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= diff --git a/pkg/config/config.go b/pkg/config/config.go index eac993eb9..3dfe20811 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -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 diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index 6f65dd469..f529e8412 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -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, diff --git a/pkg/ratelimit/doc.go b/pkg/ratelimit/doc.go new file mode 100644 index 000000000..b000c6e93 --- /dev/null +++ b/pkg/ratelimit/doc.go @@ -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 diff --git a/pkg/ratelimit/global.go b/pkg/ratelimit/global.go new file mode 100644 index 000000000..1612986e5 --- /dev/null +++ b/pkg/ratelimit/global.go @@ -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, + } +} diff --git a/pkg/ratelimit/integration.go b/pkg/ratelimit/integration.go new file mode 100644 index 000000000..c5128447b --- /dev/null +++ b/pkg/ratelimit/integration.go @@ -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 + } +} diff --git a/pkg/ratelimit/middleware.go b/pkg/ratelimit/middleware.go new file mode 100644 index 000000000..a3758ef4f --- /dev/null +++ b/pkg/ratelimit/middleware.go @@ -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 +} diff --git a/pkg/ratelimit/rate_limiter.go b/pkg/ratelimit/rate_limiter.go new file mode 100644 index 000000000..3ed1435e4 --- /dev/null +++ b/pkg/ratelimit/rate_limiter.go @@ -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 " + 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 + ) +} diff --git a/pkg/ratelimit/rate_limiter_test.go b/pkg/ratelimit/rate_limiter_test.go new file mode 100644 index 000000000..aa7011467 --- /dev/null +++ b/pkg/ratelimit/rate_limiter_test.go @@ -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) + } +} diff --git a/pkg/ratelimit/server.go b/pkg/ratelimit/server.go new file mode 100644 index 000000000..7de381586 --- /dev/null +++ b/pkg/ratelimit/server.go @@ -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 +}