feat: initial open-source release of Pipely
OTA update server designed for high-concurrency, low-bandwidth deployments. GORM-backed PostgreSQL, JWT auth, device management, artefact versioning, deployment rollout with gated rate-limited downloads, and a React admin panel.
This commit is contained in:
@@ -0,0 +1,113 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// TokenBucket implements a token bucket rate limiter per key (e.g. IP or device ID).
|
||||
type TokenBucket struct {
|
||||
mu sync.Mutex
|
||||
buckets map[string]*bucket
|
||||
rate float64 // tokens per second
|
||||
capacity float64
|
||||
}
|
||||
|
||||
type bucket struct {
|
||||
tokens float64
|
||||
lastFill time.Time
|
||||
}
|
||||
|
||||
func NewTokenBucket(rate, capacity int) *TokenBucket {
|
||||
return &TokenBucket{
|
||||
buckets: make(map[string]*bucket),
|
||||
rate: float64(rate),
|
||||
capacity: float64(capacity),
|
||||
}
|
||||
}
|
||||
|
||||
func (tb *TokenBucket) Allow(key string) bool {
|
||||
tb.mu.Lock()
|
||||
defer tb.mu.Unlock()
|
||||
|
||||
b, ok := tb.buckets[key]
|
||||
if !ok {
|
||||
b = &bucket{tokens: tb.capacity, lastFill: time.Now()}
|
||||
tb.buckets[key] = b
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
elapsed := now.Sub(b.lastFill).Seconds()
|
||||
b.tokens += elapsed * tb.rate
|
||||
if b.tokens > tb.capacity {
|
||||
b.tokens = tb.capacity
|
||||
}
|
||||
b.lastFill = now
|
||||
|
||||
if b.tokens >= 1 {
|
||||
b.tokens--
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Cleanup removes stale buckets older than the given duration.
|
||||
func (tb *TokenBucket) Cleanup(maxAge time.Duration) {
|
||||
tb.mu.Lock()
|
||||
defer tb.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
for k, b := range tb.buckets {
|
||||
if now.Sub(b.lastFill) > maxAge {
|
||||
delete(tb.buckets, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// RateLimitMiddleware returns a Gin middleware that rate-limits requests per client IP.
|
||||
func RateLimitMiddleware(tb *TokenBucket) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
ip := c.ClientIP()
|
||||
if !tb.Allow(ip) {
|
||||
c.AbortWithStatusJSON(429, gin.H{"error": "Rate limit exceeded. Please slow down."})
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// GatedSemaphore limits the number of concurrent operations (e.g. downloads).
|
||||
type GatedSemaphore struct {
|
||||
ch chan struct{}
|
||||
}
|
||||
|
||||
func NewGatedSemaphore(maxConcurrent int) *GatedSemaphore {
|
||||
return &GatedSemaphore{ch: make(chan struct{}, maxConcurrent)}
|
||||
}
|
||||
|
||||
func (gs *GatedSemaphore) Acquire() {
|
||||
gs.ch <- struct{}{}
|
||||
}
|
||||
|
||||
func (gs *GatedSemaphore) Release() {
|
||||
<-gs.ch
|
||||
}
|
||||
|
||||
func (gs *GatedSemaphore) Available() int {
|
||||
return cap(gs.ch) - len(gs.ch)
|
||||
}
|
||||
|
||||
// GatedMiddleware only allows requests through if the semaphore has capacity.
|
||||
func GatedMiddleware(gs *GatedSemaphore) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
select {
|
||||
case gs.ch <- struct{}{}:
|
||||
defer func() { <-gs.ch }()
|
||||
c.Next()
|
||||
default:
|
||||
c.AbortWithStatusJSON(503, gin.H{"error": "Server at capacity. Please retry later."})
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user