faucet: use lru cache

This commit is contained in:
emailtovamos 2024-07-23 17:44:53 +01:00
parent a9893492ba
commit cd3539ab18
2 changed files with 40 additions and 18 deletions

@ -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 !exists {
i.mu.Unlock()
return i.AddIP(ip)
if limiter, exists := i.ips.Get(ip); exists {
return limiter.(*rate.Limiter)
}
i.mu.Unlock()
return limiter
return i.AddIP(ip)
}

@ -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))