bookhoard/internal/middleware/device_rate_limiter.go

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
	ClaimAttemptsPerMinute    int
}

const (
	DefaultSyncRequestsPerMinute     = 60
	DefaultProgressUpdatesPerMinute  = 120
	DefaultMetadataRequestsPerMinute = 30
	DefaultClaimAttemptsPerMinute    = 12
	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
	case "claim":
		limit = config.ClaimAttemptsPerMinute
	default:
		limit = DefaultSyncRequestsPerMinute
	}

	if limit == 0 {
		switch requestType {
		case "progress":
			limit = DefaultProgressUpdatesPerMinute
		case "metadata":
			limit = DefaultMetadataRequestsPerMinute
		case "claim":
			limit = DefaultClaimAttemptsPerMinute
		default:
			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) 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()
}