mirror of
https://github.com/superseriousbusiness/gotosocial.git
synced 2024-05-20 09:18:13 +00:00
bee8458a2d
* feat: add rate limit middleware * chore: update vendor dir * chore: update readme with new dependency * chore: add rate limit infos to swagger.md file * refactor: add ipv6 mask limiter option Add IPv6 CIDR /64 mask * refactor: increase rate limit to 1000 Address https://github.com/superseriousbusiness/gotosocial/pull/741#discussion_r945584800 Co-authored-by: tobi <31960611+tsmethurst@users.noreply.github.com>
83 lines
2.5 KiB
Go
83 lines
2.5 KiB
Go
package memory
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/ulule/limiter/v3"
|
|
"github.com/ulule/limiter/v3/drivers/store/common"
|
|
"github.com/ulule/limiter/v3/internal/bytebuffer"
|
|
)
|
|
|
|
// Store is the in-memory store.
|
|
type Store struct {
|
|
// Prefix used for the key.
|
|
Prefix string
|
|
// cache used to store values in-memory.
|
|
cache *CacheWrapper
|
|
}
|
|
|
|
// NewStore creates a new instance of memory store with defaults.
|
|
func NewStore() limiter.Store {
|
|
return NewStoreWithOptions(limiter.StoreOptions{
|
|
Prefix: limiter.DefaultPrefix,
|
|
CleanUpInterval: limiter.DefaultCleanUpInterval,
|
|
})
|
|
}
|
|
|
|
// NewStoreWithOptions creates a new instance of memory store with options.
|
|
func NewStoreWithOptions(options limiter.StoreOptions) limiter.Store {
|
|
return &Store{
|
|
Prefix: options.Prefix,
|
|
cache: NewCache(options.CleanUpInterval),
|
|
}
|
|
}
|
|
|
|
// Get returns the limit for given identifier.
|
|
func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
|
|
buffer := bytebuffer.New()
|
|
defer buffer.Close()
|
|
buffer.Concat(store.Prefix, ":", key)
|
|
|
|
count, expiration := store.cache.Increment(buffer.String(), 1, rate.Period)
|
|
|
|
lctx := common.GetContextFromState(time.Now(), rate, expiration, count)
|
|
return lctx, nil
|
|
}
|
|
|
|
// Increment increments the limit by given count & returns the new limit value for given identifier.
|
|
func (store *Store) Increment(ctx context.Context, key string, count int64, rate limiter.Rate) (limiter.Context, error) {
|
|
buffer := bytebuffer.New()
|
|
defer buffer.Close()
|
|
buffer.Concat(store.Prefix, ":", key)
|
|
|
|
newCount, expiration := store.cache.Increment(buffer.String(), count, rate.Period)
|
|
|
|
lctx := common.GetContextFromState(time.Now(), rate, expiration, newCount)
|
|
return lctx, nil
|
|
}
|
|
|
|
// Peek returns the limit for given identifier, without modification on current values.
|
|
func (store *Store) Peek(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
|
|
buffer := bytebuffer.New()
|
|
defer buffer.Close()
|
|
buffer.Concat(store.Prefix, ":", key)
|
|
|
|
count, expiration := store.cache.Get(buffer.String(), rate.Period)
|
|
|
|
lctx := common.GetContextFromState(time.Now(), rate, expiration, count)
|
|
return lctx, nil
|
|
}
|
|
|
|
// Reset returns the limit for given identifier.
|
|
func (store *Store) Reset(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
|
|
buffer := bytebuffer.New()
|
|
defer buffer.Close()
|
|
buffer.Concat(store.Prefix, ":", key)
|
|
|
|
count, expiration := store.cache.Reset(buffer.String(), rate.Period)
|
|
|
|
lctx := common.GetContextFromState(time.Now(), rate, expiration, count)
|
|
return lctx, nil
|
|
}
|