Add initial project structure with core functionality
- 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.
This commit is contained in:
89
pkg/ratelimit/ratelimit.go
Normal file
89
pkg/ratelimit/ratelimit.go
Normal file
@@ -0,0 +1,89 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user