From b5cbed7035d548331c945773096a0be2bae2da38 Mon Sep 17 00:00:00 2001
From: plor <plor@nihil.wtf>
Date: Sun, 5 May 2024 16:10:48 -0500
Subject: [PATCH] Filters based on allowed user agent (regexp)

---
 gateway/main.go              |  1 +
 gateway/plugins/leaky.go     |  2 +-
 gateway/plugins/useragent.go | 78 ++++++++++++++++++++++++++++++++++++
 3 files changed, 80 insertions(+), 1 deletion(-)
 create mode 100644 gateway/plugins/useragent.go

diff --git a/gateway/main.go b/gateway/main.go
index 69291784..07da5f99 100644
--- a/gateway/main.go
+++ b/gateway/main.go
@@ -20,6 +20,7 @@ func main() {
         proxy.Register(&plugins.BalanceTracker{})
         proxy.Register(&plugins.LeakyBucketPlugin{"APP"})
         proxy.Register(&plugins.NoopFilter{proxy.LifecycleMask(proxy.AccountLookup|proxy.RateLimit|proxy.BalanceCheck)})
+        proxy.Register(&plugins.UserAgentFilter{})
         proxy.Register(proxy.NewReconciler(300)) // seconds
 
         gateway()
diff --git a/gateway/plugins/leaky.go b/gateway/plugins/leaky.go
index c6ac0d54..5be8a295 100644
--- a/gateway/plugins/leaky.go
+++ b/gateway/plugins/leaky.go
@@ -69,7 +69,7 @@ func (l *LeakyBucketPlugin) getBucketsForScope(ctx context.Context, app *db.App)
         return buckets
     }
     for _, rule := range rules {
-        if rule.RuleType != "rate-limits" {
+        if rule.RuleType != "rate-limits" || !rule.Active {
             continue
         }
         rate, err := utils.ParseRate(rule.Value)
diff --git a/gateway/plugins/useragent.go b/gateway/plugins/useragent.go
new file mode 100644
index 00000000..46205726
--- /dev/null
+++ b/gateway/plugins/useragent.go
@@ -0,0 +1,78 @@
+package plugins
+
+import (
+    "context"
+    "log"
+    "net/http"
+    "regexp"
+
+    "porters/db"
+    "porters/proxy"
+)
+
+const (
+    UA_TYPE_ID string = "allowed-user-agents"
+)
+
+type UserAgentFilter struct {
+
+}
+
+func (u *UserAgentFilter) Name() string {
+    return "User Agent Filter"
+}
+
+func (u *UserAgentFilter) Key() string {
+    return "USERAGENT"
+}
+
+func (u *UserAgentFilter) Load() {
+    log.Println("Loading", u.Name())
+}
+
+func (u *UserAgentFilter) HandleRequest(req *http.Request) error {
+    ctx := req.Context()
+    ua := req.UserAgent()
+    appId := proxy.PluckAppId(req)
+    app := &db.App{Id: appId}
+    err := app.Lookup(ctx)
+    if err != nil {
+        return proxy.NewHTTPError(http.StatusNotFound)
+    }
+
+    rules := u.getRulesForScope(ctx, app)
+    success := (len(rules) == 0)
+
+    for _, rule := range rules {
+        if (rule.MatchString(ua)) {
+            success = true
+            break
+        }
+    }
+
+    if !success {
+        return proxy.NewHTTPError(http.StatusUnauthorized)
+    }
+    return nil
+}
+
+func (u *UserAgentFilter) getRulesForScope(ctx context.Context, app *db.App) []regexp.Regexp {
+    useragents := make([]regexp.Regexp, 0)
+    rules, err := app.Rules(ctx)
+    if err != nil {
+        log.Println("couldn't get rules", err)
+    } else {
+        for _, rule := range rules {
+            if rule.RuleType != UA_TYPE_ID || !rule.Active {
+                continue
+            }
+            matcher, err := regexp.Compile(rule.Value)
+            if err != nil {
+                log.Println("unable to compile regexp", err)
+            } else {
+                useragents = append(useragents, *matcher)
+            }
+        }
+    }
+    return useragents
+}