feat: Add rate limiting to the lotus gateway

cmd/lotus-gateway/main.go

@@ -133,13 +133,33 @@ var runCmd = &cli.Command{
 			Usage: "maximum number of blocks to search back through for message inclusion",
 			Value: int64(gateway.DefaultStateWaitLookbackLimit),
 		},
+		&cli.Int64Flag{
+			Name:  "rate-limit",
+			Usage: "rate-limit API calls. Use 0 to disable",
+			Value: 0,
+		},
+		&cli.Int64Flag{
+			Name:  "per-conn-rate-limit",
+			Usage: "rate-limit API calls per each connection. Use 0 to disable",
+			Value: 0,
+		},
+		&cli.DurationFlag{
+			Name:  "rate-limit-timeout",
+			Usage: "the maximum time to wait for the rate limter before returning an error to clients",
+			Value: gateway.DefaultRateLimitTimeout,
+		},
+		&cli.Int64Flag{
+			Name:  "conn-per-minute",
+			Usage: "The number of incomming connections to accept from a single IP per minute.  Use 0 to disable",
+			Value: 0,
+		},
 	},
 	Action: func(cctx *cli.Context) error {
 		log.Info("Starting lotus gateway")
 
 		// Register all metric views
 		if err := view.Register(
-			metrics.ChainNodeViews...,
+			metrics.GatewayNodeViews...,
 		); err != nil {
 			log.Fatalf("Cannot register the view: %v", err)
 		}
@@ -151,9 +171,13 @@ var runCmd = &cli.Command{
 		defer closer()
 
 		var (
-			lookbackCap  = cctx.Duration("api-max-lookback")
-			address      = cctx.String("listen")
-			waitLookback = abi.ChainEpoch(cctx.Int64("api-wait-lookback-limit"))
+			lookbackCap      = cctx.Duration("api-max-lookback")
+			address          = cctx.String("listen")
+			waitLookback     = abi.ChainEpoch(cctx.Int64("api-wait-lookback-limit"))
+			rateLimit        = cctx.Int64("rate-limit")
+			perConnRateLimit = cctx.Int64("per-conn-rate-limit")
+			rateLimitTimeout = cctx.Duration("rate-limit-timeout")
+			connPerMinute    = cctx.Int64("conn-per-minute")
 		)
 
 		serverOptions := make([]jsonrpc.ServerOption, 0)
@@ -173,8 +197,8 @@ var runCmd = &cli.Command{
 			return xerrors.Errorf("failed to convert endpoint address to multiaddr: %w", err)
 		}
 
-		gwapi := gateway.NewNode(api, lookbackCap, waitLookback)
-		h, err := gateway.Handler(gwapi, api, serverOptions...)
+		gwapi := gateway.NewNode(api, lookbackCap, waitLookback, rateLimit, rateLimitTimeout)
+		h, err := gateway.Handler(gwapi, api, perConnRateLimit, connPerMinute, serverOptions...)
 		if err != nil {
 			return xerrors.Errorf("failed to set up gateway HTTP handler")
 		}

gateway/handler.go

@@ -1,7 +1,11 @@
 package gateway
 
 import (
+	"context"
+	"net"
 	"net/http"
+	"sync"
+	"time"
 
 	"contrib.go.opencensus.io/exporter/prometheus"
 	"github.com/filecoin-project/go-jsonrpc"
@@ -12,10 +16,15 @@ import (
 	"github.com/filecoin-project/lotus/node"
 	"github.com/gorilla/mux"
 	promclient "github.com/prometheus/client_golang/prometheus"
+	"golang.org/x/time/rate"
 )
 
+type perConnLimiterKeyType string
+
+const perConnLimiterKey perConnLimiterKeyType = "limiter"
+
 // Handler returns a gateway http.Handler, to be mounted as-is on the server.
-func Handler(gwapi lapi.Gateway, api lapi.FullNode, opts ...jsonrpc.ServerOption) (http.Handler, error) {
+func Handler(gwapi lapi.Gateway, api lapi.FullNode, rateLimit int64, connPerMinute int64, opts ...jsonrpc.ServerOption) (http.Handler, error) {
 	m := mux.NewRouter()
 
 	serveRpc := func(path string, hnd interface{}) {
@@ -49,5 +58,95 @@ func Handler(gwapi lapi.Gateway, api lapi.FullNode, opts ...jsonrpc.ServerOption
 		Next:   mux.ServeHTTP,
 	}*/
 
-	return m, nil
+	rlh := NewRateLimiterHandler(m, rateLimit)
+	clh := NewConnectionRateLimiterHandler(rlh, connPerMinute)
+	return clh, nil
+}
+
+func NewRateLimiterHandler(handler http.Handler, rateLimit int64) *RateLimiterHandler {
+	limiter := limiterFromRateLimit(rateLimit)
+
+	return &RateLimiterHandler{
+		handler: handler,
+		limiter: limiter,
+	}
+}
+
+// Adds a rate limiter to the request context for per-connection rate limiting
+type RateLimiterHandler struct {
+	handler http.Handler
+	limiter *rate.Limiter
+}
+
+func (h RateLimiterHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	r2 := r.WithContext(context.WithValue(r.Context(), perConnLimiterKey, h.limiter))
+	h.handler.ServeHTTP(w, r2)
+}
+
+// this blocks new connections if there have already been too many.
+func NewConnectionRateLimiterHandler(handler http.Handler, connPerMinute int64) *ConnectionRateLimiterHandler {
+	ipmap := make(map[string]int64)
+	return &ConnectionRateLimiterHandler{
+		ipmap:         ipmap,
+		connPerMinute: connPerMinute,
+		handler:       handler,
+	}
+}
+
+type ConnectionRateLimiterHandler struct {
+	mu            sync.Mutex
+	ipmap         map[string]int64
+	connPerMinute int64
+	handler       http.Handler
+}
+
+func (h *ConnectionRateLimiterHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if h.connPerMinute == 0 {
+		h.handler.ServeHTTP(w, r)
+		return
+	}
+	host, _, err := net.SplitHostPort(r.RemoteAddr)
+	if err != nil {
+		w.WriteHeader(http.StatusInternalServerError)
+		return
+	}
+
+	h.mu.Lock()
+	seen, ok := h.ipmap[host]
+	if !ok {
+		h.ipmap[host] = 1
+		h.mu.Unlock()
+		h.handler.ServeHTTP(w, r)
+		return
+	}
+	// rate limited
+	if seen > h.connPerMinute {
+		h.mu.Unlock()
+		w.WriteHeader(http.StatusTooManyRequests)
+		return
+	}
+	h.ipmap[host] = seen + 1
+	h.mu.Unlock()
+	go func() {
+		select {
+		case <-time.After(time.Minute):
+			h.mu.Lock()
+			defer h.mu.Unlock()
+			h.ipmap[host] = h.ipmap[host] - 1
+			if h.ipmap[host] <= 0 {
+				delete(h.ipmap, host)
+			}
+		}
+	}()
+	h.handler.ServeHTTP(w, r)
+}
+
+func limiterFromRateLimit(rateLimit int64) *rate.Limiter {
+	var limit rate.Limit
+	if rateLimit == 0 {
+		limit = rate.Inf
+	} else {
+		limit = rate.Every(time.Second / time.Duration(rateLimit))
+	}
+	return rate.NewLimiter(limit, stateRateLimitTokens)
 }