From 7246365a2ef8f2b152248b8b177ee9fa451c46b3 Mon Sep 17 00:00:00 2001 From: Al Cutter Date: Wed, 23 Oct 2024 10:11:52 +0100 Subject: [PATCH 1/2] Add rate-limiting. --- cmd/feedbastion/main.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/cmd/feedbastion/main.go b/cmd/feedbastion/main.go index f042b38..4faec85 100644 --- a/cmd/feedbastion/main.go +++ b/cmd/feedbastion/main.go @@ -32,6 +32,7 @@ import ( "github.com/transparency-dev/witness/internal/config" "github.com/transparency-dev/witness/omniwitness" "golang.org/x/sync/errgroup" + "golang.org/x/time/rate" "gopkg.in/yaml.v3" "k8s.io/klog/v2" ) @@ -40,6 +41,7 @@ var ( bastionURL = flag.String("bastion_url", "https://localhost:8443", "URL of the bastion service") feed = flag.String("feed", ".*", "RegEx matching log origins to feed to bastion") loopInterval = flag.Duration("loop_interval", 0, "If set to > 0, runs in looping mode sleeping this duration between feed attempts") + rateLimit = flag.Float64("max_qps", 2, "Defines maximum number of requests/s to send via bastion") ) type logFeeder struct { @@ -53,6 +55,7 @@ func main() { defer klog.Flush() ctx := context.Background() + rl := rate.NewLimiter(rate.Limit(*rateLimit), 1) httpClient := &http.Client{} insecureHttpClient := &http.Client{Transport: &http.Transport{ @@ -89,6 +92,7 @@ func main() { lf := lf if r.Match([]byte(o)) { eg.Go(func() error { + rl.Wait(ctx) return lf.info.Feeder.FeedFunc()(ctx, lf.cfg, bc, httpClient, *loopInterval) }) } From 957cd1c97471311f52f79e2c38f7740a6e4cd320 Mon Sep 17 00:00:00 2001 From: Al Cutter Date: Wed, 23 Oct 2024 10:27:46 +0100 Subject: [PATCH 2/2] lint --- cmd/feedbastion/main.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cmd/feedbastion/main.go b/cmd/feedbastion/main.go index 4faec85..d0dbf3f 100644 --- a/cmd/feedbastion/main.go +++ b/cmd/feedbastion/main.go @@ -92,7 +92,9 @@ func main() { lf := lf if r.Match([]byte(o)) { eg.Go(func() error { - rl.Wait(ctx) + if err := rl.Wait(ctx); err != nil { + return err + } return lf.info.Feeder.FeedFunc()(ctx, lf.cfg, bc, httpClient, *loopInterval) }) }