From a2512d13e4f2fd15272d3d7b6d6b31c4dd75b204 Mon Sep 17 00:00:00 2001
From: Taras Madan <tarasmadan@google.com>
Date: Mon, 22 Apr 2024 12:22:46 +0200
Subject: dashboard/app: backpressure robots

Single threaded curl and wget are Ok to use.
Current logic returns 429 client err requiring users to rate limit requests.
Let's add backpressure to guarantee acceptable request rate.
---
 dashboard/app/handler.go | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

(limited to 'dashboard/app/handler.go')

diff --git a/dashboard/app/handler.go b/dashboard/app/handler.go
index ee937f1c6..cf9bc9736 100644
--- a/dashboard/app/handler.go
+++ b/dashboard/app/handler.go
@@ -36,6 +36,7 @@ func handleContext(fn contextHandler) http.Handler {
 		if !throttleRequest(c, w, r) {
 			return
 		}
+		defer backpressureRobots(c, r)()
 		if err := fn(c, w, r); err != nil {
 			hdr := commonHeaderRaw(c, r)
 			data := &struct {
@@ -78,6 +79,35 @@ func handleContext(fn contextHandler) http.Handler {
 	})
 }
 
+func isRobot(r *http.Request) bool {
+	userAgent := strings.ToLower(strings.Join(r.Header["User-Agent"], " "))
+	if strings.HasPrefix(userAgent, "curl") ||
+		strings.HasPrefix(userAgent, "wget") {
+		return true
+	}
+	return false
+}
+
+// We don't count the request round trip time here.
+// Actual delay will be the minDelay + requestRoundTripTime.
+func backpressureRobots(c context.Context, r *http.Request) func() {
+	if !isRobot(r) {
+		return func() {}
+	}
+	cfg := getConfig(c).Throttle
+	if cfg.Empty() {
+		return func() {}
+	}
+	minDelay := cfg.Window / time.Duration(cfg.Limit)
+	delayUntil := time.Now().Add(minDelay)
+	return func() {
+		select {
+		case <-c.Done():
+		case <-time.After(time.Until(delayUntil)):
+		}
+	}
+}
+
 func throttleRequest(c context.Context, w http.ResponseWriter, r *http.Request) bool {
 	// AppEngine removes all App Engine-specific headers, which include
 	// X-Appengine-User-IP and X-Forwarded-For.
-- 
cgit mrf-deployment