package middleware import ( "sync" "time" ) type DeviceRateLimiter struct { requests map[string]*DeviceRequestInfo mu sync.RWMutex cleanup *time.Ticker } type DeviceRequestInfo struct { Requests []time.Time LastSeen time.Time DeviceID string } type DeviceRateLimitConfig struct { SyncRequestsPerMinute int ProgressUpdatesPerMinute int MetadataRequestsPerMinute int } const ( DefaultSyncRequestsPerMinute = 60 DefaultProgressUpdatesPerMinute = 120 DefaultMetadataRequestsPerMinute = 30 CleanupInterval = 5 * time.Minute RequestWindow = 1 * time.Minute ) func NewDeviceRateLimiter() *DeviceRateLimiter { limiter := &DeviceRateLimiter{ requests: make(map[string]*DeviceRequestInfo), cleanup: time.NewTicker(CleanupInterval), } go limiter.cleanupOldEntries() return limiter } func (drl *DeviceRateLimiter) CheckRateLimit(deviceID string, requestType string, config DeviceRateLimitConfig) bool { drl.mu.Lock() defer drl.mu.Unlock() now := time.Now() key := deviceID + ":" + requestType info, exists := drl.requests[key] if !exists { info = &DeviceRequestInfo{ Requests: []time.Time{}, DeviceID: deviceID, LastSeen: now, } drl.requests[key] = info } var limit int switch requestType { case "sync": limit = config.SyncRequestsPerMinute case "progress": limit = config.ProgressUpdatesPerMinute case "metadata": limit = config.MetadataRequestsPerMinute default: limit = DefaultSyncRequestsPerMinute } if limit == 0 { limit = DefaultSyncRequestsPerMinute } info.Requests = append(info.Requests, now) info.LastSeen = now cutoff := now.Add(-RequestWindow) validRequests := 0 for _, reqTime := range info.Requests { if reqTime.After(cutoff) { validRequests++ } } info.Requests = info.Requests[len(info.Requests)-validRequests:] return validRequests <= limit } func (drl *DeviceRateLimiter) GetRemainingRequests(deviceID string, requestType string, config DeviceRateLimitConfig) int { drl.mu.RLock() defer drl.mu.RUnlock() key := deviceID + ":" + requestType info, exists := drl.requests[key] if !exists { var limit int switch requestType { case "sync": limit = config.SyncRequestsPerMinute case "progress": limit = config.ProgressUpdatesPerMinute case "metadata": limit = config.MetadataRequestsPerMinute default: limit = DefaultSyncRequestsPerMinute } if limit == 0 { limit = DefaultSyncRequestsPerMinute } return limit } now := time.Now() cutoff := now.Add(-RequestWindow) validRequests := 0 for _, reqTime := range info.Requests { if reqTime.After(cutoff) { validRequests++ } } var limit int switch requestType { case "sync": limit = config.SyncRequestsPerMinute case "progress": limit = config.ProgressUpdatesPerMinute case "metadata": limit = config.MetadataRequestsPerMinute default: limit = DefaultSyncRequestsPerMinute } if limit == 0 { limit = DefaultSyncRequestsPerMinute } remaining := limit - validRequests if remaining < 0 { remaining = 0 } return remaining } func (drl *DeviceRateLimiter) Reset(deviceID string) { drl.mu.Lock() defer drl.mu.Unlock() for key := range drl.requests { if len(key) > len(deviceID) && key[:len(deviceID)] == deviceID { delete(drl.requests, key) } } } func (drl *DeviceRateLimiter) cleanupOldEntries() { for range drl.cleanup.C { drl.mu.Lock() cutoff := time.Now().Add(-10 * time.Minute) for key, info := range drl.requests { if info.LastSeen.Before(cutoff) { delete(drl.requests, key) } } drl.mu.Unlock() } } func (drl *DeviceRateLimiter) Stop() { drl.cleanup.Stop() }