diff --git a/README.md b/README.md index d409abc..d8b0251 100644 --- a/README.md +++ b/README.md @@ -117,12 +117,14 @@ type CustomStore struct { func (s *CustomStore) Limit(key string, c *gin.Context) Info { if UserWentOverLimit { return Info{ + Limit: 100, RateLimited: true, ResetTime: reset, RemainingHits: 0, } } return Info{ + Limit: 100, RateLimited: false, ResetTime: reset, RemainingHits: remaining, diff --git a/gin_rate_limit.go b/gin_rate_limit.go index 0610370..437d5f8 100644 --- a/gin_rate_limit.go +++ b/gin_rate_limit.go @@ -7,6 +7,7 @@ import ( ) type Info struct { + Limit uint RateLimited bool ResetTime time.Time RemainingHits uint @@ -31,14 +32,16 @@ func RateLimiter(s Store, options *Options) gin.HandlerFunc { } if options.ErrorHandler == nil { options.ErrorHandler = func(c *gin.Context, info Info) { - c.Header("X-Rate-Limit-Reset", fmt.Sprintf("%.2f", time.Until(info.ResetTime).Seconds())) + c.Header("X-Rate-Limit-Limit", fmt.Sprintf("%d", info.Limit)) + c.Header("X-Rate-Limit-Reset", fmt.Sprintf("%d", info.ResetTime.Unix())) c.String(429, "Too many requests") } } if options.BeforeResponse == nil { options.BeforeResponse = func(c *gin.Context, info Info) { + c.Header("X-Rate-Limit-Limit", fmt.Sprintf("%d", info.Limit)) c.Header("X-Rate-Limit-Remaining", fmt.Sprintf("%v", info.RemainingHits)) - c.Header("X-Rate-Limit-Reset", fmt.Sprintf("%.2f", time.Until(info.ResetTime).Seconds())) + c.Header("X-Rate-Limit-Reset", fmt.Sprintf("%d", info.ResetTime.Unix())) } } if options.KeyFunc == nil { diff --git a/in_memory.go b/in_memory.go index a8f3db9..a2e0f09 100644 --- a/in_memory.go +++ b/in_memory.go @@ -43,6 +43,7 @@ func (s *inMemoryStoreType) Limit(key string, c *gin.Context) Info { } if s.skip != nil && s.skip(c) { return Info{ + Limit: s.limit, RateLimited: false, ResetTime: time.Now().Add(time.Duration((s.rate - (time.Now().Unix() - u.ts)) * time.Second.Nanoseconds())), RemainingHits: u.tokens, @@ -50,6 +51,7 @@ func (s *inMemoryStoreType) Limit(key string, c *gin.Context) Info { } if u.tokens <= 0 { return Info{ + Limit: s.limit, RateLimited: true, ResetTime: time.Now().Add(time.Duration((s.rate - (time.Now().Unix() - u.ts)) * time.Second.Nanoseconds())), RemainingHits: 0, @@ -59,6 +61,7 @@ func (s *inMemoryStoreType) Limit(key string, c *gin.Context) Info { u.ts = time.Now().Unix() s.data.Store(key, u) return Info{ + Limit: s.limit, RateLimited: false, ResetTime: time.Now().Add(time.Duration((s.rate - (time.Now().Unix() - u.ts)) * time.Second.Nanoseconds())), RemainingHits: u.tokens, diff --git a/redis.go b/redis.go index 1d3e049..063359b 100644 --- a/redis.go +++ b/redis.go @@ -38,6 +38,7 @@ func (s *redisStoreType) Limit(key string, c *gin.Context) Info { } if s.skip != nil && s.skip(c) { return Info{ + Limit: s.limit, RateLimited: false, ResetTime: time.Now().Add(time.Duration((s.rate - (time.Now().Unix() - ts)) * time.Second.Nanoseconds())), RemainingHits: s.limit - uint(hits), @@ -50,6 +51,7 @@ func (s *redisStoreType) Limit(key string, c *gin.Context) Info { panic(err) } else { return Info{ + Limit: s.limit, RateLimited: false, ResetTime: time.Now().Add(time.Duration((s.rate - (time.Now().Unix() - ts)) * time.Second.Nanoseconds())), RemainingHits: 0, @@ -57,6 +59,7 @@ func (s *redisStoreType) Limit(key string, c *gin.Context) Info { } } return Info{ + Limit: s.limit, RateLimited: true, ResetTime: time.Now().Add(time.Duration((s.rate - (time.Now().Unix() - ts)) * time.Second.Nanoseconds())), RemainingHits: 0, @@ -74,6 +77,7 @@ func (s *redisStoreType) Limit(key string, c *gin.Context) Info { panic(err) } else { return Info{ + Limit: s.limit, RateLimited: false, ResetTime: time.Now().Add(time.Duration((s.rate - (time.Now().Unix() - ts)) * time.Second.Nanoseconds())), RemainingHits: s.limit - uint(hits), @@ -81,6 +85,7 @@ func (s *redisStoreType) Limit(key string, c *gin.Context) Info { } } return Info{ + Limit: s.limit, RateLimited: false, ResetTime: time.Now().Add(time.Duration((s.rate - (time.Now().Unix() - ts)) * time.Second.Nanoseconds())), RemainingHits: s.limit - uint(hits),