Justin Harms
d24d1dc5ae
- Created a new Go module named 'teleport' for secure port forwarding. - Added essential files including .gitignore, LICENSE, and README.md with project details. - Implemented configuration management with YAML support in config package. - Developed core client and server functionalities for handling port forwarding. - Introduced DNS server capabilities and integrated logging with sanitization. - Established rate limiting and metrics tracking for performance monitoring. - Included comprehensive tests for core components and functionalities. - Set up CI workflows for automated testing and release management using Gitea actions.
90 lines
1.9 KiB
Go
90 lines
1.9 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// RateLimiter implements a token bucket rate limiter
|
|
type RateLimiter struct {
|
|
requestsPerSecond int
|
|
burstSize int
|
|
windowSize time.Duration
|
|
tokens int
|
|
lastRefill time.Time
|
|
mutex sync.Mutex
|
|
}
|
|
|
|
// NewRateLimiter creates a new rate limiter
|
|
func NewRateLimiter(requestsPerSecond, burstSize int, windowSize time.Duration) *RateLimiter {
|
|
return &RateLimiter{
|
|
requestsPerSecond: requestsPerSecond,
|
|
burstSize: burstSize,
|
|
tokens: burstSize,
|
|
lastRefill: time.Now(),
|
|
windowSize: windowSize,
|
|
}
|
|
}
|
|
|
|
// Allow checks if a request is allowed under the rate limit
|
|
func (rl *RateLimiter) Allow() bool {
|
|
rl.mutex.Lock()
|
|
defer rl.mutex.Unlock()
|
|
|
|
now := time.Now()
|
|
|
|
// Calculate tokens to add based on time elapsed
|
|
elapsed := now.Sub(rl.lastRefill)
|
|
tokensToAdd := int(elapsed.Seconds() * float64(rl.requestsPerSecond))
|
|
|
|
if tokensToAdd > 0 {
|
|
rl.tokens += tokensToAdd
|
|
if rl.tokens > rl.burstSize {
|
|
rl.tokens = rl.burstSize
|
|
}
|
|
rl.lastRefill = now
|
|
}
|
|
|
|
// Check if we have tokens available
|
|
if rl.tokens > 0 {
|
|
rl.tokens--
|
|
return true
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
// AllowWithContext checks if a request is allowed with context cancellation
|
|
func (rl *RateLimiter) AllowWithContext(ctx context.Context) bool {
|
|
select {
|
|
case <-ctx.Done():
|
|
return false
|
|
default:
|
|
return rl.Allow()
|
|
}
|
|
}
|
|
|
|
// Wait blocks until a request is allowed
|
|
func (rl *RateLimiter) Wait(ctx context.Context) error {
|
|
for {
|
|
if rl.Allow() {
|
|
return nil
|
|
}
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-time.After(rl.windowSize / time.Duration(rl.requestsPerSecond)):
|
|
// Wait for next token
|
|
}
|
|
}
|
|
}
|
|
|
|
// GetStats returns current rate limiter statistics
|
|
func (rl *RateLimiter) GetStats() (tokens int, lastRefill time.Time) {
|
|
rl.mutex.Lock()
|
|
defer rl.mutex.Unlock()
|
|
return rl.tokens, rl.lastRefill
|
|
}
|