faucet: rate limit initial implementation

cmd/faucet/faucet.go

@@ -49,6 +49,7 @@
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/gorilla/websocket"
+	"golang.org/x/time/rate"
 )
 
 var (
@@ -216,6 +217,8 @@
 
 	bep2eInfos map[string]bep2eInfo
 	bep2eAbi   abi.ABI
+
+	limiter *IPRateLimiter
 }
 
 // wsConn wraps a websocket connection with a write mutex as the underlying
@@ -235,6 +238,12 @@
 		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,
@@ -245,6 +254,7 @@
 		update:     make(chan struct{}, 1),
 		bep2eInfos: bep2eInfos,
 		bep2eAbi:   bep2eAbi,
+		limiter:    limiter,
 	}, nil
 }
 
@@ -271,7 +281,23 @@
 }
 
 // apiHandler handles requests for Ether grants and transaction statuses.
 func (f *faucet) apiHandler(w http.ResponseWriter, r *http.Request) {
+
+	ip := r.RemoteAddr
+	if len(r.Header.Get("X-Forwarded-For")) > 0 {
+		ips := strings.Split(r.Header.Get("X-Forwarded-For"), ",")
+		if len(ips) > 0 {
+			ip = strings.TrimSpace(ips[len(ips)-1])
+		}
+	}
+
+	limiter := f.limiter.GetLimiter(ip)
+	if !limiter.Allow() {
+		log.Warn("Too many requests from client: ", "client", ip)
+		http.Error(w, "Too many requests", http.StatusTooManyRequests)
+		return
+	}
+
 	upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
@@ -625,19 +651,22 @@
 			balance := new(big.Int).Div(f.balance, ether)
 
 			for _, conn := range f.conns {
-				if err := send(conn, map[string]interface{}{
-					"funds":    balance,
-					"funded":   f.nonce,
-					"requests": f.reqs,
-				}, time.Second); err != nil {
-					log.Warn("Failed to send stats to client", "err", err)
-					conn.conn.Close()
-					continue
-				}
-				if err := send(conn, head, time.Second); err != nil {
-					log.Warn("Failed to send header to client", "err", err)
-					conn.conn.Close()
-				}
+				go func(conn *wsConn) {
+					if err := send(conn, map[string]interface{}{
+						"funds":    balance,
+						"funded":   f.nonce,
+						"requests": f.reqs,
+					}, time.Second); err != nil {
+						log.Warn("Failed to send stats to client", "err", err)
+						conn.conn.Close()
+						return // Exit the goroutine if the first send fails
+					}
+
+					if err := send(conn, head, time.Second); err != nil {
+						log.Warn("Failed to send header to client", "err", err)
+						conn.conn.Close()
+					}
+				}(conn)
 			}
 			f.lock.RUnlock()
 		}
@@ -656,10 +685,12 @@
 			// Pending requests updated, stream to clients
 			f.lock.RLock()
 			for _, conn := range f.conns {
-				if err := send(conn, map[string]interface{}{"requests": f.reqs}, time.Second); err != nil {
-					log.Warn("Failed to send requests to client", "err", err)
-					conn.conn.Close()
-				}
+				go func(conn *wsConn) {
+					if err := send(conn, map[string]interface{}{"requests": f.reqs}, time.Second); err != nil {
+						log.Warn("Failed to send requests to client", "err", err)
+						conn.conn.Close()
+					}
+				}(conn)
 			}
 			f.lock.RUnlock()
 		}

cmd/faucet/faucet_test.go

@@ -17,9 +17,19 @@
 package main
 
 import (
+	"net/http"
 	"testing"
 
+	"net/http/httptest"
+	"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"
 )
 
 func TestFacebook(t *testing.T) {
@@ -43,3 +53,86 @@
 		}
 	}
 }
+
+func TestFaucetRateLimiting(t *testing.T) {
+	// Create a minimal faucet instance for testing
+	privateKey, _ := crypto.GenerateKey()
+	faucetAddr := crypto.PubkeyToAddress(privateKey.PublicKey)
+
+	config := &core.Genesis{
+		Alloc: core.GenesisAlloc{
+			faucetAddr: {Balance: common.Big1},
+		},
+	}
+
+	// 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.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)
+	}
+
+	// Create a test server
+	server := httptest.NewServer(http.HandlerFunc(f.apiHandler))
+	defer server.Close()
+
+	// Helper function to make a request
+	makeRequest := func() int {
+		resp, err := http.Get(server.URL)
+		if err != nil {
+			t.Fatalf("Failed to send request: %v", err)
+		}
+		defer resp.Body.Close()
+		return resp.StatusCode
+	}
+
+	// Test rapid requests
+	var wg sync.WaitGroup
+	results := make([]int, 5)
+
+	for i := 0; i < 5; i++ {
+		wg.Add(1)
+		go func(index int) {
+			defer wg.Done()
+			results[index] = makeRequest()
+		}(i)
+	}
+
+	wg.Wait()
+
+	// Check results
+	successCount := 0
+	rateLimitCount := 0
+	for _, status := range results {
+		if status == http.StatusOK {
+			successCount++
+		} else if status == http.StatusTooManyRequests {
+			rateLimitCount++
+		}
+	}
+
+	// We expect 2 successful requests (due to burst) and 3 rate-limited requests
+	if successCount != 2 || rateLimitCount != 3 {
+		t.Errorf("Expected 2 successful and 3 rate-limited requests, got %d successful and %d rate-limited", successCount, rateLimitCount)
+	}
+
+	// Wait for rate limit to reset
+	time.Sleep(2 * time.Second)
+
+	// Make another request, it should succeed
+	status := makeRequest()
+	if status != http.StatusOK {
+		t.Errorf("Expected success after rate limit reset, got status %d", status)
+	}
+}

cmd/faucet/rate_limiter.go

@@ -0,0 +1,53 @@
+package main
+
+import (
+	"sync"
+
+	lru "github.com/hashicorp/golang-lru"
+	"golang.org/x/time/rate"
+)
+
+type IPRateLimiter struct {
+	ips *lru.Cache
+	mu  *sync.RWMutex
+	r   rate.Limit
+	b   int
+}
+
+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: cache,
+		mu:  &sync.RWMutex{},
+		r:   r,
+		b:   b,
+	}
+
+	return i, nil
+}
+
+func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	limiter := rate.NewLimiter(i.r, i.b)
+
+	i.ips.Add(ip, limiter)
+
+	return limiter
+}
+
+func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	if limiter, exists := i.ips.Get(ip); exists {
+		return limiter.(*rate.Limiter)
+	}
+
+	return i.AddIP(ip)
+}