Updating rate limits to also use TOML

config/config.go

@@ -8,22 +8,26 @@ import (
 	"github.com/BurntSushi/toml"
 )
 
+type DbConfig struct {
+	Dialect        string `toml:"dialect"`
+	Username       string `toml:"username"`
+	PasswordSecret string `toml:"password-secret"`
+	Url            string `toml:"url"`
+	Port           string `toml:"port"`
+	Name           string `toml:"name"`
+}
+
 type StackConfiguration struct {
 	ConfigLoaded bool
 
-	AllowFreshAdminGeneration bool
+	AllowFreshAdminGeneration bool     `toml:"gen-fresh-admin"`
-	AdminEmails               []string
+	AdminEmails               []string `toml:"admin-emails"`
-	AdminHmacEnv              string
+	AdminHmacEnv              string   `toml:"admin-hmac-env"`
-	UserHmacEnv               string
+	UserHmacEnv               string   `toml:"user-hmac-env"`
-	AuthedRateLimitConfig     string
+	AuthedRateLimitConfig     string   `toml:"auth-rate-limit-defs"`
-	UnauthedRateLimitConfig   string
+	UnauthedRateLimitConfig   string   `toml:"unauth-rate-limit-defs"`
 
-	DbDialect        string
+	Db DbConfig `toml:"db"`
-	DbUsername       string
-	DbPasswordSecret string
-	DbUrl            string
-	DbPort           string
-	DbName           string
 }
 
 var Environment = os.Getenv("STACK_ENVIRONMENT")
@@ -60,5 +64,5 @@ func LoadConfig() {
 
 	configInternal.ConfigLoaded = true
 
-	log.Printf("Loaded Config for stack '%s': %+v", Environment, configInternal)
+	log.Printf("Loaded Config for stack '%s':\n%+v\n", Environment, configInternal)
 }

config/dev/conf.toml

@@ -1,13 +1,14 @@
-AllowFreshAdminGeneration = true
+gen-fresh-admin = true
-AdminEmails = ["admin@admin.invalid"]
+admin-emails = ["admin@admin.invalid"]
-AdminHmacEnv = "ADMIN_HMAC_ENV"
+admin-hmac-env = "ADMIN_HMAC_ENV"
-UserHmacEnv = "USER_HMAC_ENV"
+user-hmac-env = "USER_HMAC_ENV"
-AuthedRateLimitConfig = "ratelimit.auth.json"
+auth-rate-limit-defs = "ratelimit.auth.toml"
-UnauthedRateLimitConfig = "ratelimit.unauth.json"
+unauth-rate-limit-defs = "ratelimit.unauth.toml"
 
-DbDialect = "sqlite"
+[db]
-DbUrl = "prepack.db"
+dialect = "sqlite"
-DbUsername = ""
+url = "prepack.db"
-DbPasswordSecret = ""
+username = ""
-DbPort = ""
+password-secret = ""
-DbName = ""
+port = ""
+name = ""

config/dev/ratelimit.auth.json

@@ -1,17 +0,0 @@
-{
-    "": {"seconds": 60, "max": 30, "_comment": "Global ratelimit."},
-
-    "GET:/v1/sec/doot": 
-        {"seconds": 5, "max": 3, "_comment": "One DPS (Doot Per Second) for monitoring?"},
-
-    "GET:/v1/sec/2fa-doot": 
-        {"seconds": 10, "max": 1, "_comment": "2FA doot probably doesn't need much usage at all, mainly exists as a proof of concept."},
-
-    "GET:/v1/adm/doot": 
-        {"seconds": 5, "max": 3, "_comment": "One DPS (Doot Per Second) for monitoring?"},
-
-    "GET:/v1/adm/2fa-doot": 
-        {"seconds": 10, "max": 1, "_comment": "2FA doot probably doesn't need much usage at all, mainly exists as a proof of concept."}
-    
-    
-}

config/dev/ratelimit.auth.toml

@@ -0,0 +1,29 @@
+[[Rules]]
+# Global Ratelimit
+match = ""
+seconds = 60
+max = 30
+
+[[Rules]]
+# One DPS (Doot Per Second) for monitoring
+match = "GET:/v1/sec/doot"
+seconds = 5
+max = 30
+
+[[Rules]]
+# 2FA doot probably doesn't need much usage at all, mainly exists as a proof of concept.
+match = "GET:/v1/sec/2fa-doot"
+seconds = 10
+max = 1
+
+[[Rules]]
+# One Admin DPS (Doot Per Second) for monitoring
+match = "GET:/v1/adm/doot"
+seconds = 5
+max = 3
+
+[[Rules]]
+# 2FA doot probably doesn't need much usage at all, mainly exists as a proof of concept.
+match = "GET:/v1/adm/2fa-doot"
+seconds = 10
+max = 1

config/dev/ratelimit.unauth.json

@@ -1,19 +0,0 @@
-{
-    "": 
-        {"seconds": 60, "max": 30, "_comment": "Global unauthenticated ratelimit."},
-
-    "GET:/v1/doot": 
-        {"seconds": 5, "max": 5, "_comment": "Unauthenticated DOOT for server monitoring."},
-
-    "POST:/v1/login": 
-        {"seconds": 60, "max": 3, "_comment": "Prevent bruteforce attacks on Login."},
-
-    "POST:/v1/admin": 
-        {"seconds": 60, "max": 1, "_comment": "Prevent bruteforce attacks on Admin Login."},
-
-    "POST:/v1/signup": 
-        {"seconds": 1800, "max": 1, "_comment": "Prevent spam account creation."},
-
-    "POST:/v1/forgot": 
-        {"seconds": 60, "max": 1, "_comment": "Slow down 'forgot password' enumeration/spam."}
-}

config/dev/ratelimit.unauth.toml

@@ -0,0 +1,35 @@
+[[Rules]]
+# Global unauthenticated ratelimit.
+match = ""
+seconds = 60
+max = 30
+
+[[Rules]]
+# Unauthenticated DOOT for server monitoring.
+match = "GET:/v1/doot"
+seconds = 5
+max = 5
+
+[[Rules]]
+# Prevent bruteforce attacks on Login.
+match = "POST:/v1/login"
+seconds = 60
+max = 3
+
+[[Rules]]
+# Prevent bruteforce attacks on Admin Login.
+match = "POST:/v1/admin"
+seconds = 60
+max = 1
+
+[[Rules]]
+# Prevent spam account creation.
+match = "GET:/v1/adm/2fa-doot"
+seconds = 1800
+max = 1
+
+[[Rules]]
+# Slow down 'forgot password' enumeration/spam.
+match = "POST:/v1/forgot"
+seconds = 60
+max = 1

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
+	Match   string `toml:"match"`
-	max     int
+	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{}
+	rules := RuleConfig{}
-	if err := dec.Decode(&ruleMap); err != nil {
+
+	b, err := io.ReadAll(file)
+	if err != nil {
 		panic(err)
 	}
 
-	for rkey, r := range ruleMap {
+	err = toml.Unmarshal(b, &rules)
-		m.rules[rkey] = rule{duration: time.Second * time.Duration(r.seconds), limit: r.max}
+	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))
+			panic("Unauthed rate limits not loaded")
-			unauthLoaded = true
 		}
 
 		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))
+			panic("Authed rate limits not loaded")
-			authLoaded = true
 		}
 
 		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
+}

database/database.go

@@ -18,13 +18,13 @@ var Db *gorm.DB
 var Dialect gorm.Dialector
 
 func InitDialect() gorm.Dialector {
-	if config.Config().DbDialect == "sqlite" {
+	if config.Config().Db.Dialect == "sqlite" {
-		return sqlite.Open(config.Config().DbUrl)
+		return sqlite.Open(config.Config().Db.Url)
-	} else if config.Config().DbDialect == "postgres" {
+	} else if config.Config().Db.Dialect == "postgres" {
 		return postgres.New(postgres.Config{
 			DSN: fmt.Sprintf("user=%v password=%v dbname=%v port=%v sslmode=disable TimeZone=UTC",
-				config.Config().DbUsername, config.Config().DbPasswordSecret, config.Config().DbName,
+				config.Config().Db.Username, config.Config().Db.PasswordSecret, config.Config().Db.Name,
-				config.Config().DbPort),
+				config.Config().Db.Port),
 		})
 	} else {
 		panic("No valid DB config set up.")

main.go

@@ -54,6 +54,9 @@ func main() {
 		}
 	}
 
+	// Pre-load rate limits
+	controllers.LoadRateLimits()
+
 	v1 := r.Group("/v1")
 
 	v1.GET("/doot", controllers.UnauthRateLimit(), core.Doot())