In memory rate limiter. (#767)

* implement in memory rate limiter

* make it configurable

* make it configurable

* shuttles

api/v1/api.go

@@ -14,6 +14,7 @@ import (
 	"github.com/application-research/filclient"
 	"github.com/filecoin-project/lotus/api"
 	"github.com/labstack/echo/v4"
+	"github.com/labstack/echo/v4/middleware"
 	"github.com/whyrusleeping/memo"
 	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
@@ -91,6 +92,7 @@ func NewAPIV1(
 // @securityDefinitions.Bearer.name Authorization
 func (s *apiV1) RegisterRoutes(e *echo.Echo) {
 
+	e.Use(middleware.RateLimiterWithConfig(util.ConfigureRateLimiter(s.cfg.RateLimit)))
 	e.POST("/register", s.handleRegisterUser)
 	e.POST("/login", s.handleLoginUser)
 	e.GET("/health", s.handleHealth)

cmd/estuary-shuttle/main.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"flag"
 	"fmt"
+	"golang.org/x/time/rate"
 	"io"
 	"io/ioutil"
 	"net/http"
@@ -86,7 +87,7 @@ const (
 	ColDir  = "dir"
 )
 
-//#nosec G104 - it's not common to treat SetLogLevel error return
+// #nosec G104 - it's not common to treat SetLogLevel error return
 func before(cctx *cli.Context) error {
 	level := util.LogLevel
 
@@ -196,6 +197,8 @@ func overrideSetOptions(flags []cli.Flag, cctx *cli.Context, cfg *config.Shuttle
 			cfg.RpcEngine.Queue.Enabled = cctx.Bool("queue-eng-enabled")
 		case "queue-eng-consumers":
 			cfg.RpcEngine.Queue.Consumers = cctx.Int("queue-eng-consumers")
+		case "rate-limit":
+			cfg.RateLimit = rate.Limit(cctx.Float64("rate-limit"))
 		default:
 		}
 	}
@@ -325,7 +328,7 @@ func main() {
 			Value: cfg.Dev,
 		},
 		&cli.StringSliceFlag{
-			Name: "announce-addr",
+			Name:  "announce-addr",
 			Usage: "specify multiaddrs that this node can be connected to	",
 			Value: cli.NewStringSlice(cfg.Node.AnnounceAddrs...),
 		},
@@ -1154,6 +1157,8 @@ func (s *Shuttle) ServeAPI() error {
 		e.Use(middleware.Logger())
 	}
 
+	e.Use(middleware.RateLimiterWithConfig(util.ConfigureRateLimiter(s.shuttleConfig.RateLimit)))
+
 	e.Use(s.tracingMiddleware)
 	e.Use(util.AppVersionMiddleware(s.shuttleConfig.AppVersion))
 	e.HTTPErrorHandler = util.ErrorHandler

config/estuary.go

@@ -2,6 +2,7 @@ package config
 
 import (
 	"errors"
+	"golang.org/x/time/rate"
 	"path/filepath"
 	"time"
 
@@ -37,6 +38,7 @@ type Estuary struct {
 	Replication            int           `json:"replication"`
 	RpcEngine              RpcEngine     `json:"rpc_engine"`
 	Pinning                Pinning       `json:"pinning"`
+	RateLimit              rate.Limit    `json:"rate_limit"`
 }
 
 func (cfg *Estuary) Load(filename string) error {
@@ -84,6 +86,7 @@ func NewEstuary(appVersion string) *Estuary {
 		DisableFilecoinStorage: false,
 		DisableSwaggerEndpoint: false,
 		DisableAutoRetrieve:    true,
+		RateLimit:              rate.Limit(20),
 
 		Deal: Deal{
 			IsDisabled:            false,

config/shuttle.go

@@ -2,6 +2,7 @@ package config
 
 import (
 	"errors"
+	"golang.org/x/time/rate"
 	"path/filepath"
 
 	rcmgr "github.com/libp2p/go-libp2p/p2p/host/resource-manager"
@@ -28,6 +29,7 @@ type Shuttle struct {
 	Private            bool          `json:"private"`
 	Dev                bool          `json:"dev"`
 	NoReloadPinQueue   bool          `json:"no_reload_pin_queue"`
+	RateLimit          rate.Limit    `json:"rate_limit"`
 	Node               Node          `json:"node"`
 	Jaeger             Jaeger        `json:"jaeger"`
 	Content            Content       `json:"content"`
@@ -82,6 +84,7 @@ func NewShuttle(appVersion string) *Shuttle {
 		Private:            false,
 		Dev:                false,
 		NoReloadPinQueue:   false,
+		RateLimit:          rate.Limit(20),
 
 		Content: Content{
 			DisableLocalAdding: false,

main.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	"golang.org/x/time/rate"
 	"os"
 	"path/filepath"
 	"strings"
@@ -223,6 +224,9 @@ func overrideSetOptions(flags []cli.Flag, cctx *cli.Context, cfg *config.Estuary
 			}
 			cfg.Deal.MaxVerifiedPrice = abi.TokenAmount(maxVerifiedPrice)
 
+		case "rate-limit":
+			cfg.RateLimit = rate.Limit(cctx.Float64("rate-limit"))
+
 		default:
 		}
 	}

util/rate_limiter.go

@@ -0,0 +1,26 @@
+package util
+
+import (
+	"github.com/labstack/echo/v4"
+	"github.com/labstack/echo/v4/middleware"
+	"golang.org/x/time/rate"
+)
+
+func ConfigureRateLimiter(rateLimit rate.Limit) middleware.RateLimiterConfig {
+	config := middleware.DefaultRateLimiterConfig
+
+	config.IdentifierExtractor = func(ctx echo.Context) (string, error) {
+		var id string
+		user, ok := ctx.Get("user").(User)
+		if !ok {
+			id = ctx.RealIP()
+		} else {
+			id = user.UUID
+		}
+		return id, nil
+	}
+
+	config.Store = middleware.NewRateLimiterMemoryStore(rateLimit)
+
+	return config
+}