Simple rate-limiting added

controllers/ratelimit.go

@@ -0,0 +1,108 @@
+package controllers
+
+import (
+	"net/http"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/yxzzy-wtf/gin-gonic-prepack/util"
+)
+
+type rule struct {
+	duration time.Duration
+	limit    int
+}
+
+type bucket struct {
+	rules  map[string]rule
+	access map[string]int
+}
+
+func (b *bucket) take(resource string) bool {
+	r, ex := b.rules[resource]
+	if !ex {
+		resource = "*"
+		r = b.rules[resource]
+	}
+	max := r.limit
+	duration := r.duration
+
+	remaining, ex := b.access[resource]
+	if !ex {
+		b.access[resource] = max
+		remaining = max
+	}
+
+	if remaining > 0 {
+		remaining = remaining - 1
+		b.access[resource] = remaining
+
+		go func(b *bucket, res string, d time.Duration) {
+			time.Sleep(d)
+			b.access[resource] = b.access[resource] + 1
+		}(b, resource, duration)
+
+		return true
+	}
+
+	return false
+}
+
+type megabucket struct {
+	buckets map[string]bucket
+	rules   map[string]rule
+}
+
+func (m *megabucket) take(signature string, resource string) bool {
+	b, ex := m.buckets[signature]
+	if !ex {
+		b = bucket{
+			rules:  m.rules,
+			access: map[string]int{},
+		}
+		m.buckets[signature] = b
+	}
+
+	return b.take(resource)
+}
+
+var unauthed = megabucket{
+	buckets: map[string]bucket{},
+	rules: map[string]rule{
+		"*": {duration: time.Second * 10, limit: 20},
+	},
+}
+
+func UnauthRateLimit() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		ip := c.ClientIP()
+
+		if !unauthed.take(ip, "") {
+			c.AbortWithStatus(http.StatusTooManyRequests)
+			return
+		}
+	}
+}
+
+var authed = megabucket{
+	buckets: map[string]bucket{},
+	rules: map[string]rule{
+		"*": {duration: time.Second * 10, limit: 5},
+	},
+}
+
+func AuthedRateLimit() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		pif, exists := c.Get("principal")
+		p := pif.(util.PrincipalInfo)
+		if !exists {
+			c.AbortWithStatus(http.StatusUnauthorized)
+			return
+		}
+
+		if !authed.take(p.Uid.String(), c.FullPath()) {
+			c.AbortWithStatus(http.StatusTooManyRequests)
+			return
+		}
+	}
+}