Add rate limits to the fountain by @travisperson

cmd/lotus-fountain/main.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"time"
 
 	rice "github.com/GeertJohan/go.rice"
 	logging "github.com/ipfs/go-log"
@@ -88,6 +89,22 @@ var runCmd = &cli.Command{
 			ctx:  ctx,
 			api:  nodeApi,
 			from: from,
+			limiter: NewLimiter(LimiterConfig{
+				TotalRate:   time.Second,
+				TotalBurst:  20,
+				IPRate:      5 * time.Minute,
+				IPBurst:     5,
+				WalletRate:  time.Hour,
+				WalletBurst: 1,
+			}),
+			colLimiter: NewLimiter(LimiterConfig{
+				TotalRate:   time.Second,
+				TotalBurst:  20,
+				IPRate:      24 * time.Hour,
+				IPBurst:     1,
+				WalletRate:  24 * 364 * time.Hour,
+				WalletBurst: 1,
+			}),
 		}
 
 		http.Handle("/", http.FileServer(rice.MustFindBox("site").HTTPBox()))
@@ -110,9 +127,25 @@ type handler struct {
 	api api.FullNode
 
 	from address.Address
+
+	limiter    *Limiter
+	colLimiter *Limiter
 }
 
 func (h *handler) send(w http.ResponseWriter, r *http.Request) {
+	// General limiter to allow throttling all messages that can make it into the mpool
+	if !h.limiter.Allow() {
+		http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
+		return
+	}
+
+	// Limit based on IP
+	limiter := h.limiter.GetIPLimiter(r.RemoteAddr)
+	if !limiter.Allow() {
+		http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
+		return
+	}
+
 	to, err := address.NewFromString(r.FormValue("address"))
 	if err != nil {
 		w.WriteHeader(400)
@@ -120,6 +153,13 @@ func (h *handler) send(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	// Limit based on wallet address
+	limiter = h.limiter.GetWalletLimiter(to.String())
+	if !limiter.Allow() {
+		http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
+		return
+	}
+
 	smsg, err := h.api.MpoolPushMessage(h.ctx, &types.Message{
 		Value: sendPerRequest,
 		From:  h.from,

cmd/lotus-fountain/rate_limiter.go

@@ -0,0 +1,95 @@
+
+package main
+
+import (
+	"sync"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+type Limiter struct {
+	control *rate.Limiter
+
+	ips     map[string]*rate.Limiter
+	wallets map[string]*rate.Limiter
+	mu      *sync.RWMutex
+
+	config LimiterConfig
+}
+
+type LimiterConfig struct {
+	TotalRate  time.Duration
+	TotalBurst int
+
+	IPRate  time.Duration
+	IPBurst int
+
+	WalletRate  time.Duration
+	WalletBurst int
+}
+
+func NewLimiter(c LimiterConfig) *Limiter {
+	return &Limiter{
+		control: rate.NewLimiter(rate.Every(c.TotalRate), c.TotalBurst),
+		mu:      &sync.RWMutex{},
+		ips:     make(map[string]*rate.Limiter),
+		wallets: make(map[string]*rate.Limiter),
+
+		config: c,
+	}
+}
+
+func (i *Limiter) Allow() bool {
+	return i.control.Allow()
+}
+
+func (i *Limiter) AddIPLimiter(ip string) *rate.Limiter {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	limiter := rate.NewLimiter(rate.Every(i.config.IPRate), i.config.IPBurst)
+
+	i.ips[ip] = limiter
+
+	return limiter
+}
+
+func (i *Limiter) GetIPLimiter(ip string) *rate.Limiter {
+	i.mu.Lock()
+	limiter, exists := i.ips[ip]
+
+	if !exists {
+		i.mu.Unlock()
+		return i.AddIPLimiter(ip)
+	}
+
+	i.mu.Unlock()
+
+	return limiter
+}
+
+func (i *Limiter) AddWalletLimiter(addr string) *rate.Limiter {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	limiter := rate.NewLimiter(rate.Every(i.config.WalletRate), i.config.WalletBurst)
+
+	i.wallets[addr] = limiter
+
+	return limiter
+}
+
+func (i *Limiter) GetWalletLimiter(wallet string) *rate.Limiter {
+	i.mu.Lock()
+	limiter, exists := i.wallets[wallet]
+
+	if !exists {
+		i.mu.Unlock()
+		return i.AddWalletLimiter(wallet)
+	}
+
+	i.mu.Unlock()
+
+	return limiter
+}

cmd/lotus-fountain/rate_limiter_test.go

@@ -0,0 +1,38 @@
+package main
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRateLimit(t *testing.T) {
+	limiter := NewLimiter(LimiterConfig{
+		TotalRate:   time.Second,
+		TotalBurst:  20,
+		IPRate:      time.Second,
+		IPBurst:     1,
+		WalletRate:  time.Second,
+		WalletBurst: 1,
+	})
+
+	for i := 0; i < 20; i++ {
+		assert.True(t, limiter.Allow())
+	}
+
+	assert.False(t, limiter.Allow())
+
+	time.Sleep(time.Second)
+	assert.True(t, limiter.Allow())
+
+	assert.True(t, limiter.GetIPLimiter("127.0.0.1").Allow())
+	assert.False(t, limiter.GetIPLimiter("127.0.0.1").Allow())
+	time.Sleep(time.Second)
+	assert.True(t, limiter.GetIPLimiter("127.0.0.1").Allow())
+
+	assert.True(t, limiter.GetWalletLimiter("abc123").Allow())
+	assert.False(t, limiter.GetWalletLimiter("abc123").Allow())
+	time.Sleep(time.Second)
+	assert.True(t, limiter.GetWalletLimiter("abc123").Allow())
+}