Skip to content

Commit

Permalink
shuttles
Browse files Browse the repository at this point in the history
  • Loading branch information
anjor committed Dec 14, 2022
1 parent 905d6af commit dda8fc7
Show file tree
Hide file tree
Showing 5 changed files with 37 additions and 24 deletions.
2 changes: 1 addition & 1 deletion api/v1/api.go
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ func NewAPIV1(
// @securityDefinitions.Bearer.name Authorization
func (s *apiV1) RegisterRoutes(e *echo.Echo) {

e.Use(middleware.RateLimiterWithConfig(configureRateLimiter(s.cfg.RateLimit)))
e.Use(middleware.RateLimiterWithConfig(util.ConfigureRateLimiter(s.cfg.RateLimit)))
e.POST("/register", s.handleRegisterUser)
e.POST("/login", s.handleLoginUser)
e.GET("/health", s.handleHealth)
Expand Down
21 changes: 0 additions & 21 deletions api/v1/handlers.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,6 @@ import (
"encoding/hex"
"encoding/json"
"fmt"
"github.com/labstack/echo/v4/middleware"
"golang.org/x/time/rate"
"io"
"io/ioutil"
"math/rand"
Expand Down Expand Up @@ -121,25 +119,6 @@ func withUser(f func(echo.Context, *util.User) error) func(echo.Context) error {
}
}

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").(*util.User)
if !ok {
id = ctx.RealIP()
} else {
id = user.UUID
}
return id, nil
}

config.Store = middleware.NewRateLimiterMemoryStore(rateLimit)

return config
}

// handleStats godoc
// @Summary Get content statistics
// @Description This endpoint is used to get content statistics. Every content stored in the network (estuary) is tracked by a unique ID which can be used to get information about the content. This endpoint will allow the consumer to get the collected stats of a content
Expand Down
9 changes: 7 additions & 2 deletions cmd/estuary-shuttle/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"encoding/json"
"flag"
"fmt"
"golang.org/x/time/rate"
"io"
"io/ioutil"
"net/http"
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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:
}
}
Expand Down Expand Up @@ -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...),
},
Expand Down Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions config/shuttle.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"`
Expand Down Expand Up @@ -82,6 +84,7 @@ func NewShuttle(appVersion string) *Shuttle {
Private: false,
Dev: false,
NoReloadPinQueue: false,
RateLimit: rate.Limit(20),

Content: Content{
DisableLocalAdding: false,
Expand Down
26 changes: 26 additions & 0 deletions util/rate_limiter.go
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit dda8fc7

Please sign in to comment.