rate-limiting tests for server-policy and app-inbound crates

linkerd/app/inbound/src/policy/http/tests.rs

@@ -1,6 +1,7 @@
 use super::*;
 use crate::policy::{Authentication, Authorization, Meta, Protocol, ServerPolicy};
 use linkerd_app_core::{svc::Service, Infallible};
+use linkerd_proxy_server_policy::local_rate_limit::LocalRateLimit;
 
 macro_rules! conn {
     ($client:expr, $dst:expr) => {{
@@ -19,7 +20,7 @@ macro_rules! conn {
 }
 
 macro_rules! new_svc {
-    ($proto:expr, $conn:expr, $rsp:expr) => {{
+    ($proto:expr, $conn:expr, $rsp:expr, $rl: expr) => {{
         let (policy, tx) = AllowPolicy::for_test(
             $conn.dst,
             ServerPolicy {
@@ -29,7 +30,7 @@ macro_rules! new_svc {
                     kind: "Server".into(),
                     name: "testsrv".into(),
                 }),
-                local_rate_limit: Arc::new(Default::default()),
+                local_rate_limit: Arc::new($rl),
             },
         );
         let svc = HttpPolicyService {
@@ -47,7 +48,11 @@ macro_rules! new_svc {
         (svc, tx)
     }};
 
-    ($proto:expr) => {{
+    ($proto:expr, $conn:expr, $rsp:expr) => {{
+        new_svc!($proto, $conn, $rsp, Default::default())
+    }};
+
+    ($proto:expr, $rl:expr) => {{
         new_svc!(
             $proto,
             conn!(),
@@ -57,9 +62,14 @@ macro_rules! new_svc {
                     .unwrap();
                 rsp.extensions_mut().insert(permit.clone());
                 Ok::<_, Infallible>(rsp)
-            }
+            },
+            $rl
         )
     }};
+
+    ($proto:expr) => {{
+        new_svc!($proto, Default::default())
+    }};
 }
 
 #[tokio::test(flavor = "current_thread")]
@@ -365,6 +375,38 @@ async fn http_filter_inject_failure() {
     );
 }
 
+#[tokio::test(flavor = "current_thread")]
+async fn rate_limit() {
+    use linkerd_app_core::{Ipv4Net, Ipv6Net};
+
+    let rmeta = Meta::new_default("default");
+    let rl = LocalRateLimit::new_no_overrides(Some(10), Some(5));
+    let authorizations = Arc::new([Authorization {
+        meta: rmeta.clone(),
+        networks: vec![Ipv4Net::default().into(), Ipv6Net::default().into()],
+        authentication: Authentication::Unauthenticated,
+    }]);
+
+    let (mut svc, _tx) = new_svc!(
+        Protocol::Http1(Arc::new([http::default(authorizations.clone())])),
+        rl
+    );
+
+    let rsp = svc
+        .call(
+            ::http::Request::builder()
+                .body(hyper::Body::default())
+                .unwrap(),
+        )
+        .await
+        .expect("serves");
+    let permit = rsp
+        .extensions()
+        .get::<HttpRoutePermit>()
+        .expect("permitted");
+    assert_eq!(permit.labels.route.route, rmeta);
+}
+
 #[tokio::test(flavor = "current_thread")]
 async fn grpc_route() {
     use linkerd_proxy_server_policy::grpc::{

linkerd/proxy/server-policy/src/local_rate_limit/tests.rs

@@ -0,0 +1,186 @@
+use std::time::Duration;
+
+use super::*;
+use maplit::hashmap;
+
+#[cfg(feature = "proto")]
+#[tokio::test(flavor = "current_thread")]
+async fn from_proto() {
+    use linkerd2_proxy_api::{
+        inbound::{self, http_local_rate_limit},
+        meta,
+    };
+
+    let client_1: Id = "client-1".parse().unwrap();
+    let client_2: Id = "client-2".parse().unwrap();
+    let client_3: Id = "client-3".parse().unwrap();
+    let client_4: Id = "client-4".parse().unwrap();
+    let rl_proto = inbound::HttpLocalRateLimit {
+        metadata: Some(meta::Metadata {
+            kind: Some(meta::metadata::Kind::Default("ratelimit-1".into())),
+        }),
+        total: Some(http_local_rate_limit::Limit {
+            requests_per_second: 100,
+        }),
+        identity: Some(http_local_rate_limit::Limit {
+            requests_per_second: 20,
+        }),
+        overrides: vec![
+            http_local_rate_limit::Override {
+                limit: Some(http_local_rate_limit::Limit {
+                    requests_per_second: 50,
+                }),
+                clients: Some(http_local_rate_limit::r#override::ClientIdentities {
+                    identities: vec![
+                        inbound::Identity {
+                            name: client_1.to_string(),
+                        },
+                        inbound::Identity {
+                            name: client_2.to_string(),
+                        },
+                    ],
+                }),
+            },
+            http_local_rate_limit::Override {
+                limit: Some(http_local_rate_limit::Limit {
+                    requests_per_second: 75,
+                }),
+                clients: Some(http_local_rate_limit::r#override::ClientIdentities {
+                    identities: vec![
+                        inbound::Identity {
+                            name: client_3.to_string(),
+                        },
+                        inbound::Identity {
+                            name: client_4.to_string(),
+                        },
+                    ],
+                }),
+            },
+        ],
+    };
+
+    let rl = Into::<LocalRateLimit>::into(rl_proto);
+    assert_eq!(rl.total.as_ref().unwrap().rps.get(), 100);
+    assert_eq!(rl.per_identity.as_ref().unwrap().rps.get(), 20);
+
+    assert_eq!(rl.overrides.get(&client_1).unwrap().rps.get(), 50);
+    assert_eq!(rl.overrides.get(&client_2).unwrap().rps.get(), 50);
+    assert_eq!(rl.overrides.get(&client_3).unwrap().rps.get(), 75);
+    assert_eq!(rl.overrides.get(&client_4).unwrap().rps.get(), 75);
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn check_rate_limits() {
+    let total = RateLimit::<Direct, FakeRelativeClock>::new(35).unwrap();
+    let per_identity = RateLimit::<Keyed, FakeRelativeClock>::new(5).unwrap();
+    let overrides = hashmap! {
+        "client-3".parse().unwrap() => Arc::new(RateLimit::<Direct, FakeRelativeClock>::new(10).unwrap()),
+        "client-4".parse().unwrap() => Arc::new(RateLimit::<Direct, FakeRelativeClock>::new(15).unwrap()),
+    };
+    let rl = LocalRateLimit {
+        total: Some(total),
+        per_identity: Some(per_identity),
+        overrides,
+    };
+
+    // These clients will be rate-limited by the per_identity rate-limiter
+    let client_1: Id = "client-1".parse().unwrap();
+    let client_2: Id = "client-2".parse().unwrap();
+
+    // These clients will be rate-limited by the overrides rate-limiters
+    let client_3: Id = "client-3".parse().unwrap();
+    let client_4: Id = "client-4".parse().unwrap();
+
+    let total_clock = rl.total.as_ref().unwrap().limiter.clock();
+    let per_identity_clock = rl.per_identity.as_ref().unwrap().limiter.clock();
+    let client_3_clock = rl.overrides.get(&client_3).unwrap().limiter.clock();
+    let client_4_clock = rl.overrides.get(&client_4).unwrap().limiter.clock();
+
+    // This loop checks that:
+    // - client_1 gets rate-limited first via the per_identity rate-limiter
+    // - then client_3 and client_4 via the overrides rate-limiters
+    // - then client_2 via the total rate-limiter
+    // - then we advance time to replenish the rate-limiters buckets and repeat the checks
+    for i in 1..=5 {
+        // Requests per-client: 5
+        // Total requests: 15
+        // All clients should NOT be rate-limited
+        for _ in 1..=5 {
+            assert!(rl.check(Some(&client_1)).is_ok());
+            assert!(rl.check(Some(&client_3)).is_ok());
+            assert!(rl.check(Some(&client_4)).is_ok());
+        }
+        // Reached per_identity limit for client_1
+        // Total requests: 16
+        assert_eq!(
+            rl.check(Some(&client_1)),
+            Err(RateLimitError::PerIdentity(NonZeroU32::new(5).unwrap()))
+        );
+
+        // Requests per-client: 10
+        // Total requests thus far: 26
+        // client_3 and client_4 should NOT be rate-limited
+        for _ in 1..=5 {
+            assert!(rl.check(Some(&client_3)).is_ok());
+            assert!(rl.check(Some(&client_4)).is_ok());
+        }
+        // Total requests thus far: 27
+        // Reached override limit for client_3
+        assert_eq!(
+            rl.check(Some(&client_3)),
+            Err(RateLimitError::Override(NonZeroU32::new(10).unwrap()))
+        );
+
+        // Requests per-client: 5
+        // Total requests thus far: 32
+        // client_4 should NOT be rate-limited
+        for _ in 1..=5 {
+            assert!(rl.check(Some(&client_4)).is_ok());
+        }
+        // Total requests: 33
+        // Reached override limit for client_4
+        assert_eq!(
+            rl.check(Some(&client_4)),
+            Err(RateLimitError::Override(NonZeroU32::new(15).unwrap()))
+        );
+
+        if i == 1 {
+            // Total requests: 35
+            // Only 2 requests for client_2 allowed as we're reaching the total rate-limit
+            // See note below about why this is only run for the first iteration
+            for _ in 1..=2 {
+                assert!(rl.check(Some(&client_2)).is_ok());
+            }
+        }
+
+        // Total requests: 36
+        // Reached total limit for all clients
+        assert_eq!(
+            rl.check(Some(&client_2)),
+            Err(RateLimitError::Total(NonZeroU32::new(35).unwrap()))
+        );
+
+        // Advance time for a couple of seconds to replenish the rate-limiters buckets
+        total_clock.advance(Duration::from_secs(2));
+        per_identity_clock.advance(Duration::from_secs(2));
+        client_3_clock.advance(Duration::from_secs(2));
+        client_4_clock.advance(Duration::from_secs(2));
+
+        // Bug?
+        // Each rate-limiter gets a free request after the initial iteration, so we consume those
+        // free requests here.
+        // See https://github.com/boinkor-net/governor/issues/249
+        //
+        // Free request for the per_client rate-limiter, and free request for the global
+        // rate-limiter
+        assert!(rl.check(Some(&client_1)).is_ok());
+        // Free request for the first override rate-limiter, but accumulates request to the global
+        // rate-limiter
+        assert!(rl.check(Some(&client_3)).is_ok());
+        // Free request for the second override rate-limiter, but accumulates request to the global
+        // rate-limiter
+        assert!(rl.check(Some(&client_4)).is_ok());
+        // We accumulated 2 non-free requests for the global rate-limiter that will impact the next
+        // iteration; that's why we don't make those 2 checks above for i > 1
+    }
+}