cerc-io/lotus
16
0
Fork 0
Code Issues Pull Requests Packages Projects Releases27 Wiki Activity

use key type

Browse Source
This commit is contained in:
Cory Schwartz 2022-06-06 16:37:11 -07:00
parent 4c51e4ccda
commit ea51732163
2 changed files with 6 additions and 2 deletions

6
gateway/handler.go
View File

@ -19,6 +19,10 @@ import (
"golang.org/x/time/rate" "golang.org/x/time/rate"
) )
type perConnLimiterKeyType string
const perConnLimiterKey = "limiter"
// Handler returns a gateway http.Handler, to be mounted as-is on the server. // Handler returns a gateway http.Handler, to be mounted as-is on the server.
func Handler(gwapi lapi.Gateway, api lapi.FullNode, rateLimit int64, connPerMinute int64, 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() m := mux.NewRouter()
@ -75,7 +79,7 @@ type RateLimiterHandler struct {
} }
func (h RateLimiterHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { func (h RateLimiterHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
r2 := r.WithContext(context.WithValue(r.Context(), "limiter", h.limiter)) r2 := r.WithContext(context.WithValue(r.Context(), perConnLimiterKey, h.limiter))
h.handler.ServeHTTP(w, r2) h.handler.ServeHTTP(w, r2)
} }

2
gateway/node.go
View File

@ -166,7 +166,7 @@ func (gw *Node) checkTimestamp(at time.Time) error {
func (gw *Node) limit(ctx context.Context, tokens int) error { func (gw *Node) limit(ctx context.Context, tokens int) error {
ctx2, cancel := context.WithTimeout(ctx, gw.rateLimitTimeout) ctx2, cancel := context.WithTimeout(ctx, gw.rateLimitTimeout)
defer cancel() defer cancel()
if perConnLimiter, ok := ctx2.Value("limiter").(*rate.Limiter); ok { if perConnLimiter, ok := ctx2.Value(perConnLimiterKey).(*rate.Limiter); ok {
err := perConnLimiter.WaitN(ctx2, tokens) err := perConnLimiter.WaitN(ctx2, tokens)
if err != nil { if err != nil {
return fmt.Errorf("connection limited. %w", err) return fmt.Errorf("connection limited. %w", err)