From bd307d317087e55f6af4e44d9837d07a02b347a5 Mon Sep 17 00:00:00 2001 From: Jake McGinty Date: Fri, 4 May 2018 01:59:47 -0700 Subject: ratelimiter: implement ratelimiter and tests --- Cargo.lock | 69 +++++++++++---------- src/lib.rs | 1 + src/ratelimiter.rs | 178 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 214 insertions(+), 34 deletions(-) create mode 100644 src/ratelimiter.rs diff --git a/Cargo.lock b/Cargo.lock index d1ee78d..ee5b5f7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -39,13 +39,13 @@ dependencies = [ [[package]] name = "backtrace" -version = "0.3.6" +version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ "backtrace-sys 0.1.16 (registry+https://github.com/rust-lang/crates.io-index)", - "cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", "libc 0.2.40 (registry+https://github.com/rust-lang/crates.io-index)", - "rustc-demangle 0.1.7 (registry+https://github.com/rust-lang/crates.io-index)", + "rustc-demangle 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)", "winapi 0.3.4 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -126,7 +126,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" [[package]] name = "cfg-if" -version = "0.1.2" +version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" [[package]] @@ -197,8 +197,8 @@ dependencies = [ "itertools 0.7.8 (registry+https://github.com/rust-lang/crates.io-index)", "itertools-num 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)", "log 0.4.1 (registry+https://github.com/rust-lang/crates.io-index)", - "serde 1.0.44 (registry+https://github.com/rust-lang/crates.io-index)", - "serde_derive 1.0.44 (registry+https://github.com/rust-lang/crates.io-index)", + "serde 1.0.45 (registry+https://github.com/rust-lang/crates.io-index)", + "serde_derive 1.0.45 (registry+https://github.com/rust-lang/crates.io-index)", "serde_json 1.0.17 (registry+https://github.com/rust-lang/crates.io-index)", "simplelog 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -249,7 +249,7 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ "arrayvec 0.4.7 (registry+https://github.com/rust-lang/crates.io-index)", - "cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", "crossbeam-utils 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)", "lazy_static 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)", "memoffset 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)", @@ -263,7 +263,7 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ "arrayvec 0.4.7 (registry+https://github.com/rust-lang/crates.io-index)", - "cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", "crossbeam-utils 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)", "lazy_static 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)", "memoffset 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)", @@ -275,7 +275,7 @@ name = "crossbeam-utils" version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ - "cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", ] [[package]] @@ -283,7 +283,7 @@ name = "crossbeam-utils" version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ - "cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", ] [[package]] @@ -339,7 +339,7 @@ name = "failure" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ - "backtrace 0.3.6 (registry+https://github.com/rust-lang/crates.io-index)", + "backtrace 0.3.7 (registry+https://github.com/rust-lang/crates.io-index)", "failure_derive 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -367,7 +367,7 @@ name = "filetime" version = "0.1.15" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ - "cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", "libc 0.2.40 (registry+https://github.com/rust-lang/crates.io-index)", "redox_syscall 0.1.37 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -438,7 +438,7 @@ dependencies = [ "pest_derive 1.0.7 (registry+https://github.com/rust-lang/crates.io-index)", "quick-error 1.2.1 (registry+https://github.com/rust-lang/crates.io-index)", "regex 0.2.11 (registry+https://github.com/rust-lang/crates.io-index)", - "serde 1.0.44 (registry+https://github.com/rust-lang/crates.io-index)", + "serde 1.0.45 (registry+https://github.com/rust-lang/crates.io-index)", "serde_json 1.0.17 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -532,7 +532,7 @@ name = "log" version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ - "cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", ] [[package]] @@ -584,9 +584,10 @@ dependencies = [ [[package]] name = "mio-uds" -version = "0.6.4" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ + "iovec 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", "libc 0.2.40 (registry+https://github.com/rust-lang/crates.io-index)", "mio 0.6.14 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -628,7 +629,7 @@ name = "net2" version = "0.2.32" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ - "cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", "libc 0.2.40 (registry+https://github.com/rust-lang/crates.io-index)", "winapi 0.3.4 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -649,7 +650,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ "bitflags 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)", "bytes 0.4.7 (registry+https://github.com/rust-lang/crates.io-index)", - "cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", "gcc 0.3.54 (registry+https://github.com/rust-lang/crates.io-index)", "libc 0.2.40 (registry+https://github.com/rust-lang/crates.io-index)", "void 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)", @@ -662,7 +663,7 @@ source = "git+https://github.com/mcginty/nix?branch=ipv6-pktinfo#d5c3de15edc8147 dependencies = [ "bitflags 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)", "cc 1.0.15 (registry+https://github.com/rust-lang/crates.io-index)", - "cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", "libc 0.2.40 (git+https://github.com/rust-lang/libc)", "void 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -895,7 +896,7 @@ dependencies = [ [[package]] name = "rustc-demangle" -version = "0.1.7" +version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" [[package]] @@ -921,7 +922,7 @@ dependencies = [ [[package]] name = "scoped-tls" -version = "0.1.1" +version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" [[package]] @@ -944,12 +945,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index" [[package]] name = "serde" -version = "1.0.44" +version = "1.0.45" source = "registry+https://github.com/rust-lang/crates.io-index" [[package]] name = "serde_derive" -version = "1.0.44" +version = "1.0.45" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ "proc-macro2 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)", @@ -974,7 +975,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ "dtoa 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)", "itoa 0.4.1 (registry+https://github.com/rust-lang/crates.io-index)", - "serde 1.0.44 (registry+https://github.com/rust-lang/crates.io-index)", + "serde 1.0.45 (registry+https://github.com/rust-lang/crates.io-index)", ] [[package]] @@ -1021,7 +1022,7 @@ name = "socket2" version = "0.3.5" source = "registry+https://github.com/rust-lang/crates.io-index" dependencies = [ - "cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)", "libc 0.2.40 (registry+https://github.com/rust-lang/crates.io-index)", "winapi 0.3.4 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -1232,7 +1233,7 @@ dependencies = [ "iovec 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", "log 0.4.1 (registry+https://github.com/rust-lang/crates.io-index)", "mio 0.6.14 (registry+https://github.com/rust-lang/crates.io-index)", - "scoped-tls 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)", + "scoped-tls 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", "tokio 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)", "tokio-executor 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", "tokio-io 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)", @@ -1289,7 +1290,7 @@ dependencies = [ "futures 0.1.21 (registry+https://github.com/rust-lang/crates.io-index)", "libc 0.2.40 (registry+https://github.com/rust-lang/crates.io-index)", "mio 0.6.14 (registry+https://github.com/rust-lang/crates.io-index)", - "mio-uds 0.6.4 (registry+https://github.com/rust-lang/crates.io-index)", + "mio-uds 0.6.5 (registry+https://github.com/rust-lang/crates.io-index)", "tokio-core 0.1.17 (registry+https://github.com/rust-lang/crates.io-index)", "tokio-io 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)", "winapi 0.3.4 (registry+https://github.com/rust-lang/crates.io-index)", @@ -1354,7 +1355,7 @@ dependencies = [ "libc 0.2.40 (registry+https://github.com/rust-lang/crates.io-index)", "log 0.3.9 (registry+https://github.com/rust-lang/crates.io-index)", "mio 0.6.14 (registry+https://github.com/rust-lang/crates.io-index)", - "mio-uds 0.6.4 (registry+https://github.com/rust-lang/crates.io-index)", + "mio-uds 0.6.5 (registry+https://github.com/rust-lang/crates.io-index)", "tokio-core 0.1.17 (registry+https://github.com/rust-lang/crates.io-index)", "tokio-io 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)", ] @@ -1537,7 +1538,7 @@ dependencies = [ "checksum arrayref 0.3.4 (registry+https://github.com/rust-lang/crates.io-index)" = "0fd1479b7c29641adbd35ff3b5c293922d696a92f25c8c975da3e0acbc87258f" "checksum arrayvec 0.4.7 (registry+https://github.com/rust-lang/crates.io-index)" = "a1e964f9e24d588183fcb43503abda40d288c8657dfc27311516ce2f05675aef" "checksum atty 0.2.10 (registry+https://github.com/rust-lang/crates.io-index)" = "2fc4a1aa4c24c0718a250f0681885c1af91419d242f29eb8f2ab28502d80dbd1" -"checksum backtrace 0.3.6 (registry+https://github.com/rust-lang/crates.io-index)" = "ebbe525f66f42d207968308ee86bc2dd60aa5fab535b22e616323a173d097d8e" +"checksum backtrace 0.3.7 (registry+https://github.com/rust-lang/crates.io-index)" = "8ea58cd16fd6c9d120b5bcb01d63883ae4cc7ba2aed35c1841b862a3c7ef6639" "checksum backtrace-sys 0.1.16 (registry+https://github.com/rust-lang/crates.io-index)" = "44585761d6161b0f57afc49482ab6bd067e4edef48c12a152c237eb0203f7661" "checksum base64 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)" = "30e93c03064e7590d0466209155251b90c22e37fab1daf2771582598b5827557" "checksum bitflags 0.4.0 (registry+https://github.com/rust-lang/crates.io-index)" = "8dead7461c1127cf637931a1e50934eb6eee8bff2f74433ac7909e9afcee04a3" @@ -1550,7 +1551,7 @@ dependencies = [ "checksum bytes 0.4.7 (registry+https://github.com/rust-lang/crates.io-index)" = "2f1d50c876fb7545f5f289cd8b2aee3f359d073ae819eed5d6373638e2c61e59" "checksum cast 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "926013f2860c46252efceabb19f4a6b308197505082c609025aa6706c011d427" "checksum cc 1.0.15 (registry+https://github.com/rust-lang/crates.io-index)" = "0ebb87d1116151416c0cf66a0e3fb6430cccd120fd6300794b4dfaa050ac40ba" -"checksum cfg-if 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "d4c819a1287eb618df47cc647173c5c4c66ba19d888a6e50d605672aed3140de" +"checksum cfg-if 0.1.3 (registry+https://github.com/rust-lang/crates.io-index)" = "405216fd8fe65f718daa7102ea808a946b6ce40c742998fbfd3463645552de18" "checksum chacha20-poly1305-aead 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "77d2058ba29594f69c75e8a9018e0485e3914ca5084e3613cd64529042f5423b" "checksum chrono 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)" = "1cce36c92cb605414e9b824f866f5babe0a0368e39ea07393b9b63cf3844c0e6" "checksum clap 2.31.2 (registry+https://github.com/rust-lang/crates.io-index)" = "f0f16b89cbb9ee36d87483dc939fe9f1e13c05898d56d7b230a0d4dff033a536" @@ -1603,7 +1604,7 @@ dependencies = [ "checksum memoffset 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "0f9dc261e2b62d7a622bf416ea3c5245cdd5d9a7fcc428c0d06804dfce1775b3" "checksum mio 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)" = "a637d1ca14eacae06296a008fa7ad955347e34efcb5891cfd8ba05491a37907e" "checksum mio 0.6.14 (registry+https://github.com/rust-lang/crates.io-index)" = "6d771e3ef92d58a8da8df7d6976bfca9371ed1de6619d9d5a5ce5b1f29b85bfe" -"checksum mio-uds 0.6.4 (registry+https://github.com/rust-lang/crates.io-index)" = "1731a873077147b626d89cc6c2a0db6288d607496c5d10c0cfcf3adc697ec673" +"checksum mio-uds 0.6.5 (registry+https://github.com/rust-lang/crates.io-index)" = "914c9238b09a99d89afe4a98625a160d3e946e5dbda65c6486695deb119f005e" "checksum mio-utun 0.6.15 (registry+https://github.com/rust-lang/crates.io-index)" = "480124e6fb01b3275a2f20234616a062219c615b7856646e21f167aa46201697" "checksum miow 0.1.5 (registry+https://github.com/rust-lang/crates.io-index)" = "3e690c5df6b2f60acd45d56378981e827ff8295562fc8d34f573deb267a59cd1" "checksum miow 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "8c1f2f3b1cf331de6896aabf6e9d55dca90356cc9960cca7eaaf408a355ae919" @@ -1637,16 +1638,16 @@ dependencies = [ "checksum regex-syntax 0.5.6 (registry+https://github.com/rust-lang/crates.io-index)" = "7d707a4fa2637f2dca2ef9fd02225ec7661fe01a53623c1e6515b6916511f7a7" "checksum ring 0.13.0-alpha (registry+https://github.com/rust-lang/crates.io-index)" = "946e5e2b336032275e23152755ec2a0610ede0302101d3666fdb686795d26535" "checksum rust-crypto 0.2.36 (git+https://github.com/mcginty/rust-crypto)" = "" -"checksum rustc-demangle 0.1.7 (registry+https://github.com/rust-lang/crates.io-index)" = "11fb43a206a04116ffd7cfcf9bcb941f8eb6cc7ff667272246b0a1c74259a3cb" +"checksum rustc-demangle 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)" = "76d7ba1feafada44f2d38eed812bd2489a03c0f5abb975799251518b68848649" "checksum rustc-serialize 0.3.24 (registry+https://github.com/rust-lang/crates.io-index)" = "dcf128d1287d2ea9d80910b5f1120d0b8eede3fbf1abe91c40d39ea7d51e6fda" "checksum rustc_version 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "a54aa04a10c68c1c4eacb4337fd883b435997ede17a9385784b990777686b09a" "checksum same-file 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)" = "cfb6eded0b06a0b512c8ddbcf04089138c9b4362c2f696f3c3d76039d68f3637" -"checksum scoped-tls 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "8674d439c964889e2476f474a3bf198cc9e199e77499960893bac5de7e9218a4" +"checksum scoped-tls 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)" = "332ffa32bf586782a3efaeb58f127980944bbc8c4d6913a86107ac2a5ab24b28" "checksum scopeguard 0.3.3 (registry+https://github.com/rust-lang/crates.io-index)" = "94258f53601af11e6a49f722422f6e3425c52b06245a5cf9bc09908b174f5e27" "checksum semver 0.9.0 (registry+https://github.com/rust-lang/crates.io-index)" = "1d7eb9ef2c18661902cc47e535f9bc51b78acd254da71d375c2f6720d9a40403" "checksum semver-parser 0.7.0 (registry+https://github.com/rust-lang/crates.io-index)" = "388a1df253eca08550bef6c72392cfe7c30914bf41df5269b68cbd6ff8f570a3" -"checksum serde 1.0.44 (registry+https://github.com/rust-lang/crates.io-index)" = "ce67a48047802238bfc88687272de48fd6d7af256b0097f110e968b0017235a5" -"checksum serde_derive 1.0.44 (registry+https://github.com/rust-lang/crates.io-index)" = "5d870179775231857959d909fd4ab5e8c4e69d537a1e9698707b2ba2c92d370e" +"checksum serde 1.0.45 (registry+https://github.com/rust-lang/crates.io-index)" = "6a49d806123bcdaacdefe7aab3721c64ec11d05921bf64d888a857d3a92024a0" +"checksum serde_derive 1.0.45 (registry+https://github.com/rust-lang/crates.io-index)" = "0409f5130e9b06444e07d4c71f55d6a2c4d1290d79faa612d9b0b540a9703fcd" "checksum serde_derive_internals 0.23.1 (registry+https://github.com/rust-lang/crates.io-index)" = "9d30c4596450fd7bbda79ef15559683f9a79ac0193ea819db90000d7e1cae794" "checksum serde_json 1.0.17 (registry+https://github.com/rust-lang/crates.io-index)" = "f3ad6d546e765177cf3dded3c2e424a8040f870083a0e64064746b958ece9cb1" "checksum simplelog 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)" = "ce595117de34b75e057b41e99079e43e9fcc4e5ec9c7ba5f2fea55321f0c624e" diff --git a/src/lib.rs b/src/lib.rs index 5c39c68..7b038b7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -50,6 +50,7 @@ mod cookie; mod error; mod ip_packet; mod message; +mod ratelimiter; mod router; mod timer; mod udp; diff --git a/src/ratelimiter.rs b/src/ratelimiter.rs new file mode 100644 index 0000000..26faaa4 --- /dev/null +++ b/src/ratelimiter.rs @@ -0,0 +1,178 @@ +#![allow(dead_code)] + +use timestamp::Timestamp; + +use failure::Error; +use futures::{unsync::mpsc, Async, Future, Poll, Stream, Sink}; +use tokio::timer::Interval; +use tokio_core::reactor::Handle; +use std::collections::HashMap; +use std::net::IpAddr; +use std::time::{Duration, Instant}; + +const PACKETS_PER_SECOND : u64 = 20; +const PACKETS_BURSTABLE : u64 = 5; +const PACKET_COST : u64 = 1_000_000_000 / PACKETS_PER_SECOND; +const MAX_TOKENS : u64 = PACKET_COST * PACKETS_BURSTABLE; + +lazy_static! { + pub static ref GC_INTERVAL: Duration = Duration::new(1, 0); +} + +struct Entry { + pub last_time : Timestamp, + pub tokens : u64, +} + +struct RateLimiter { + table: HashMap, + rx: mpsc::Receiver<()>, +} + +impl RateLimiter { + pub fn new(handle: &Handle) -> Result { + let (tx, rx) = mpsc::channel(128); + let i_handle = handle.clone(); + + let gc = Interval::new(Instant::now(), *GC_INTERVAL) + .map_err(|e| panic!("timer failed; err={:?}", e)) + .for_each(move |_| { + i_handle.spawn(tx.clone().send(()).then(|_| Ok(()))); + Ok(()) + }); + handle.spawn(gc); + + Ok(Self { + table: HashMap::new(), + rx + }) + } + + fn _new_for_test() -> Self { + let (_tx, rx) = mpsc::channel(1); + Self { table: HashMap::new(), rx } + } + + pub fn allow(&mut self, addr: &IpAddr) -> bool { + if let Some(entry) = self.table.get_mut(addr) { + entry.tokens = MAX_TOKENS.min(entry.tokens + entry.last_time.elapsed().subsec_nanos() as u64); + entry.last_time = Timestamp::now(); + + if entry.tokens > PACKET_COST { + entry.tokens -= PACKET_COST; + return true; + } else { + return false; + } + } + + self.table.insert(*addr, Entry { + last_time: Timestamp::now(), + tokens: MAX_TOKENS - PACKET_COST + }); + true + } + + fn handle_gc(&mut self) { + self.table.retain(|_, ref mut entry| entry.last_time.elapsed() <= *GC_INTERVAL); + } +} + +impl Future for RateLimiter { + type Item = (); + type Error = (); + + fn poll(&mut self) -> Poll { + match self.rx.poll() { + Ok(Async::Ready(Some(()))) => self.handle_gc(), + _ => {}, + } + Ok(Async::NotReady) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std; + + struct Result { + allowed: bool, + text: &'static str, + wait: Duration, + } + + #[test] + fn test_ratelimiter() { + let mut ratelimiter = RateLimiter::_new_for_test(); + let mut expected = vec![]; + let ips = vec![ + "127.0.0.1".parse().unwrap(), + "192.168.1.1".parse().unwrap(), + "172.167.2.3".parse().unwrap(), + "97.231.252.215".parse().unwrap(), + "248.97.91.167".parse().unwrap(), + "188.208.233.47".parse().unwrap(), + "104.2.183.179".parse().unwrap(), + "72.129.46.120".parse().unwrap(), + "2001:0db8:0a0b:12f0:0000:0000:0000:0001".parse().unwrap(), + "f5c2:818f:c052:655a:9860:b136:6894:25f0".parse().unwrap(), + "b2d7:15ab:48a7:b07c:a541:f144:a9fe:54fc".parse().unwrap(), + "a47b:786e:1671:a22b:d6f9:4ab0:abc7:c918".parse().unwrap(), + "ea1e:d155:7f7a:98fb:2bf5:9483:80f6:5445".parse().unwrap(), + "3f0e:54a2:f5b4:cd19:a21d:58e1:3746:84c4".parse().unwrap(),]; + + for _ in 0..PACKETS_BURSTABLE { + expected.push(Result { + allowed : true, + wait : Duration::new(0, 0), + text : "inital burst", + }); + } + + expected.push(Result { + allowed : false, + wait : Duration::new(0, 0), + text : "after burst", + }); + + expected.push(Result { + allowed : true, + wait : Duration::new(0, PACKET_COST as u32), + text : "filling tokens for single packet", + }); + + expected.push(Result { + allowed : false, + wait : Duration::new(0, 0), + text : "not having refilled enough", + }); + + expected.push(Result { + allowed : true, + wait : Duration::new(0, 2 * PACKET_COST as u32), + text : "filling tokens for 2 * packet burst", + }); + + expected.push(Result { + allowed : true, + wait : Duration::new(0, 0), + text : "second packet in 2 packet burst", + }); + + expected.push(Result { + allowed : false, + wait : Duration::new(0, 0), + text : "packet following 2 packet burst", + }); + + for item in expected { + std::thread::sleep(item.wait); + for ip in ips.iter() { + if ratelimiter.allow(&ip) != item.allowed { + panic!("test failed for {} on {}. expected: {}, got: {}", ip, item.text, item.allowed, !item.allowed) + } + } + } + } +} \ No newline at end of file -- cgit v1.2.3-59-g8ed1b