faucet: rate limit initial implementation

cmd/faucet/faucet.go

@@ -49,6 +49,7 @@ import (
 	"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 @@ type faucet struct {
 
 	bep2eInfos map[string]bep2eInfo
 	bep2eAbi   abi.ABI
+
+	limiter *IPRateLimiter
 }
 
 // wsConn wraps a websocket connection with a write mutex as the underlying
@@ -245,6 +248,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
 	}, nil
 }
 
@@ -272,6 +276,19 @@ func (f *faucet) webHandler(w http.ResponseWriter, r *http.Request) {
 
 // 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"), ",")
+		ip = strings.TrimSpace(ips[len(ips)-1])
+	}
+
+	limiter := f.limiter.GetLimiter(ip)
+	if !limiter.Allow() {
+		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 {
@@ -901,3 +918,46 @@ func getGenesis(genesisFlag string, goerliFlag bool, sepoliaFlag bool) (*core.Ge
 		return nil, errors.New("no genesis flag provided")
 	}
 }
+
+type IPRateLimiter struct {
+	ips map[string]*rate.Limiter
+	mu  *sync.RWMutex
+	r   rate.Limit
+	b   int
+}
+
+func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
+	i := &IPRateLimiter{
+		ips: make(map[string]*rate.Limiter),
+		mu:  &sync.RWMutex{},
+		r:   r,
+		b:   b,
+	}
+
+	return i
+}
+
+func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	limiter := rate.NewLimiter(i.r, i.b)
+
+	i.ips[ip] = limiter
+
+	return limiter
+}
+
+func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
+	i.mu.Lock()
+	limiter, exists := i.ips[ip]
+
+	if !exists {
+		i.mu.Unlock()
+		return i.AddIP(ip)
+	}
+
+	i.mu.Unlock()
+
+	return limiter
+}

cmd/faucet/faucet_rate_limit_test.go

@@ -0,0 +1,125 @@
+package main
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+// Mock mockfaucet struct
+type mockfaucet struct {
+	limiter *MockIPRateLimiter
+}
+
+// Mock MockIPRateLimiter struct and methods
+type MockIPRateLimiter struct {
+	ips map[string]*rate.Limiter
+	mu  *sync.RWMutex
+	r   rate.Limit
+	b   int
+}
+
+func MockNewIPRateLimiter(r rate.Limit, b int) *MockIPRateLimiter {
+	i := &MockIPRateLimiter{
+		ips: make(map[string]*rate.Limiter),
+		mu:  &sync.RWMutex{},
+		r:   r,
+		b:   b,
+	}
+	return i
+}
+
+func (i *MockIPRateLimiter) AddIP(ip string) *rate.Limiter {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+	limiter := rate.NewLimiter(i.r, i.b)
+	i.ips[ip] = limiter
+	return limiter
+}
+
+func (i *MockIPRateLimiter) GetLimiter(ip string) *rate.Limiter {
+	i.mu.Lock()
+	limiter, exists := i.ips[ip]
+	if !exists {
+		i.mu.Unlock()
+		return i.AddIP(ip)
+	}
+	i.mu.Unlock()
+	return limiter
+}
+
+// Mock apiHandler
+func (f *mockfaucet) apiHandler(w http.ResponseWriter, r *http.Request) {
+	//ip := r.RemoteAddr
+	ip := "test-ip" // Use a constant IP for testing
+	limiter := f.limiter.GetLimiter(ip)
+	if !limiter.Allow() {
+		http.Error(w, "Too many requests", http.StatusTooManyRequests)
+		return
+	}
+	w.WriteHeader(http.StatusOK)
+}
+
+// Mock mocknewFaucet function
+func mocknewFaucet() *mockfaucet {
+	return &mockfaucet{
+		limiter: MockNewIPRateLimiter(rate.Limit(1), 2), // 1 request per second, burst of 2
+	}
+}
+
+func TestMockFaucetRateLimiting(t *testing.T) {
+	// Create a mockfaucet with rate limiting
+	f := mocknewFaucet()
+
+	// 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
+
+	results := make([]int, 5)
+
+	for i := 0; i < 5; i++ {
+		results[i] = makeRequest()
+		time.Sleep(10 * time.Millisecond) // Small delay to ensure order
+	}
+
+	// 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/faucet_test.go

@@ -17,9 +17,17 @@
 package main
 
 import (
+	"net/http"
 	"testing"
 
+	"net/http/httptest"
+	"sync"
+	"time"
+
 	"github.com/ethereum/go-ethereum/common"
+
+	"github.com/ethereum/go-ethereum/core"
+	"github.com/ethereum/go-ethereum/crypto"
 )
 
 func TestFacebook(t *testing.T) {
@@ -43,3 +51,75 @@ func TestFacebook(t *testing.T) {
 		}
 	}
 }
+
+func TestFaucetRateLimiting(t *testing.T) {
+	// Create a minimal mockfaucet instance for testing
+	privateKey, _ := crypto.GenerateKey()
+	faucetAddr := crypto.PubkeyToAddress(privateKey.PublicKey)
+
+	config := &core.Genesis{
+		Alloc: core.GenesisAlloc{
+			faucetAddr: {Balance: common.Big1},
+		},
+	}
+
+	// Create a mockfaucet with rate limiting (1 request per second, burst of 2)
+	f, err := newFaucet(config, "http://localhost:8545", nil, []byte{}, nil)
+	if err != nil {
+		t.Fatalf("Failed to create mockfaucet: %v", err)
+	}
+	f.limiter = NewIPRateLimiter(1, 2)
+
+	// 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)
+	}
+}