tornadocontrib/infra
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Move base rate rate limit check inside handleBatchRPC() (#37)

Browse Source 
* move base rate rate limit check

* fix comment
This commit is contained in:
cody-wang-cb 2024-07-31 16:25:53 -04:00 committed by GitHub
parent ec84826c4c
commit 0fb094feb4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 33 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
proxyd/integration_tests/rate_limit_test.go
View File

@ -16,7 +16,6 @@ type resWithCode struct {
res []byte
}
const frontendOverLimitResponse = `{"error":{"code":-32016,"message":"over rate limit with special message"},"id":null,"jsonrpc":"2.0"}`
const frontendOverLimitResponseWithID = `{"error":{"code":-32016,"message":"over rate limit with special message"},"id":999,"jsonrpc":"2.0"}`
var ethChainID = "eth_chainId"
@ -37,7 +36,7 @@ func TestFrontendMaxRPSLimit(t *testing.T) {
limitedRes, codes := spamReqs(t, client, ethChainID, 429, 3)
require.Equal(t, 1, codes[429])
require.Equal(t, 2, codes[200])
RequireEqualJSON(t, []byte(frontendOverLimitResponse), limitedRes)
RequireEqualJSON(t, []byte(frontendOverLimitResponseWithID), limitedRes)
})
t.Run("exempt user agent over limit", func(t *testing.T) {
@ -106,6 +105,25 @@ func TestFrontendMaxRPSLimit(t *testing.T) {
time.Sleep(time.Second)
t.Run("Batch RPC with some requests rate limited", func(t *testing.T) {
client := NewProxydClient("http://127.0.0.1:8545")
req := NewRPCReq("123", "eth_chainId", nil)
out, code, err := client.SendBatchRPC(req, req, req)
require.NoError(t, err)
var res []proxyd.RPCRes
require.NoError(t, json.Unmarshal(out, &res))
expCode := proxyd.ErrOverRateLimit.Code
require.Equal(t, 200, code)
require.Equal(t, 3, len(res))
require.Nil(t, res[0].Error)
require.Nil(t, res[1].Error)
// base rate = 2, so the third request should be rate limited
require.Equal(t, expCode, res[2].Error.Code)
})
time.Sleep(time.Second)
t.Run("RPC override in batch exempt", func(t *testing.T) {
h := make(http.Header)
h.Set("User-Agent", "exempt_agent")

30
proxyd/server.go
View File

@ -301,20 +301,6 @@ func (s *Server) HandleRPC(w http.ResponseWriter, r *http.Request) {
return !ok
}
if isLimited("") {
RecordRPCError(ctx, BackendProxyd, "unknown", ErrOverRateLimit)
log.Warn(
"rate limited request",
"req_id", GetReqID(ctx),
"auth", GetAuthCtx(ctx),
"user_agent", userAgent,
"origin", origin,
"remote_ip", xff,
)
writeRPCError(ctx, w, nil, ErrOverRateLimit)
return
}
log.Info(
"received RPC request",
"req_id", GetReqID(ctx),
@ -469,10 +455,20 @@ func (s *Server) handleBatchRPC(ctx context.Context, reqs []json.RawMessage, isL
continue
}
// Take base rate limit first
if isLimited("") {
log.Info(
"rate limited individual RPC in a batch request",
"source", "rpc",
"req_id", parsedReq.ID,
"method", parsedReq.Method,
)
RecordRPCError(ctx, BackendProxyd, parsedReq.Method, ErrOverRateLimit)
responses[i] = NewRPCErrorRes(parsedReq.ID, ErrOverRateLimit)
continue
}
// Take rate limit for specific methods.
// NOTE: eventually, this should apply to all batch requests. However,
// since we don't have data right now on the size of each batch, we
// only apply this to the methods that have an additional rate limit.
if _, ok := s.overrideLims[parsedReq.Method]; ok && isLimited(parsedReq.Method) {
log.Info(
"rate limited specific RPC",