faucet: use lru cache
This commit is contained in:
parent
a9893492ba
commit
cd3539ab18
@ -49,6 +49,7 @@ import (
|
||||
"github.com/ethereum/go-ethereum/log"
|
||||
"github.com/ethereum/go-ethereum/params"
|
||||
"github.com/gorilla/websocket"
|
||||
lru "github.com/hashicorp/golang-lru"
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
@ -238,6 +239,12 @@ func newFaucet(genesis *core.Genesis, url string, ks *keystore.KeyStore, index [
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Allow 1 request per minute with burst of 5, and cache up to 10000 IPs
|
||||
limiter, err := NewIPRateLimiter(rate.Limit(1.0), 5, 1000)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &faucet{
|
||||
config: genesis.Config,
|
||||
client: client,
|
||||
@ -248,7 +255,7 @@ func newFaucet(genesis *core.Genesis, url string, ks *keystore.KeyStore, index [
|
||||
update: make(chan struct{}, 1),
|
||||
bep2eInfos: bep2eInfos,
|
||||
bep2eAbi: bep2eAbi,
|
||||
limiter: NewIPRateLimiter(rate.Limit(1), 5), // Allow 1 request per second with a burst of 5
|
||||
limiter: limiter,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@ -926,21 +933,26 @@ func getGenesis(genesisFlag string, goerliFlag bool, sepoliaFlag bool) (*core.Ge
|
||||
}
|
||||
|
||||
type IPRateLimiter struct {
|
||||
ips map[string]*rate.Limiter
|
||||
ips *lru.Cache
|
||||
mu *sync.RWMutex
|
||||
r rate.Limit
|
||||
b int
|
||||
}
|
||||
|
||||
func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
|
||||
func NewIPRateLimiter(r rate.Limit, b int, size int) (*IPRateLimiter, error) {
|
||||
cache, err := lru.New(size)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
i := &IPRateLimiter{
|
||||
ips: make(map[string]*rate.Limiter),
|
||||
ips: cache,
|
||||
mu: &sync.RWMutex{},
|
||||
r: r,
|
||||
b: b,
|
||||
}
|
||||
|
||||
return i
|
||||
return i, nil
|
||||
}
|
||||
|
||||
func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {
|
||||
@ -949,21 +961,18 @@ func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {
|
||||
|
||||
limiter := rate.NewLimiter(i.r, i.b)
|
||||
|
||||
i.ips[ip] = limiter
|
||||
i.ips.Add(ip, limiter)
|
||||
|
||||
return limiter
|
||||
}
|
||||
|
||||
func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
|
||||
i.mu.Lock()
|
||||
limiter, exists := i.ips[ip]
|
||||
defer i.mu.Unlock()
|
||||
|
||||
if limiter, exists := i.ips.Get(ip); exists {
|
||||
return limiter.(*rate.Limiter)
|
||||
}
|
||||
|
||||
if !exists {
|
||||
i.mu.Unlock()
|
||||
return i.AddIP(ip)
|
||||
}
|
||||
|
||||
i.mu.Unlock()
|
||||
|
||||
return limiter
|
||||
}
|
||||
|
@ -24,7 +24,9 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/ethereum/go-ethereum/accounts/keystore"
|
||||
"github.com/ethereum/go-ethereum/common"
|
||||
"golang.org/x/time/rate"
|
||||
|
||||
"github.com/ethereum/go-ethereum/core"
|
||||
"github.com/ethereum/go-ethereum/crypto"
|
||||
@ -53,7 +55,7 @@ func TestFacebook(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestFaucetRateLimiting(t *testing.T) {
|
||||
// Create a minimal mockfaucet instance for testing
|
||||
// Create a minimal faucet instance for testing
|
||||
privateKey, _ := crypto.GenerateKey()
|
||||
faucetAddr := crypto.PubkeyToAddress(privateKey.PublicKey)
|
||||
|
||||
@ -63,12 +65,23 @@ func TestFaucetRateLimiting(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
// Create a mockfaucet with rate limiting (1 request per second, burst of 2)
|
||||
f, err := newFaucet(config, "http://localhost:8545", nil, []byte{}, nil)
|
||||
// Create a faucet with rate limiting (1 request per second, burst of 2)
|
||||
ks := keystore.NewKeyStore(t.TempDir(), keystore.LightScryptN, keystore.LightScryptP)
|
||||
_, err := ks.NewAccount("password")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create mockfaucet: %v", err)
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(ks.Accounts()) == 0 {
|
||||
t.Fatalf("No accounts %v", ks)
|
||||
}
|
||||
f, err := newFaucet(config, "http://localhost:8545", ks, []byte{}, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create faucet: %v", err)
|
||||
}
|
||||
f.limiter, err = NewIPRateLimiter(rate.Limit(1.0), 1, 2)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create NewIPRateLimiter: %v", err)
|
||||
}
|
||||
f.limiter = NewIPRateLimiter(1, 2)
|
||||
|
||||
// Create a test server
|
||||
server := httptest.NewServer(http.HandlerFunc(f.apiHandler))
|
||||
|
Loading…
Reference in New Issue
Block a user