From 126ef04852d364545e6c4f60e70fd99be210e7d5 Mon Sep 17 00:00:00 2001
From: Edson Zegarra <xcr33d@gmail.com>
Date: Thu, 29 Dec 2022 22:26:08 +0100
Subject: [PATCH 1/3] first approach for a rate limiter, onyl for generic
 commands

---
 cli/cli.go      | 37 +++++++++++++++++++++++++++----------
 cli/throttle.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 10 deletions(-)
 create mode 100644 cli/throttle.go

diff --git a/cli/cli.go b/cli/cli.go
index bb67a12..bf1fa62 100644
--- a/cli/cli.go
+++ b/cli/cli.go
@@ -56,18 +56,25 @@ var au aurora.Aurora
 var currentConfig *APIConfig
 
 func generic(method string, addr string, args []string) {
-	var body io.Reader
-
-	d, err := GetBody("application/json", args)
-	if err != nil {
-		panic(err)
-	}
-	if len(d) > 0 {
-		body = strings.NewReader(d)
+	rateLimiter := Throttle{
+		rate:     viper.GetInt("rsh-load-rate"),
+		requests: viper.GetInt("rsh-requests"),
 	}
 
-	req, _ := http.NewRequest(method, fixAddress(addr), body)
-	MakeRequestAndFormat(req)
+	rateLimiter.RateLimit(func() {
+		var body io.Reader
+
+		d, err := GetBody("application/json", args)
+		if err != nil {
+			panic(err)
+		}
+		if len(d) > 0 {
+			body = strings.NewReader(d)
+		}
+
+		req, _ := http.NewRequest(method, fixAddress(addr), body)
+		MakeRequestAndFormat(req)
+	})
 }
 
 // templateVarRegex used to find/replace variables `/{foo}/bar/{baz}` in a
@@ -544,6 +551,8 @@ Not after (expires): %s (%s)
 	AddGlobalFlag("rsh-raw", "r", "Output result of query as raw rather than an escaped JSON string or list", false, false)
 	AddGlobalFlag("rsh-server", "s", "Override scheme://server:port for an API", "", false)
 	AddGlobalFlag("rsh-header", "H", "Add custom header", []string{}, true)
+	AddGlobalFlag("rsh-requests", "R", "Do a number of requests", 1, false)
+	AddGlobalFlag("rsh-load-rate", "l", "Set a load rate in requests per second. Overrides the HTTP cache to false", 1, false)
 	AddGlobalFlag("rsh-query", "q", "Add custom query param", []string{}, true)
 	AddGlobalFlag("rsh-no-paginate", "", "Disable auto-pagination", false, false)
 	AddGlobalFlag("rsh-profile", "p", "API auth profile", "default", false)
@@ -735,6 +744,14 @@ func Run() (returnErr error) {
 	if headers, _ := GlobalFlags.GetStringArray("rsh-header"); len(headers) > 0 {
 		viper.Set("rsh-header", headers)
 	}
+	if requests, _ := GlobalFlags.GetInt("rsh-requests"); requests > 0 {
+		viper.Set("rsh-requests", requests)
+	}
+	if rps, _ := GlobalFlags.GetInt("rsh-load-rate"); rps > 0 {
+		viper.Set("rsh-load-rate", rps)
+		// doesn't make sense to flood with requests and obtain a cached response
+		viper.Set("rsh-no-cache", false)
+	}
 	profile, _ := GlobalFlags.GetString("rsh-profile")
 	viper.Set("rsh-profile", profile)
 
diff --git a/cli/throttle.go b/cli/throttle.go
new file mode 100644
index 0000000..df3cf36
--- /dev/null
+++ b/cli/throttle.go
@@ -0,0 +1,33 @@
+package cli
+
+import (
+	"sync"
+	"time"
+)
+
+type Throttle struct {
+	rate     int // requests per second
+	requests int // total requests
+}
+
+// simple rate limitter
+func (l Throttle) RateLimit(f func()) {
+	executed := 0
+
+	ticker := time.Tick(time.Second)
+	var wg = &sync.WaitGroup{}
+	wg.Add(l.requests)
+	for ; true; <-ticker {
+		if executed >= l.requests {
+			break
+		}
+		for i := 0; i < l.rate; i++ {
+			go func() {
+				f()
+				executed++
+				wg.Done()
+			}()
+		}
+	}
+	wg.Wait()
+}

From eb9d9284e32db131adf670967087a94eb9b23c62 Mon Sep 17 00:00:00 2001
From: Edson Zegarra <xcr33d@gmail.com>
Date: Fri, 30 Dec 2022 12:56:54 +0100
Subject: [PATCH 2/3] add rate limiter for api commands

---
 cli/operation.go | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/cli/operation.go b/cli/operation.go
index 8c316e8..ffe5aa2 100644
--- a/cli/operation.go
+++ b/cli/operation.go
@@ -139,7 +139,13 @@ func (o Operation) command() *cobra.Command {
 
 			req, _ := http.NewRequest(o.Method, uri, body)
 			req.Header = headers
-			MakeRequestAndFormat(req)
+			rateLimiter := Throttle{
+				rate:     viper.GetInt("rsh-load-rate"),
+				requests: viper.GetInt("rsh-requests"),
+			}
+			rateLimiter.RateLimit(func() {
+				MakeRequestAndFormat(req)
+			})
 		},
 	}
 

From 0a99a29e3461ee76a492595a6d670bf5df785b92 Mon Sep 17 00:00:00 2001
From: Edson Zegarra <xcr33d@gmail.com>
Date: Thu, 16 Mar 2023 18:36:09 +0100
Subject: [PATCH 3/3] fix cache flag

---
 cli/cli.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cli/cli.go b/cli/cli.go
index bf1fa62..fa47501 100644
--- a/cli/cli.go
+++ b/cli/cli.go
@@ -750,7 +750,7 @@ func Run() (returnErr error) {
 	if rps, _ := GlobalFlags.GetInt("rsh-load-rate"); rps > 0 {
 		viper.Set("rsh-load-rate", rps)
 		// doesn't make sense to flood with requests and obtain a cached response
-		viper.Set("rsh-no-cache", false)
+		viper.Set("rsh-no-cache", true)
 	}
 	profile, _ := GlobalFlags.GetString("rsh-profile")
 	viper.Set("rsh-profile", profile)