feat: Enhance Starskey rate limiter with IP jail reporting and management

- Add periodic reporting of jailed IPs every 5 minutes
- Implement startup reporting of jailed IPs count
- Improve jail status tracking with more detailed logging
- Add manual reset logging for jailed IPs
- Introduce stopReports channel for graceful goroutine termination
- Enhance Allow method with more comprehensive jail status checks

commit: 086e0dfe4e6c37897b9c2c51118f968a65308b30
author: bamen · 19mo ago · 2025-03-10 14:03 UTC
parent: 6b4468e
signature: unsigned

1 file · +83 -12

M pkg/kv/ratelimiter.go +83 -12


@@ -15,10 +15,11 @@ type RateLimitInfo struct {
15	15		}
16	16
17	17		type StarskeyRateLimiterStore struct {
18		-	db *starskey.Starskey
19		-	rate float64
20		-	burst int
21		-	expiresIn time.Duration
	18	+	db *starskey.Starskey
	19	+	rate float64
	20	+	burst int
	21	+	expiresIn time.Duration
	22	+	stopReports chan struct{} // Channel to signal stopping the report goroutine
22	23		}
23	24
24	25		// NewStarskeyRateLimiterStore creates a new rate limiter store backed by Starskey

@@ -47,12 +48,18 @@ func NewStarskeyRateLimiterStore(dbPath string, rate float64, burst int, expires
47	48		"burst", burst,
48	49		"expiration", expiresIn)
49	50
50		-	return &StarskeyRateLimiterStore{
51		-	db: db,
52		-	rate: rate,
53		-	burst: burst,
54		-	expiresIn: expiresIn,
55		-	}, nil
	51	+	store := &StarskeyRateLimiterStore{
	52	+	db: db,
	53	+	rate: rate,
	54	+	burst: burst,
	55	+	expiresIn: expiresIn,
	56	+	stopReports: make(chan struct{}),
	57	+	}
	58	+
	59	+	// Start the jailed IPs reporting
	60	+	go store.startJailReporting()
	61	+
	62	+	return store, nil
56	63		}
57	64
58	65		// Implement Echo's RateLimiterStore interface

@@ -84,6 +91,12 @@ func (s *StarskeyRateLimiterStore) Allow(identifier string) (bool, error) {
84	91		return nil
85	92		}
86	93
	94	+	// If IP was jailed but now should be freed (reset time passed)
	95	+	if info.Jailed && !now.Before(info.ResetTime) {
	96	+	log.Info("IP address released from jail", "ip", identifier)
	97	+	info.Jailed = false
	98	+	}
	99	+
87	100		// Calculate time-based allowance
88	101		timePassed := now.Sub(info.ResetTime).Seconds()
89	102		tokensToAdd := timePassed * s.rate

@@ -107,6 +120,8 @@ func (s *StarskeyRateLimiterStore) Allow(identifier string) (bool, error) {
107	120		info.Jailed = true
108	121		info.ResetTime = now.Add(time.Second) // 1 second penalty
109	122		log.Debug("IP jailed (burst exceeded)", "ip", identifier, "reset_at", info.ResetTime)
	123	+	// Add info level log when an IP is jailed
	124	+	log.Info("IP address jailed due to rate limit violation", "ip", identifier, "reset_at", info.ResetTime)
110	125		} else {
111	126		log.Debug("Request allowed", "ip", identifier, "attempts", info.Attempts)
112	127		}

@@ -134,7 +149,25 @@ func (s *StarskeyRateLimiterStore) Allow(identifier string) (bool, error) {
134	149		}
135	150
136	151		func (s *StarskeyRateLimiterStore) Reset(identifier string) (bool, error) {
137		-	err := s.db.Delete([]byte(identifier))
	152	+	// Check if the IP was jailed before resetting
	153	+	var wasJailed bool
	154	+
	155	+	// Access data before deletion to check jail status
	156	+	value, err := s.db.Get([]byte(identifier))
	157	+	if err == nil && value != nil {
	158	+	var info RateLimitInfo
	159	+	if err := json.Unmarshal(value, &info); err == nil {
	160	+	wasJailed = info.Jailed
	161	+	}
	162	+	}
	163	+
	164	+	err = s.db.Delete([]byte(identifier))
	165	+
	166	+	// Log if we freed a jailed IP
	167	+	if wasJailed {
	168	+	log.Info("IP address manually reset and released from jail", "ip", identifier)
	169	+	}
	170	+
138	171		return err == nil, err
139	172		}
140	173

@@ -177,9 +210,47 @@ func (s *StarskeyRateLimiterStore) GetJailedIPs() (map[string]time.Time, error)
177	210		return jailed, nil
178	211		}
179	212
180		-	// Close closes the underlying Starskey database
	213	+	// startJailReporting starts a goroutine that logs jailed IPs count on startup and every 5 minutes
	214	+	func (s *StarskeyRateLimiterStore) startJailReporting() {
	215	+	// Initial report on startup
	216	+	s.reportJailedIPs()
	217	+
	218	+	// Set up ticker for periodic reporting (every 5 minutes)
	219	+	ticker := time.NewTicker(5 * time.Minute)
	220	+	defer ticker.Stop()
	221	+
	222	+	for {
	223	+	select {
	224	+	case <-ticker.C:
	225	+	s.reportJailedIPs()
	226	+	case <-s.stopReports:
	227	+	log.Debug("Stopping jailed IPs reporting")
	228	+	return
	229	+	}
	230	+	}
	231	+	}
	232	+
	233	+	// reportJailedIPs gets and logs the current jailed IPs count
	234	+	func (s *StarskeyRateLimiterStore) reportJailedIPs() {
	235	+	jailed, err := s.GetJailedIPs()
	236	+	if err != nil {
	237	+	log.Error("Failed to get jailed IPs", "error", err)
	238	+	return
	239	+	}
	240	+
	241	+	count := len(jailed)
	242	+	if count > 0 {
	243	+	log.Info("Currently jailed IPs", "count", count)
	244	+	} else {
	245	+	log.Debug("No IPs currently in jail")
	246	+	}
	247	+	}
	248	+
	249	+	// Close closes the underlying Starskey database and stops the reporting goroutine
181	250		func (s *StarskeyRateLimiterStore) Close() error {
182	251		log.Debug("Closing rate limiter store")
	252	+	// Signal the reporting goroutine to stop
	253	+	close(s.stopReports)
183	254		return s.db.Close()
184	255		}
185	256


@@ -15,10 +15,11 @@ type RateLimitInfo struct {
15		}	15		}
16			16
17		type StarskeyRateLimiterStore struct {	17		type StarskeyRateLimiterStore struct {
18	-	db *starskey.Starskey	18	+	db *starskey.Starskey
19	-	rate float64	19	+	rate float64
20	-	burst int	20	+	burst int
21	-	expiresIn time.Duration	21	+	expiresIn time.Duration
			22	+	stopReports chan struct{} // Channel to signal stopping the report goroutine
22		}	23		}
23			24
24		// NewStarskeyRateLimiterStore creates a new rate limiter store backed by Starskey	25		// NewStarskeyRateLimiterStore creates a new rate limiter store backed by Starskey

@@ -47,12 +48,18 @@ func NewStarskeyRateLimiterStore(dbPath string, rate float64, burst int, expires
47		"burst", burst,	48		"burst", burst,
48		"expiration", expiresIn)	49		"expiration", expiresIn)
49			50
50	-	return &StarskeyRateLimiterStore{	51	+	store := &StarskeyRateLimiterStore{
51	-	db: db,	52	+	db: db,
52	-	rate: rate,	53	+	rate: rate,
53	-	burst: burst,	54	+	burst: burst,
54	-	expiresIn: expiresIn,	55	+	expiresIn: expiresIn,
55	-	}, nil	56	+	stopReports: make(chan struct{}),
			57	+	}
			58	+
			59	+	// Start the jailed IPs reporting
			60	+	go store.startJailReporting()
			61	+
			62	+	return store, nil
56		}	63		}
57			64
58		// Implement Echo's RateLimiterStore interface	65		// Implement Echo's RateLimiterStore interface

@@ -84,6 +91,12 @@ func (s *StarskeyRateLimiterStore) Allow(identifier string) (bool, error) {
84		return nil	91		return nil
85		}	92		}
86			93
			94	+	// If IP was jailed but now should be freed (reset time passed)
			95	+	if info.Jailed && !now.Before(info.ResetTime) {
			96	+	log.Info("IP address released from jail", "ip", identifier)
			97	+	info.Jailed = false
			98	+	}
			99	+
87		// Calculate time-based allowance	100		// Calculate time-based allowance
88		timePassed := now.Sub(info.ResetTime).Seconds()	101		timePassed := now.Sub(info.ResetTime).Seconds()
89		tokensToAdd := timePassed * s.rate	102		tokensToAdd := timePassed * s.rate

@@ -107,6 +120,8 @@ func (s *StarskeyRateLimiterStore) Allow(identifier string) (bool, error) {
107		info.Jailed = true	120		info.Jailed = true
108		info.ResetTime = now.Add(time.Second) // 1 second penalty	121		info.ResetTime = now.Add(time.Second) // 1 second penalty
109		log.Debug("IP jailed (burst exceeded)", "ip", identifier, "reset_at", info.ResetTime)	122		log.Debug("IP jailed (burst exceeded)", "ip", identifier, "reset_at", info.ResetTime)
			123	+	// Add info level log when an IP is jailed
			124	+	log.Info("IP address jailed due to rate limit violation", "ip", identifier, "reset_at", info.ResetTime)
110		} else {	125		} else {
111		log.Debug("Request allowed", "ip", identifier, "attempts", info.Attempts)	126		log.Debug("Request allowed", "ip", identifier, "attempts", info.Attempts)
112		}	127		}

@@ -134,7 +149,25 @@ func (s *StarskeyRateLimiterStore) Allow(identifier string) (bool, error) {
134		}	149		}
135			150
136		func (s *StarskeyRateLimiterStore) Reset(identifier string) (bool, error) {	151		func (s *StarskeyRateLimiterStore) Reset(identifier string) (bool, error) {
137	-	err := s.db.Delete([]byte(identifier))	152	+	// Check if the IP was jailed before resetting
			153	+	var wasJailed bool
			154	+
			155	+	// Access data before deletion to check jail status
			156	+	value, err := s.db.Get([]byte(identifier))
			157	+	if err == nil && value != nil {
			158	+	var info RateLimitInfo
			159	+	if err := json.Unmarshal(value, &info); err == nil {
			160	+	wasJailed = info.Jailed
			161	+	}
			162	+	}
			163	+
			164	+	err = s.db.Delete([]byte(identifier))
			165	+
			166	+	// Log if we freed a jailed IP
			167	+	if wasJailed {
			168	+	log.Info("IP address manually reset and released from jail", "ip", identifier)
			169	+	}
			170	+
138		return err == nil, err	171		return err == nil, err
139		}	172		}
140			173

@@ -177,9 +210,47 @@ func (s *StarskeyRateLimiterStore) GetJailedIPs() (map[string]time.Time, error)
177		return jailed, nil	210		return jailed, nil
178		}	211		}
179			212
180	-	// Close closes the underlying Starskey database	213	+	// startJailReporting starts a goroutine that logs jailed IPs count on startup and every 5 minutes
			214	+	func (s *StarskeyRateLimiterStore) startJailReporting() {
			215	+	// Initial report on startup
			216	+	s.reportJailedIPs()
			217	+
			218	+	// Set up ticker for periodic reporting (every 5 minutes)
			219	+	ticker := time.NewTicker(5 * time.Minute)
			220	+	defer ticker.Stop()
			221	+
			222	+	for {
			223	+	select {
			224	+	case <-ticker.C:
			225	+	s.reportJailedIPs()
			226	+	case <-s.stopReports:
			227	+	log.Debug("Stopping jailed IPs reporting")
			228	+	return
			229	+	}
			230	+	}
			231	+	}
			232	+
			233	+	// reportJailedIPs gets and logs the current jailed IPs count
			234	+	func (s *StarskeyRateLimiterStore) reportJailedIPs() {
			235	+	jailed, err := s.GetJailedIPs()
			236	+	if err != nil {
			237	+	log.Error("Failed to get jailed IPs", "error", err)
			238	+	return
			239	+	}
			240	+
			241	+	count := len(jailed)
			242	+	if count > 0 {
			243	+	log.Info("Currently jailed IPs", "count", count)
			244	+	} else {
			245	+	log.Debug("No IPs currently in jail")
			246	+	}
			247	+	}
			248	+
			249	+	// Close closes the underlying Starskey database and stops the reporting goroutine
181		func (s *StarskeyRateLimiterStore) Close() error {	250		func (s *StarskeyRateLimiterStore) Close() error {
182		log.Debug("Closing rate limiter store")	251		log.Debug("Closing rate limiter store")
			252	+	// Signal the reporting goroutine to stop
			253	+	close(s.stopReports)
183		return s.db.Close()	254		return s.db.Close()
184		}	255		}
185			256