From d1c281febcb13f856aa4559f47628ffeacb34e69 Mon Sep 17 00:00:00 2001
From: Maartje Eyskens <maartje@eyskens.me>
Date: Sat, 13 Jul 2024 16:32:04 +0200
Subject: [PATCH] Force iRail rate limit

---
 apiserver/pkg/ris/irail/liveboard.go |  2 +-
 apiserver/pkg/ris/irail/ratelimit.go | 37 ++++++++++++++++++++++++++++
 apiserver/pkg/ris/irail/vehicle.go   |  2 +-
 3 files changed, 39 insertions(+), 2 deletions(-)
 create mode 100644 apiserver/pkg/ris/irail/ratelimit.go

diff --git a/apiserver/pkg/ris/irail/liveboard.go b/apiserver/pkg/ris/irail/liveboard.go
index 2590e4d..f34840d 100644
--- a/apiserver/pkg/ris/irail/liveboard.go
+++ b/apiserver/pkg/ris/irail/liveboard.go
@@ -109,7 +109,7 @@ func GetLiveboard(station, arriveOrDeparture, lang string, from time.Time) (Live
 	}
 	req.Header.Set("User-Agent", USER_AGENT)
 
-	client := &http.Client{}
+	client := newClient()
 	resp, err := client.Do(req)
 	if err != nil {
 		return Liveboard{}, err
diff --git a/apiserver/pkg/ris/irail/ratelimit.go b/apiserver/pkg/ris/irail/ratelimit.go
new file mode 100644
index 0000000..12071cd
--- /dev/null
+++ b/apiserver/pkg/ris/irail/ratelimit.go
@@ -0,0 +1,37 @@
+package irail
+
+import (
+	"context"
+	"net/http"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+var rateLimit = rate.NewLimiter(rate.Every(time.Second), 3)
+
+type RLHTTPClient struct {
+	client      *http.Client
+	Ratelimiter *rate.Limiter
+}
+
+func (c *RLHTTPClient) Do(req *http.Request) (*http.Response, error) {
+	ctx := context.Background()
+	err := c.Ratelimiter.Wait(ctx) // This is a blocking call. Honors the rate limit
+	if err != nil {
+		return nil, err
+	}
+	resp, err := c.client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	return resp, nil
+}
+
+func newClient() *RLHTTPClient {
+	c := &RLHTTPClient{
+		client:      http.DefaultClient,
+		Ratelimiter: rateLimit,
+	}
+	return c
+}
diff --git a/apiserver/pkg/ris/irail/vehicle.go b/apiserver/pkg/ris/irail/vehicle.go
index 14a14c9..406e917 100644
--- a/apiserver/pkg/ris/irail/vehicle.go
+++ b/apiserver/pkg/ris/irail/vehicle.go
@@ -93,7 +93,7 @@ func GetVehicleCached(id, lang string, date time.Time) (Vehicle, error) {
 	}
 	req.Header.Set("User-Agent", USER_AGENT)
 
-	client := &http.Client{}
+	client := newClient()
 	resp, err := client.Do(req)
 	if err != nil {
 		return Vehicle{}, err