base-golang/pkg/limiter/ip.go

package limiter

import (
	"golang.org/x/time/rate"
	"sync"
	"time"
)

var _ IPRateLimiter = (*ipRateLimiter)(nil)

type IPRateLimiter interface {
	addIP(ip string) *rate.Limiter
	GetLimiter(ip string) *rate.Limiter
}

type ipRateLimiter struct {
	ips   map[string]*rate.Limiter
	mu    *sync.RWMutex
	limit rate.Limit
	burst int
}

func NewIPRateLimiter(limit time.Duration, burst int) IPRateLimiter {
	return &ipRateLimiter{
		ips:   make(map[string]*rate.Limiter),
		mu:    new(sync.RWMutex),
		limit: rate.Every(limit),
		burst: burst,
	}
}

func (i *ipRateLimiter) addIP(ip string) *rate.Limiter {
	i.mu.Lock()
	defer i.mu.Unlock()
	limiter := rate.NewLimiter(i.limit, i.burst)
	i.ips[ip] = limiter
	return limiter
}

func (i *ipRateLimiter) GetLimiter(ip string) *rate.Limiter {
	i.mu.Lock()
	limiter, exists := i.ips[ip]
	if !exists {
		i.mu.Unlock()
		return i.addIP(ip)
	}
	i.mu.Unlock()
	return limiter
}