Updating rate limits to also use TOML

controllers/ratelimit.go

@@ -1,21 +1,28 @@
 package controllers
 
 import (
-	"encoding/json"
+	"fmt"
+	"io"
 	"log"
 	"net/http"
 	"os"
 	"regexp"
 	"time"
 
+	"github.com/BurntSushi/toml"
 	"github.com/gin-gonic/gin"
 	"github.com/yxzzy-wtf/gin-gonic-prepack/config"
 	"github.com/yxzzy-wtf/gin-gonic-prepack/util"
 )
 
+type RuleConfig struct {
+	Rules []ruleDescription
+}
+
 type ruleDescription struct {
-	seconds int
-	max     int
+	Match   string `toml:"match"`
+	Seconds int    `toml:"seconds"`
+	Max     int    `toml:"max"`
 }
 
 type rule struct {
@@ -80,16 +87,27 @@ type megabucket struct {
 }
 
 func (m *megabucket) loadFromConfig(filename string) {
-	file, _ := os.Open("conf.json")
+	file, err := os.Open(filename)
+	if err != nil {
+		panic(err)
+	}
 	defer file.Close()
-	dec := json.NewDecoder(file)
-	ruleMap := map[string]ruleDescription{}
-	if err := dec.Decode(&ruleMap); err != nil {
+
+	rules := RuleConfig{}
+
+	b, err := io.ReadAll(file)
+	if err != nil {
 		panic(err)
 	}
 
-	for rkey, r := range ruleMap {
-		m.rules[rkey] = rule{duration: time.Second * time.Duration(r.seconds), limit: r.max}
+	err = toml.Unmarshal(b, &rules)
+	if err != nil {
+		panic(err)
+	}
+
+	for _, r := range rules.Rules {
+		fmt.Printf("Loading ratelimit rule: %+v\n", r)
+		m.rules[r.Match] = rule{duration: time.Second * time.Duration(r.Seconds), limit: r.Max}
 	}
 }
 
@@ -119,8 +137,7 @@ var unauthLoaded = false
 func UnauthRateLimit() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		if !unauthLoaded {
-			unauthed.loadFromConfig(config.GetConfigPath(config.Config().UnauthedRateLimitConfig))
-			unauthLoaded = true
+			panic("Unauthed rate limits not loaded")
 		}
 
 		ip := c.ClientIP()
@@ -145,8 +162,7 @@ var authLoaded = false
 func AuthedRateLimit() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		if !authLoaded {
-			authed.loadFromConfig(config.GetConfigPath(config.Config().AuthedRateLimitConfig))
-			authLoaded = true
+			panic("Authed rate limits not loaded")
 		}
 
 		pif, exists := c.Get("principal")
@@ -162,3 +178,10 @@ func AuthedRateLimit() gin.HandlerFunc {
 		}
 	}
 }
+
+func LoadRateLimits() {
+	authed.loadFromConfig(config.GetConfigPath(config.Config().AuthedRateLimitConfig))
+	authLoaded = true
+	unauthed.loadFromConfig(config.GetConfigPath(config.Config().UnauthedRateLimitConfig))
+	unauthLoaded = true
+}