From 2da46a641367331d864e9b0bacce2756f4cd2db6 Mon Sep 17 00:00:00 2001 From: Jan Michael Auer Date: Fri, 26 Jun 2020 12:26:38 +0200 Subject: [PATCH 1/2] feat(server): Enforce attachment size in rate limits --- relay-quotas/src/is_rate_limited.lua | 67 +++++++++++++-------- relay-quotas/src/redis.rs | 87 ++++++++++++++++++++++----- relay-server/src/actors/events.rs | 4 +- tests/integration/test_attachments.py | 9 ++- 4 files changed, 124 insertions(+), 43 deletions(-) diff --git a/relay-quotas/src/is_rate_limited.lua b/relay-quotas/src/is_rate_limited.lua index 4fbb670b759..a4004ed57a9 100644 --- a/relay-quotas/src/is_rate_limited.lua +++ b/relay-quotas/src/is_rate_limited.lua @@ -1,45 +1,64 @@ -- Check a collection of quota counters to identify if an item should be rate --- limited. Values provided as ``KEYS`` specify the keys of the counters to --- check and the keys of counters to subtract, and values provided as ``ARGV`` --- specify the maximum value (quota limit) and expiration time for each key. +-- limited. For each quota, repeat the same set of ``KEYS`` and ``ARGV``: -- --- For example, to check a quota ``foo`` that has a corresponding refund/negative --- counter "subtract_from_foo", a limit of 10 items and expires at the Unix timestamp --- ``100``, as well as a quota ``bar`` that has a corresponding refund/negative --- counter "subtract_from_bar" limit of 20 items and should expire at the Unix --- timestamp ``100``, the ``KEYS`` and ``ARGV`` values would be as follows: +-- ``KEYS`` (2 per quota): +-- * [string] Key of the counter. +-- * [string] Key of the refund counter. -- --- KEYS = {"foo", "subtract_from_foo", "bar", "subtract_from_bar"} --- ARGV = {10, 100, 20, 100} +-- ``ARGV`` (3 per quota): +-- * [number] Quota limit. Can be ``-1`` for unlimited quotas. +-- * [number] Expiration time in seconds for the key. +-- * [number] Quantity to increment the quota by. -- --- If all checks pass (the item is accepted), the counters for all quotas are --- incremented. If any checks fail (the item is rejected), the counters for all --- quotas are unaffected. The result is a Lua table/array (Redis multi bulk --- reply) that specifies whether or not the item was *rejected* based on the --- provided limit. -assert(#KEYS == #ARGV, "incorrect number of keys and arguments provided") -assert(#KEYS % 2 == 0, "there must be an even number of keys") +-- For example, to check the following two quotas each with a timeout of 10 minutes: +-- * Key ``foo``, refund key ``foo_refund``, limit ``10``; quantity ``5`` +-- * Key ``bar``, refund key ``bar_refund``, limit ``20``; quantity ``1`` +-- +-- Send these values: +-- +-- KEYS = {"foo", "foo_refund", "bar", "bar_refund"} +-- ARGV = {10, 600, 5, 20, 600, 1} +-- +-- The script applies the following logic: +-- * If all checks pass, the item is accepted and the counters for all quotas +-- are incremented. +-- * If any check fails, the item is rejected and the counters for all remain +-- unchanged. +-- +-- The result is a Lua table/array (Redis multi bulk reply) that specifies +-- whether or not the item was *rejected* based on the provided limit. +assert(#KEYS % 2 == 0, "there must be 2 keys per quota") +assert(#ARGV % 3 == 0, "there must be 3 args per quota") +assert(#KEYS / 2 == #ARGV / 3, "incorrect number of keys and arguments provided") local results = {} local failed = false -for i=1, #KEYS, 2 do - local limit = tonumber(ARGV[i]) +local num_quotas = #KEYS / 2 +for i=0, num_quotas - 1 do + local k = i * 2 + 1 + local v = i * 3 + 1 + + local limit = tonumber(ARGV[v]) + local quantity = tonumber(ARGV[v+2]) local rejected = false -- limit=-1 means "no limit" if limit >= 0 then - rejected = (redis.call('GET', KEYS[i]) or 0) - (redis.call('GET', KEYS[i + 1]) or 0) + 1 > limit + rejected = (redis.call('GET', KEYS[k]) or 0) - (redis.call('GET', KEYS[k + 1]) or 0) + quantity > limit end if rejected then failed = true end - results[(i + 1) / 2] = rejected + results[i + 1] = rejected end if not failed then - for i=1, #KEYS, 2 do - redis.call('INCR', KEYS[i]) - redis.call('EXPIREAT', KEYS[i], ARGV[i + 1]) + for i=0, num_quotas - 1 do + local k = i * 2 + 1 + local v = i * 3 + 1 + + redis.call('INCRBY', KEYS[k], ARGV[v + 2]) + redis.call('EXPIREAT', KEYS[k], ARGV[v + 1]) end end diff --git a/relay-quotas/src/redis.rs b/relay-quotas/src/redis.rs index d5da70f269d..db93ad94ef6 100644 --- a/relay-quotas/src/redis.rs +++ b/relay-quotas/src/redis.rs @@ -171,6 +171,7 @@ impl RedisRateLimiter { &self, quotas: &[Quota], item_scoping: ItemScoping<'_>, + quantity: usize, ) -> Result { let timestamp = UnixTimestamp::now(); @@ -197,6 +198,7 @@ impl RedisRateLimiter { invocation.arg(quota.limit()); invocation.arg(quota.expiry().as_secs()); + invocation.arg(quantity); tracked_quotas.push(quota); } else { @@ -296,7 +298,7 @@ mod tests { }; let rate_limits: Vec = RATE_LIMITER - .is_rate_limited(quotas, scoping) + .is_rate_limited(quotas, scoping, 1) .expect("rate limiting failed") .into_iter() .collect(); @@ -336,7 +338,7 @@ mod tests { for i in 0..10 { let rate_limits: Vec = RATE_LIMITER - .is_rate_limited(quotas, scoping) + .is_rate_limited(quotas, scoping, 1) .expect("rate limiting failed") .into_iter() .collect(); @@ -370,7 +372,7 @@ mod tests { }; let rate_limits: Vec = RATE_LIMITER - .is_rate_limited(&[], scoping) + .is_rate_limited(&[], scoping, 1) .expect("rate limiting failed") .into_iter() .collect(); @@ -413,7 +415,7 @@ mod tests { for i in 0..1 { let rate_limits: Vec = RATE_LIMITER - .is_rate_limited(quotas, scoping) + .is_rate_limited(quotas, scoping, 1) .expect("rate limiting failed") .into_iter() .collect(); @@ -434,6 +436,51 @@ mod tests { } } + #[test] + fn test_quota_with_quantity() { + let quotas = &[Quota { + id: Some(format!("test_quantity_quota_{:?}", SystemTime::now())), + categories: DataCategories::new(), + scope: QuotaScope::Organization, + scope_id: None, + limit: Some(500), + window: Some(60), + reason_code: Some(ReasonCode::new("get_lost")), + }]; + + let scoping = ItemScoping { + category: DataCategory::Error, + scoping: &Scoping { + organization_id: 42, + project_id: ProjectId::new(43), + public_key: "a94ae32be2584e0bbd7a4cbb95971fee".to_owned(), + key_id: Some(44), + }, + }; + + for i in 0..10 { + let rate_limits: Vec = RATE_LIMITER + .is_rate_limited(quotas, scoping, 100) + .expect("rate limiting failed") + .into_iter() + .collect(); + + if i >= 5 { + assert_eq!( + rate_limits, + vec![RateLimit { + categories: DataCategories::new(), + scope: RateLimitScope::Organization(42), + reason_code: Some(ReasonCode::new("get_lost")), + retry_after: rate_limits[0].retry_after, + }] + ); + } else { + assert_eq!(rate_limits, vec![]); + } + } + } + #[test] fn test_get_redis_key_scoped() { let quota = Quota { @@ -512,14 +559,16 @@ mod tests { let mut invocation = script.prepare_invoke(); invocation - .key(&foo) - .key(&r_foo) - .key(&bar) - .key(&r_bar) - .arg(1) - .arg(now + 60) - .arg(2) - .arg(now + 120); + .key(&foo) // key + .key(&r_foo) // refund key + .key(&bar) // key + .key(&r_bar) // refund key + .arg(1) // limit + .arg(now + 60) // expiry + .arg(1) // quantity + .arg(2) // limit + .arg(now + 120) // expiry + .arg(1); // quantity // The item should not be rate limited by either key. assert_eq!( @@ -560,7 +609,12 @@ mod tests { let () = conn.set(&apple, 5).unwrap(); let mut invocation = script.prepare_invoke(); - invocation.key(&orange).key(&baz).arg(1).arg(now + 60); + invocation + .key(&orange) // key + .key(&baz) // refund key + .arg(1) // limit + .arg(now + 60) // expiry + .arg(1); // quantity // increment assert_eq!( @@ -575,7 +629,12 @@ mod tests { ); let mut invocation = script.prepare_invoke(); - invocation.key(&orange).key(&apple).arg(1).arg(now + 60); + invocation + .key(&orange) // key + .key(&apple) // refund key + .arg(1) // limit + .arg(now + 60) // expiry + .arg(1); // quantity // test that refund key is used assert_eq!( diff --git a/relay-server/src/actors/events.rs b/relay-server/src/actors/events.rs index 53946d0117d..bce5b3779a0 100644 --- a/relay-server/src/actors/events.rs +++ b/relay-server/src/actors/events.rs @@ -912,8 +912,8 @@ impl EventProcessor { // When invoking the rate limiter, capture if the event item has been rate limited to also // remove it from the processing state eventually. - let mut envelope_limiter = EnvelopeLimiter::new(|item_scope, _quantity| { - let limits = rate_limiter.is_rate_limited(quotas, item_scope)?; + let mut envelope_limiter = EnvelopeLimiter::new(|item_scope, quantity| { + let limits = rate_limiter.is_rate_limited(quotas, item_scope, quantity)?; remove_event ^= Some(item_scope.category) == category && limits.is_limited(); Ok(limits) }); diff --git a/tests/integration/test_attachments.py b/tests/integration/test_attachments.py index 3fd106da8a3..54066015f96 100644 --- a/tests/integration/test_attachments.py +++ b/tests/integration/test_attachments.py @@ -171,6 +171,7 @@ def test_attachments_quotas( mini_sentry, relay_with_processing, attachments_consumer, outcomes_consumer, ): event_id = "515539018c9b4260a6f999572f1661ee" + attachment_body = b"blabla" relay = relay_with_processing() relay.wait_relay_healthcheck() @@ -181,17 +182,19 @@ def test_attachments_quotas( "id": "test_rate_limiting_{}".format(uuid.uuid4().hex), "categories": ["attachment"], "window": 3600, - "limit": 5, # TODO: Test attachment size quota once implemented + "limit": 5 * len(attachment_body), "reasonCode": "attachments_exceeded", } ] attachments_consumer = attachments_consumer() outcomes_consumer = outcomes_consumer() - attachments = [("att_1", "foo.txt", b"blabla")] + attachments = [("att_1", "foo.txt", attachment_body)] for i in range(5): - relay.send_attachments(42, event_id, [("att_1", "%s.txt" % i, b"")]) + relay.send_attachments(42, event_id, [("att_1", "%s.txt" % i, attachment_body)]) + chunk, _ = attachments_consumer.get_attachment_chunk() + assert chunk == attachment_body attachment = attachments_consumer.get_individual_attachment() assert attachment["attachment"]["name"] == "%s.txt" % i From 148c58e250c042feb1c0c3db00efa296863a2f50 Mon Sep 17 00:00:00 2001 From: Jan Michael Auer Date: Fri, 26 Jun 2020 12:30:24 +0200 Subject: [PATCH 2/2] meta: Changelog entry --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a999a6e5fff..8c86bc9594d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ **Internal**: - Restructure the envelope and event ingestion paths into a pipeline and apply rate limits to all envelopes. ([#635](https://github.com/getsentry/relay/pull/635), [#636](https://github.com/getsentry/relay/pull/636)) +- Pass the combined size of all attachments in an envelope to the Redis rate limiter as quantity to enforce attachment quotas. ([#639](https://github.com/getsentry/relay/pull/639)) ## 20.6.0