├── .devcontainer
    ├── Dockerfile
    └── devcontainer.json
├── .github
    ├── FUNDING.yml
    └── workflows
    │   ├── publish.yml
    │   └── rust.yml
├── .gitignore
├── Cargo.toml
├── README.md
├── git-hooks
    ├── pre-commit
    └── pre-push
├── rocket-governor.code-workspace
├── src
    ├── header.rs
    ├── lib.rs
    ├── limit_error.rs
    ├── limit_error
    │   └── catcher.rs
    ├── limit_header_gen.rs
    ├── logger.rs
    ├── registry.rs
    ├── req_state.rs
    └── rocket_governable.rs
└── tests
    └── tests.rs


/.devcontainer/Dockerfile:
--------------------------------------------------------------------------------
 1 | # See here for image contents: https://github.com/microsoft/vscode-dev-containers/tree/v0.187.0/containers/rust/.devcontainer/base.Dockerfile
 2 | 
 3 | FROM mcr.microsoft.com/vscode/devcontainers/rust:0-1
 4 | 
 5 | # [Optional] Uncomment this section to install additional packages.
 6 | # RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \
 7 | #     && apt-get -y install --no-install-recommends <your-package-list-here>
 8 | 
 9 | RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \
10 |     && apt upgrade \
11 |     && rustup update && \
12 |     && cargo install sccache && \
13 |     && echo -e "[build]\nrustc-wrapper = \"/usr/local/cargo/bin/sccache\"\n" > /usr/local/cargo/config.toml && \
14 |     && cargo install cargo-audit cargo-outdated cargo-install-update cargo-watch cargo-crev && \
15 |     && sccache --stop-server && \
16 |     && apt-get clean -y && rm -rf /var/lib/apt/lists/* /usr/local/cargo/registry ~/.cache/sccache
17 | 


--------------------------------------------------------------------------------
/.devcontainer/devcontainer.json:
--------------------------------------------------------------------------------
 1 | // For format details, see https://aka.ms/devcontainer.json. For config options, see the README at:
 2 | // https://github.com/microsoft/vscode-dev-containers/tree/v0.187.0/containers/rust
 3 | {
 4 | 	"name": "Rust",
 5 | 	"build": {
 6 | 		"dockerfile": "Dockerfile"
 7 | 	},
 8 | 	"runArgs": [ "--cap-add=SYS_PTRACE", "--security-opt", "seccomp=unconfined" ],
 9 | 
10 | 	// Set *default* container specific settings.json values on container create.
11 | 	"settings": { 
12 | 		"lldb.executable": "/usr/bin/lldb",
13 | 		// VS Code don't watch files under ./target
14 | 		"files.watcherExclude": {
15 | 			"**/target/**": true
16 | 		}
17 | 	},
18 | 
19 | 	// Add the IDs of extensions you want installed when the container is created.
20 | 	"extensions": [
21 | 		"bungcip.better-toml",
22 | 		"vadimcn.vscode-lldb",
23 | 		"mutantdino.resourcemonitor",
24 | 		"serayuzgur.crates",
25 | 		"matklad.rust-analyzer",
26 | 	],
27 | 
28 | 	// Use 'forwardPorts' to make a list of ports inside the container available locally.
29 | 	"forwardPorts": [],
30 | 
31 | 	// Use 'postCreateCommand' to run commands after the container is created.
32 | 	// "postCreateCommand": "rustc --version",
33 | 
34 | 	// Comment out connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root.
35 | 	"remoteUser": "vscode"
36 | }
37 | 


--------------------------------------------------------------------------------
/.github/FUNDING.yml:
--------------------------------------------------------------------------------
 1 | # These are supported funding model platforms
 2 | 
 3 | github: [kolbma] # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]
 4 | patreon: # Replace with a single Patreon username
 5 | open_collective: # Replace with a single Open Collective username
 6 | ko_fi: # Replace with a single Ko-fi username
 7 | tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel
 8 | community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry
 9 | liberapay: # Replace with a single Liberapay username
10 | issuehunt: # Replace with a single IssueHunt username
11 | otechie: # Replace with a single Otechie username
12 | lfx_crowdfunding: # Replace with a single LFX Crowdfunding project-name e.g., cloud-foundry
13 | custom: ['https://www.buymeacoffee.com/arlecchino'] # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2']
14 | 


--------------------------------------------------------------------------------
/.github/workflows/publish.yml:
--------------------------------------------------------------------------------
 1 | name: Publish crates.io
 2 | 
 3 | on:
 4 |   push:
 5 |     # branches: [ main ]
 6 |     tags:
 7 |       - v[0-9]+.[0-9]+.[0-9]+
 8 |       - '!v[0-9]+.[0-9]+.[0-9]+-*'
 9 |   workflow_dispatch:
10 | 
11 | jobs:
12 |   publish:
13 |     runs-on: ubuntu-latest
14 |     steps:
15 |     - uses: actions/checkout@v4
16 |     - uses: dtolnay/rust-toolchain@master
17 |       with:
18 |           toolchain: stable
19 |     - uses: katyo/publish-crates@v2
20 |       with:
21 |           dry-run: false
22 |           registry-token: ${{ secrets.CRATES_IO_ROCKET_GOVERNOR }}
23 | 


--------------------------------------------------------------------------------
/.github/workflows/rust.yml:
--------------------------------------------------------------------------------
  1 | name: Rust
  2 | 
  3 | on:
  4 |   push:
  5 |     branches: [ main ]
  6 |   pull_request:
  7 |     branches: [ main ]
  8 | 
  9 | env:
 10 |   CARGO_TERM_COLOR: always
 11 |   RUST_MSRV_TOOLCHAIN: 1.69.0
 12 | 
 13 | jobs:
 14 |   source-check:
 15 |     runs-on: ubuntu-latest
 16 | 
 17 |     permissions:
 18 |       checks: write
 19 |       contents: read
 20 | 
 21 |     steps:
 22 |     - uses: dtolnay/rust-toolchain@master
 23 |       with:
 24 |         toolchain: stable
 25 |         components: clippy, rustfmt
 26 |     - uses: actions/checkout@v4
 27 |     - uses: loliGothicK/clippy-check@v0.3.2
 28 |       with:
 29 |         token: ${{ secrets.GITHUB_TOKEN }}
 30 |     - uses: loliGothicK/clippy-check@v0.3.2
 31 |       with:
 32 |         token: ${{ secrets.GITHUB_TOKEN }}
 33 |         options: --all-features
 34 |     - run: cargo fmt --check
 35 | 
 36 |   cargo-doc-stable:
 37 |     runs-on: ubuntu-latest
 38 |     needs: source-check
 39 | 
 40 |     steps:
 41 |     - uses: dtolnay/rust-toolchain@master
 42 |       with:
 43 |         toolchain: stable
 44 |     - uses: actions/checkout@v4
 45 |     - run: cargo doc --verbose --all-features --no-deps
 46 | 
 47 |   cargo-doc-msrv:
 48 |     runs-on: ubuntu-latest
 49 |     needs: source-check
 50 | 
 51 |     steps:
 52 |     - uses: dtolnay/rust-toolchain@master
 53 |       with:
 54 |         toolchain: ${{ env.RUST_MSRV_TOOLCHAIN }}
 55 |     - uses: actions/checkout@v4
 56 |     - run: cargo update -p tokio@1.43.0 --precise 1.38.1
 57 |     - run: cargo update -p tokio-stream@0.1.17 --precise 0.1.15
 58 |     - run: cargo update -p tokio-util@0.7.13 --precise 0.7.11
 59 |     - run: cargo doc --verbose --all-features --no-deps
 60 | 
 61 |   build-stable:
 62 |     runs-on: ubuntu-latest
 63 |     needs: source-check
 64 | 
 65 |     steps:
 66 |     - uses: dtolnay/rust-toolchain@master
 67 |       with:
 68 |         toolchain: stable
 69 |     - uses: actions/checkout@v4
 70 |     - run: cargo --version
 71 |     - run: cargo build --verbose --release --all-features
 72 | 
 73 |   test-stable:
 74 |     runs-on: ubuntu-latest
 75 |     needs: build-stable
 76 | 
 77 |     steps:
 78 |     - uses: dtolnay/rust-toolchain@master
 79 |       with:
 80 |         toolchain: stable
 81 |         default: true
 82 |     - uses: actions/checkout@v4
 83 |     - run: cargo test --verbose
 84 | 
 85 |   test-stable-all-features:
 86 |     runs-on: ubuntu-latest
 87 |     needs: test-stable
 88 | 
 89 |     steps:
 90 |     - uses: dtolnay/rust-toolchain@master
 91 |       with:
 92 |         toolchain: stable
 93 |         default: true
 94 |     - uses: actions/checkout@v4
 95 |     - run: cargo test --verbose --all-features
 96 | 
 97 |   build-msrv:
 98 |     runs-on: ubuntu-latest
 99 |     needs: build-stable
100 | 
101 |     steps:
102 |     - uses: dtolnay/rust-toolchain@master
103 |       with:
104 |         toolchain: ${{ env.RUST_MSRV_TOOLCHAIN }}
105 |         default: true
106 |     - uses: actions/checkout@v4
107 |     - run: cargo --version
108 |     - run: cargo update -p tokio@1.43.0 --precise 1.38.1
109 |     - run: cargo update -p tokio-stream@0.1.17 --precise 0.1.15
110 |     - run: cargo update -p tokio-util@0.7.13 --precise 0.7.11
111 |     - run: cargo build --verbose --release --all-features
112 | 
113 |   test-msrv-all-features:
114 |     runs-on: ubuntu-latest
115 |     needs: build-msrv
116 | 
117 |     steps:
118 |     - uses: dtolnay/rust-toolchain@master
119 |       with:
120 |         toolchain: ${{ env.RUST_MSRV_TOOLCHAIN }}
121 |         default: true
122 |     - uses: actions/checkout@v4
123 |     - run: cargo update -p tokio@1.43.0 --precise 1.38.1
124 |     - run: cargo update -p tokio-stream@0.1.17 --precise 0.1.15
125 |     - run: cargo update -p tokio-util@0.7.13 --precise 0.7.11
126 |     - run: cargo test --verbose --all-features
127 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | /target
2 | /.vscode
3 | Cargo.lock
4 | 


--------------------------------------------------------------------------------
/Cargo.toml:
--------------------------------------------------------------------------------
 1 | [package]
 2 | name = "rocket-governor"
 3 | version = "0.2.0-rc.4"
 4 | edition = "2021"
 5 | license = "MIT OR Apache-2.0"
 6 | description = "Implementation of Governor rate limiter for Rocket"
 7 | documentation = "https://docs.rs/rocket-governor"
 8 | homepage = "https://github.com/kolbma/rocket-governor"
 9 | repository = "https://github.com/kolbma/rocket-governor"
10 | keywords = ["rocket", "governor", "rate-limit", "rate-limiting", "ratelimit"]
11 | categories = ["web-programming"]
12 | rust-version = "1.69.0"
13 | 
14 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
15 | 
16 | [package.metadata]
17 | # minimum supported rust version 1.67.0 because of 3rd level dependency crate rocket->time
18 | # minimum supported rust version 1.69.0 because of incompatibility in macro expansion
19 | msrv = "1.69.0"
20 | 
21 | [package.metadata.docs.rs]
22 | all-features = true # get api doc for all features on https://docs.rs
23 | 
24 | [dependencies]
25 | governor = { version = "0.6", default-features = false, features = ["std"] }
26 | lazy_static = "1.4"
27 | # rocket = { path = "../rocket/core/lib" }
28 | # rocket = { git = "https://github.com/SergioBenitez/Rocket.git", branch = "v0.5" }
29 | rocket = { version = "0.5", default-features = false }
30 | 
31 | [features]
32 | limit_info = []
33 | logger = []
34 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | [![Rust](https://github.com/kolbma/rocket-governor/actions/workflows/rust.yml/badge.svg)](https://github.com/kolbma/rocket-governor/actions/workflows/rust.yml)
  2 | [![crates.io](https://img.shields.io/crates/v/rocket-governor)](https://crates.io/crates/rocket-governor)
  3 | [![docs](https://docs.rs/rocket-governor/badge.svg)](https://docs.rs/rocket-governor)
  4 | 
  5 | ## Description
  6 | 
  7 | Implementation of the [Governor](https://github.com/boinkor-net/governor.git) rate limiter for [Rocket](https://rocket.rs) web framework.
  8 | 
  9 | Rate limiting is used to control the rate requests are received and handled by an endpoint of the web application 
 10 | or web service.
 11 | 
 12 | ## Rocket specific features
 13 | 
 14 | Define as many rate limits with [Quota](https://docs.rs/governor/latest/governor/struct.Quota.html) of Governor
 15 | as you like and need in your Rocket web application/service.  
 16 | It is implemented as a Rocket [Request Guard](https://rocket.rs/v0.5/guide/requests/#request-guards) and provides
 17 | also an implementation of an [Error Catcher](https://rocket.rs/v0.5/guide/requests/#error-catchers).  
 18 | The Error Catcher can be registered on any path to handle [`Status::TooManyRequests`](https://api.rocket.rs/v0.5/rocket/http/struct.Status.html#associatedconstant.TooManyRequests) and provide HTTP headers in the response.
 19 | 
 20 | ## Usage
 21 | 
 22 | Add dependencies to [rocket-governor](https://crates.io/crates/rocket-governor) crate to your _Cargo.toml_.
 23 | 
 24 | Implement `RocketGovernable` for a _guard struct_ as you like: 
 25 | 
 26 | ```rust
 27 | use rocket_governor::{Method, Quota, RocketGovernable, RocketGovernor};
 28 | 
 29 | pub struct RateLimitGuard;
 30 | 
 31 | impl<'r> RocketGovernable<'r> for RateLimitGuard {
 32 |     fn quota(_method: Method, _route_name: &str) -> Quota {
 33 |         Quota::per_second(Self::nonzero(1u32))
 34 |     }
 35 | }
 36 | ```
 37 | 
 38 | This requires to implement the method `fn quota(_: Method, _: &str) -> Quota`.  
 39 | You can vary your `Quota` on any combination of __method__ and __route_name__, but the returned `Quota` should be a _static-like_. It should __not change__ between invocations of the `quota()`-method with equal parameters.
 40 | 
 41 | There is a small helper function `nonzero(u32)` for creating Quotas in your `quota()`-implementation e.g.:
 42 | ```rust
 43 |     Quota::per_second(Self::nonzero(1u32))
 44 | ```
 45 | 
 46 | After implementing the minimal requirements of trait `RocketGovernable` you can add your _Guard_ to your route-methods like:
 47 | 
 48 | ```rust
 49 | #[get("/")]
 50 | fn route_test(_limitguard: RocketGovernor<RateLimitGuard>) -> Status {
 51 |     Status::Ok
 52 | }
 53 | ```
 54 | 
 55 | ### Register Catcher
 56 | 
 57 | To handle HTTP Status 429 TooManyRequests there is an catcher-function implementation.
 58 | 
 59 | It is __called__ __`rocket_governor_catcher`__.  
 60 | 
 61 | Register it with `register(<path>, <catchers>)`-method of Rocket:
 62 | 
 63 | ```rust
 64 | use rocket_governor::rocket_governor_catcher;
 65 | 
 66 | #[launch]
 67 | fn launch_rocket() -> _ {
 68 |     rocket::build()
 69 |         .mount("/", routes![route_test])
 70 |         .register("/", catchers!(rocket_governor_catcher))
 71 | }
 72 | ```
 73 | 
 74 | ### Optional feature __limit_info__
 75 | 
 76 | There is the optional feature __limit_info__ which enables reporting about
 77 | rate limits in HTTP headers of requests.
 78 | 
 79 | The implementation is based on headers of
 80 | https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers
 81 | 
 82 | The feature provides a default implementation of a Rocket fairing
 83 | which need to be used to get the HTTP headers set.
 84 | 
 85 | See API documentation for [LimitHeaderGen](https://docs.rs/rocket-governor/latest/rocket_governor/struct.LimitHeaderGen.html).
 86 | 
 87 | For usage depend on it in Cargo.toml
 88 | ```toml
 89 | [dependencies]
 90 | rocket-governor = { version = "...", features = ["limit_info"] }
 91 | ```
 92 | 
 93 | ### Optional feature __logger__
 94 | 
 95 | There is the optional feature __logger__ which enables some logging output.
 96 | 
 97 | For usage depend on it in Cargo.toml
 98 | ```toml
 99 | [dependencies]
100 | rocket-governor = { version = "...", features = ["logger"] }
101 | ```
102 | 
103 | ### Additional information
104 | 
105 | To understand the basics of Rocket, please visit the _Rocket Guide_:
106 | * https://rocket.rs/v0.5/guide/requests/#request-guards
107 | * https://rocket.rs/v0.5/guide/requests/#error-catchers
108 | 
109 | ## Licenses
110 | 
111 | You can choose between __[MIT License](https://opensource.org/licenses/MIT)__ or __[Apache License 2.0](http://www.apache.org/licenses/LICENSE-2.0)__.
112 | 
113 | ### MIT License
114 | 
115 | Copyright (c) 2023 Markus Kolb
116 | 
117 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
118 | 
119 | The above copyright notice and this permission notice (including the next paragraph) shall be included in all copies or substantial portions of the Software.
120 | 
121 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
122 | 
123 | ### Apache License 2.0
124 | 
125 | Copyright 2023 Markus Kolb
126 | 
127 | Licensed under the Apache License, Version 2.0 (the "License");
128 | you may not use this file except in compliance with the License.
129 | You may obtain a copy of the License at
130 | 
131 | http://www.apache.org/licenses/LICENSE-2.0
132 | 
133 | Unless required by applicable law or agreed to in writing, software
134 | distributed under the License is distributed on an "AS IS" BASIS,
135 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
136 | See the License for the specific language governing permissions and
137 | limitations under the License.
138 | 


--------------------------------------------------------------------------------
/git-hooks/pre-commit:
--------------------------------------------------------------------------------
 1 | #!/bin/sh
 2 | #
 3 | 
 4 | script=$(basename "$0")
 5 | 
 6 | fail() {
 7 | 	echo "$script: pre-commit hook failed" >&2
 8 | 	exit 1
 9 | }
10 | 
11 | cargo fmt --check && \
12 | cargo clippy --all-features --no-deps -- -Dwarnings && \
13 | cargo msrv verify || fail
14 | 


--------------------------------------------------------------------------------
/git-hooks/pre-push:
--------------------------------------------------------------------------------
 1 | #!/bin/sh
 2 | 
 3 | script=$(basename "$0")
 4 | 
 5 | fail() {
 6 | 	echo "$script: pre-push hook failed" >&2
 7 | 	exit 1
 8 | }
 9 | 
10 | version=v$(sed -e '/^version[ =]/!d' -e 's/^version[ =]*//' -e 's/"//g' Cargo.toml)
11 | 
12 | git tag | grep -q "^$version$"
13 | [ "$?" -eq 0 ] || {
14 | 	git tag -s -m "build: release $version" "$version" || fail
15 | }
16 | 
17 | exit 0
18 | 


--------------------------------------------------------------------------------
/rocket-governor.code-workspace:
--------------------------------------------------------------------------------
 1 | {
 2 | 	"folders": [
 3 | 		{
 4 | 			"path": "."
 5 | 		}
 6 | 	],
 7 | 	"settings": {
 8 | 		"rust-analyzer.check.extraArgs": [
 9 | 			"--all-features"
10 | 		]
11 | 	}
12 | }
13 | 


--------------------------------------------------------------------------------
/src/header.rs:
--------------------------------------------------------------------------------
  1 | //! The headers used when a [RocketGovernor](super::RocketGovernor) guarded
  2 | //! path responds with [`TooManyRequests`](http::Status::TooManyRequests).  
  3 | //!
  4 | //! Depending on setup some headers are also set on responding successful
  5 | //! to let the client know about request limitations in the near future.
  6 | //!
  7 | //! There is an [RFC Draft](https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers).
  8 | // TODO: Check RFC Draft for publication
  9 | 
 10 | use rocket::http;
 11 | 
 12 | /// HTTP headers used for rate-limiting.
 13 | pub enum Header {
 14 |     /// Standard header for status 429 Too Many Requests
 15 |     /// ([RFC 6585](https://tools.ietf.org/html/rfc6585#section-4)).
 16 |     /// This should indicate a client for how long it should wait in seconds
 17 |     /// for retry.
 18 |     RetryAfter(u64),
 19 | 
 20 |     /// Custom header for reporting problems with rate limiter.
 21 |     XRateLimitError(&'static str),
 22 | 
 23 |     /// Header provides information about limitation of the route.
 24 |     XRateLimitLimit(u64),
 25 | 
 26 |     /// Header provides information about how many requests are left for the
 27 |     /// endpoint.
 28 |     XRateLimitRemaining(u64),
 29 | 
 30 |     /// Header provides the time in seconds when a request to the route is not
 31 |     /// rate limited and the rate limiter bucket is full again.
 32 |     XRateLimitReset(u64),
 33 | }
 34 | 
 35 | /// Standard header for status 429 Too Many Requests
 36 | /// ([RFC 6585](https://tools.ietf.org/html/rfc6585#section-4)).
 37 | /// This should indicate a client for how long it should wait in seconds for
 38 | /// retry.
 39 | pub const RETRY_AFTER: &str = "retry-after";
 40 | 
 41 | /// Custom header for reporting problems with rate limiter.
 42 | pub const X_RATELIMIT_ERROR: &str = "x-ratelimit-error";
 43 | 
 44 | // TODO: https://github.com/kolbma/rocket-governor/issues/2
 45 | 
 46 | // TODO: Check ratelimit-headers draft for publication
 47 | // https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers
 48 | 
 49 | /// Header provides information about limitation of the route.
 50 | pub const X_RATELIMIT_LIMIT: &str = "x-ratelimit-limit";
 51 | 
 52 | /// Header provides information about how many requests are left for the
 53 | /// endpoint.
 54 | pub const X_RATELIMIT_REMAINING: &str = "x-ratelimit-remaining";
 55 | 
 56 | /// Header provides the time in seconds when a request to the route is not
 57 | /// rate limited and the rate limiter bucket is full again.
 58 | pub const X_RATELIMIT_RESET: &str = "x-ratelimit-reset";
 59 | 
 60 | #[doc(hidden)]
 61 | impl From<Header> for http::Header<'_> {
 62 |     fn from(header: Header) -> Self {
 63 |         match header {
 64 |             Header::RetryAfter(sec) => http::Header::new(RETRY_AFTER, sec.to_string()),
 65 |             Header::XRateLimitError(err) => http::Header::new(X_RATELIMIT_ERROR, err),
 66 |             Header::XRateLimitLimit(limit) => {
 67 |                 http::Header::new(X_RATELIMIT_LIMIT, limit.to_string())
 68 |             }
 69 |             Header::XRateLimitRemaining(remaining) => {
 70 |                 http::Header::new(X_RATELIMIT_REMAINING, remaining.to_string())
 71 |             }
 72 |             Header::XRateLimitReset(sec) => http::Header::new(X_RATELIMIT_RESET, sec.to_string()),
 73 |         }
 74 |     }
 75 | }
 76 | 
 77 | #[cfg(test)]
 78 | mod tests {
 79 |     use super::Header;
 80 |     use super::{
 81 |         RETRY_AFTER, X_RATELIMIT_ERROR, X_RATELIMIT_LIMIT, X_RATELIMIT_REMAINING, X_RATELIMIT_RESET,
 82 |     };
 83 |     use rocket::http;
 84 |     use std::str::FromStr;
 85 | 
 86 |     #[test]
 87 |     fn test_header() {
 88 |         let h: http::Header = Header::RetryAfter(10).into();
 89 |         assert_eq!(RETRY_AFTER, h.name());
 90 |         assert_eq!(10, u64::from_str(h.value()).unwrap());
 91 | 
 92 |         let h: http::Header = Header::XRateLimitError("some error").into();
 93 |         assert_eq!(X_RATELIMIT_ERROR, h.name());
 94 |         assert_eq!("some error", h.value());
 95 | 
 96 |         let h: http::Header = Header::XRateLimitLimit(100).into();
 97 |         assert_eq!(X_RATELIMIT_LIMIT, h.name());
 98 |         assert_eq!(100, u64::from_str(h.value()).unwrap());
 99 | 
100 |         let h: http::Header = Header::XRateLimitRemaining(1).into();
101 |         assert_eq!(X_RATELIMIT_REMAINING, h.name());
102 |         assert_eq!(1, u64::from_str(h.value()).unwrap());
103 | 
104 |         let h: http::Header = Header::XRateLimitReset(5).into();
105 |         assert_eq!(X_RATELIMIT_RESET, h.name());
106 |         assert_eq!(5, u64::from_str(h.value()).unwrap());
107 |     }
108 | }
109 | 


--------------------------------------------------------------------------------
/src/lib.rs:
--------------------------------------------------------------------------------
  1 | //! # rocket-governor - rate-limiting implementation for Rocket web framework
  2 | //!
  3 | //! Provides the [rocket] guard implementing rate-limiting (based on [governor]).
  4 | //!
  5 | //! Declare a struct and use it with the generic [RocketGovernor] guard.  
  6 | //! This requires to implement trait [RocketGovernable] for your struct.
  7 | //!
  8 | //! ## Example
  9 | //!
 10 | //! ```rust
 11 | //! use rocket::{catchers, get, http::Status, launch, routes};
 12 | //! use rocket_governor::{rocket_governor_catcher, Method, Quota, RocketGovernable, RocketGovernor};
 13 | //!
 14 | //! pub struct RateLimitGuard;
 15 | //!
 16 | //! impl<'r> RocketGovernable<'r> for RateLimitGuard {
 17 | //!     fn quota(_method: Method, _route_name: &str) -> Quota {
 18 | //!         Quota::per_second(Self::nonzero(1u32))
 19 | //!     }
 20 | //! }
 21 | //!
 22 | //! #[get("/")]
 23 | //! fn route_example(_limitguard: RocketGovernor<RateLimitGuard>) -> Status {
 24 | //!     Status::Ok
 25 | //! }
 26 | //!
 27 | //! #[launch]
 28 | //! fn launch_rocket() -> _ {
 29 | //!     rocket::build()
 30 | //!         .mount("/", routes![route_example])
 31 | //!         .register("/", catchers![rocket_governor_catcher])
 32 | //! }
 33 | //! ```
 34 | //!
 35 | //! See [rocket-governor] Github project for more information.
 36 | //!
 37 | //! ## Features
 38 | //!
 39 | //! ### Optional feature __limit_info__
 40 | //!
 41 | //! There is the optional feature __limit_info__ which enables reporting about
 42 | //! rate limits in HTTP headers of requests.
 43 | //!
 44 | //! The implementation is based on headers of
 45 | //! [https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers]().
 46 | //! The feature provides a default implementation of a Rocket fairing
 47 | //! which need to be used to get the HTTP headers set.
 48 | //!
 49 | //! See API documentation for [LimitHeaderGen].
 50 | //!
 51 | //! For usage depend on it in Cargo.toml
 52 | //! ```toml
 53 | //! [dependencies]
 54 | //! rocket-governor = { version = "...", features = ["limit_info"] }
 55 | //! ```
 56 | //!
 57 | //! ### Optional feature __logger__
 58 | //!
 59 | //! There is the optional feature __logger__ which enables some logging output.
 60 | //!
 61 | //! For usage depend on it in Cargo.toml
 62 | //! ```toml
 63 | //! [dependencies]
 64 | //! rocket-governor = { version = "...", features = ["logger"] }
 65 | //! ```
 66 | //!
 67 | //! [governor]: https://docs.rs/governor/
 68 | //! [rocket]: https://docs.rs/rocket/
 69 | //! [rocket-governor]: https://github.com/kolbma/rocket-governor/
 70 | 
 71 | #![deny(clippy::all)]
 72 | #![deny(keyword_idents)]
 73 | #![deny(missing_docs)]
 74 | #![deny(non_ascii_idents)]
 75 | #![deny(unreachable_pub)]
 76 | #![deny(unsafe_code)]
 77 | #![deny(unused_crate_dependencies)]
 78 | #![deny(unused_qualifications)]
 79 | //#![deny(unused_results)]
 80 | #![deny(warnings)]
 81 | 
 82 | use governor::clock::{Clock, DefaultClock};
 83 | pub use governor::Quota;
 84 | use lazy_static::lazy_static;
 85 | pub use limit_error::LimitError;
 86 | #[cfg(feature = "limit_info")]
 87 | pub use limit_header_gen::LimitHeaderGen;
 88 | use logger::{error, info, trace};
 89 | use registry::Registry;
 90 | #[cfg(feature = "limit_info")]
 91 | pub use req_state::ReqState;
 92 | pub use rocket::http::Method;
 93 | use rocket::{
 94 |     async_trait, catch,
 95 |     http::Status,
 96 |     request::{FromRequest, Outcome},
 97 |     Request,
 98 | };
 99 | pub use rocket_governable::RocketGovernable;
100 | use std::marker::PhantomData;
101 | pub use std::num::NonZeroU32;
102 | 
103 | pub mod header;
104 | mod limit_error;
105 | #[cfg(feature = "limit_info")]
106 | mod limit_header_gen;
107 | mod logger;
108 | mod registry;
109 | #[cfg(feature = "limit_info")]
110 | mod req_state;
111 | mod rocket_governable;
112 | 
113 | /// Generic [RocketGovernor] implementation.
114 | ///
115 | /// [rocket_governor](crate) is a [rocket] guard implementation of the
116 | /// [governor] rate limiter.
117 | ///
118 | /// Declare a struct and use it with the generic [RocketGovernor] guard.
119 | /// This requires to implement [RocketGovernable] for your struct.
120 | ///
121 | /// See the top level [crate] documentation.
122 | ///
123 | /// [governor]: https://docs.rs/governor/
124 | /// [rocket]: https://docs.rs/rocket/
125 | ///
126 | pub struct RocketGovernor<'r, T>
127 | where
128 |     T: RocketGovernable<'r>,
129 | {
130 |     _phantom: PhantomData<&'r T>,
131 | }
132 | 
133 | lazy_static! {
134 |     static ref CLOCK: DefaultClock = DefaultClock::default();
135 | }
136 | 
137 | #[doc(hidden)]
138 | impl<'r, T> RocketGovernor<'r, T>
139 | where
140 |     T: RocketGovernable<'r>,
141 | {
142 |     /// Handler used in `FromRequest::from_request(request: &'r Request)`.
143 |     #[inline(always)]
144 |     pub fn handle_from_request(request: &'r Request) -> Outcome<Self, LimitError> {
145 |         let res = request.local_cache(|| {
146 |             if let Some(route) = request.route() {
147 |                 if let Some(route_name) = &route.name {
148 |                     let limiter = Registry::get_or_insert::<T>(
149 |                         route.method,
150 |                         route_name,
151 |                         T::quota(route.method, route_name),
152 |                     );
153 |                     if let Some(client_ip) = request.client_ip() {
154 |                         let limit_check_res = limiter.check_key(&client_ip);
155 |                         match limit_check_res {
156 |                             Ok(state) => {
157 |                                 #[allow(unused_variables)] // only used in trace or when feature limit_info
158 |                                 let request_capacity = state.remaining_burst_capacity();
159 |                                 trace!(
160 |                                     "not governed ip {} method {} route {}: remaining request capacity {}",
161 |                                     &client_ip,
162 |                                     &route.method,
163 |                                     route_name,
164 |                                     request_capacity
165 |                                 );
166 | 
167 |                                 #[cfg(feature = "limit_info")] {
168 |                                     // `local_cache` lookup works by type and so it doesn't work to catch
169 |                                     // `LimitError` and handle different Ok objects:
170 |                                     // See https://rocket.rs/v0.5/guide/state/#request-local-state
171 |                                     // State wrapper is so cached separate...
172 |                                     let req_state = ReqState::new(state.quota(), request_capacity);
173 |                                     let is_req_state_allowed = T::limit_info_allow(Some(route.method), Some(route_name), &req_state);
174 |                                     if is_req_state_allowed {
175 |                                         // For safety and speed this is used by default in a limited way, see:
176 |                                         // * Information disclosure:
177 |                                         //   https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers#section-6.2
178 |                                         //
179 |                                         let _ = request.local_cache(|| req_state);
180 |                                     }
181 |                                 }
182 | 
183 |                                 Ok(()) // needs to be something not changing during request
184 |                             }
185 |                             Err(notuntil) => {
186 |                                 let wait_time = notuntil.wait_time_from(CLOCK.now()).as_secs();
187 |                                 info!(
188 |                                     "ip {} method {} route {} limited {} sec",
189 |                                     &client_ip, &route.method, route_name, &wait_time
190 |                                 );
191 |                                 Err(LimitError::GovernedRequest(wait_time, notuntil.quota()))
192 |                             }
193 |                         }
194 |                     } else {
195 |                         error!(
196 |                             "missing ip - method {} route {}: request: {:?}",
197 |                             &route.method, route_name, request
198 |                         );
199 |                         Err(LimitError::MissingClientIpAddr)
200 |                     }
201 |                 } else {
202 |                     error!("route without name: request: {:?}", request);
203 |                     Err(LimitError::MissingRouteName)
204 |                 }
205 |             } else {
206 |                 error!("routing failure: request: {:?}", request);
207 |                 Err(LimitError::MissingRoute)
208 |             }
209 |         });
210 | 
211 |         match res {
212 |             Ok(_) => {
213 |                 #[cfg(feature = "limit_info")]
214 |                 {
215 |                     // available if `T::limit_info_allow()` is true
216 |                     let state_opt = ReqState::get_or_default(request);
217 |                     #[allow(unused_variables)] // state only used in trace
218 |                     if let Some(state) = state_opt {
219 |                         trace!(
220 |                             "request_capacity: {} rate-limit: {}",
221 |                             state.request_capacity,
222 |                             state.quota.burst_size().get()
223 |                         );
224 |                     }
225 |                 }
226 | 
227 |                 // Forward request
228 |                 Outcome::Success(Self::default())
229 |             }
230 |             Err(e) => {
231 |                 let e = e.clone();
232 |                 match e {
233 |                     LimitError::GovernedRequest(_, _) => {
234 |                         Outcome::Error((Status::TooManyRequests, e))
235 |                     }
236 |                     _ => Outcome::Error((Status::BadRequest, e)),
237 |                 }
238 |             }
239 |         }
240 |     }
241 | }
242 | 
243 | #[doc(hidden)]
244 | impl<'r, T> Default for RocketGovernor<'r, T>
245 | where
246 |     T: RocketGovernable<'r>,
247 | {
248 |     fn default() -> Self {
249 |         Self {
250 |             _phantom: PhantomData,
251 |         }
252 |     }
253 | }
254 | 
255 | #[doc(hidden)]
256 | #[async_trait]
257 | impl<'r, T> FromRequest<'r> for RocketGovernor<'r, T>
258 | where
259 |     T: RocketGovernable<'r>,
260 | {
261 |     type Error = LimitError;
262 | 
263 |     async fn from_request(request: &'r Request<'_>) -> Outcome<Self, LimitError> {
264 |         Self::handle_from_request(request)
265 |     }
266 | }
267 | 
268 | /// A default implementation for Rocket [Catcher] handling HTTP TooManyRequests responses.
269 | ///
270 | /// ## Example
271 | ///
272 | /// ```rust
273 | /// use rocket::{catchers, launch};
274 | /// use rocket_governor::rocket_governor_catcher;
275 | ///
276 | /// #[launch]
277 | /// fn launch_rocket() -> _ {
278 | ///     rocket::build()
279 | ///         .register("/", catchers![rocket_governor_catcher])
280 | /// }
281 | /// ```
282 | ///
283 | /// [Catcher]: https://api.rocket.rs/v0.5/rocket/struct.Catcher.html
284 | #[catch(429)]
285 | pub fn rocket_governor_catcher<'r>(request: &'r Request) -> &'r LimitError {
286 |     let cached_res: &Result<(), LimitError> = request.local_cache(|| Err(LimitError::Error));
287 |     if let Err(limit_err) = cached_res {
288 |         limit_err
289 |     } else {
290 |         &LimitError::Error
291 |     }
292 | }
293 | 


--------------------------------------------------------------------------------
/src/limit_error.rs:
--------------------------------------------------------------------------------
 1 | //! Errors for governed requests which implement
 2 | //! [Responder](rocket::response::Responder).
 3 | 
 4 | use super::{header::Header, Quota};
 5 | use rocket::{
 6 |     response::{self, Responder},
 7 |     Request,
 8 | };
 9 | 
10 | mod catcher;
11 | 
12 | /// Errors for governed requests which implement
13 | /// [Responder](rocket::response::Responder).
14 | #[derive(Clone, Debug)]
15 | pub enum LimitError {
16 |     /// Any other undefined LimitError
17 |     Error,
18 | 
19 |     // TODO: Check ratelimit headers DRAFT for publication
20 |     /// Governed request for the next provided seconds.
21 |     /// Provided `Quota` will be used for setting additional
22 |     /// HTTP headers defined by
23 |     /// [draft-ietf-httpapi-ratelimit-headers](https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers).
24 |     /// These headers could be used in service clients to use the service
25 |     /// in a more compliant way for its resources.
26 |     GovernedRequest(u64, Quota),
27 | 
28 |     /// There is no remote client IP address known in the request. Might be
29 |     /// a misconfigured server environment.
30 |     MissingClientIpAddr,
31 | 
32 |     /// Route is not available which might be only the case in fairings
33 |     MissingRoute,
34 | 
35 |     /// There is a route without name and this can not be matched for
36 |     /// rate limiting
37 |     MissingRouteName,
38 | }
39 | 
40 | /// Implements [Responder] to provide
41 | /// [Result](rocket::response::Result) possibilities.
42 | impl<'r, 'o: 'r> Responder<'r, 'o> for &LimitError {
43 |     fn respond_to(self, request: &'r Request<'_>) -> response::Result<'o> {
44 |         let mut handler = catcher::too_many_requests_handler(request);
45 | 
46 |         match self {
47 |             LimitError::Error => {
48 |                 handler.set_header(Header::XRateLimitError("rate limiter error"));
49 |             }
50 |             LimitError::GovernedRequest(wait_time, quota) => {
51 |                 handler.set_header(Header::RetryAfter(*wait_time));
52 |                 // TODO: x-ratelimit-limit can describe the time window of limit
53 |                 //       https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers#section-5.1
54 |                 handler.set_header(Header::XRateLimitLimit(quota.burst_size().get() as u64));
55 |                 // XRateLimitRemaining makes no sense here in LimitError
56 |                 // because `state.remaining_burst_capacity()` should be
57 |                 // always 0
58 |                 handler.set_header(Header::XRateLimitReset(
59 |                     quota.burst_size_replenished_in().as_secs(),
60 |                 ));
61 |             }
62 |             LimitError::MissingClientIpAddr => {
63 |                 handler.set_header(Header::XRateLimitError(
64 |                     "application not retrieving client ip",
65 |                 ));
66 |             }
67 |             LimitError::MissingRoute => {
68 |                 handler.set_header(Header::XRateLimitError("routing failure"));
69 |             }
70 |             LimitError::MissingRouteName => {
71 |                 handler.set_header(Header::XRateLimitError("route without name"));
72 |             }
73 |         };
74 | 
75 |         Ok(handler)
76 |     }
77 | }
78 | 


--------------------------------------------------------------------------------
/src/limit_error/catcher.rs:
--------------------------------------------------------------------------------
  1 | //! Provides a similiar catcher for html/json like the Rocket framework
  2 | 
  3 | use rocket::{
  4 |     http::{ContentType, Status},
  5 |     Request, Response,
  6 | };
  7 | use std::{borrow::Cow, io::Cursor};
  8 | 
  9 | // copied from rocket/src/catcher/catcher.rs because it is also pub(crate)
 10 | 
 11 | macro_rules! html_error_template {
 12 |     ($code:expr, $reason:expr, $description:expr) => {
 13 |         concat!(
 14 |             r#"<!DOCTYPE html>
 15 | <html lang="en">
 16 | <head>
 17 |     <meta charset="utf-8">
 18 |     <title>"#,
 19 |             $code,
 20 |             " ",
 21 |             $reason,
 22 |             r#"</title>
 23 | </head>
 24 | <body align="center">
 25 |     <div role="main" align="center">
 26 |         <h1>"#,
 27 |             $code,
 28 |             ": ",
 29 |             $reason,
 30 |             r#"</h1>
 31 |         <p>"#,
 32 |             $description,
 33 |             r#"</p>
 34 |         <hr />
 35 |     </div>
 36 |     <div role="contentinfo" align="center">
 37 |         <small>Rocket</small>
 38 |     </div>
 39 | </body>
 40 | </html>"#
 41 |         )
 42 |     };
 43 | }
 44 | 
 45 | macro_rules! json_error_template {
 46 |     ($code:expr, $reason:expr, $description:expr) => {
 47 |         concat!(
 48 |             r#"{
 49 |   "error": {
 50 |     "code": "#,
 51 |             $code,
 52 |             r#",
 53 |     "reason": ""#,
 54 |             $reason,
 55 |             r#"",
 56 |     "description": ""#,
 57 |             $description,
 58 |             r#""
 59 |   }
 60 | }"#
 61 |         )
 62 |     };
 63 | }
 64 | 
 65 | // also copied and modified
 66 | 
 67 | /// Create the handler for `Status::TooManyRequests`.
 68 | pub(crate) fn too_many_requests_handler<'r, 'o: 'r>(req: &'r Request<'_>) -> Response<'o> {
 69 |     let preferred = req.accept().map(|a| a.preferred());
 70 |     let (mime, text) = if preferred.map_or(false, |a| a.is_json()) {
 71 |         let json: Cow<'_, str> = json_error_template!(
 72 |             429,
 73 |             "Too Many Requests",
 74 |             "Too many requests have been received recently."
 75 |         )
 76 |         .into();
 77 | 
 78 |         (ContentType::JSON, json)
 79 |     } else {
 80 |         let html: Cow<'_, str> = html_error_template!(
 81 |             429,
 82 |             "Too Many Requests",
 83 |             "Too many requests have been received recently."
 84 |         )
 85 |         .into();
 86 | 
 87 |         (ContentType::HTML, html)
 88 |     };
 89 | 
 90 |     let mut r = Response::build()
 91 |         .status(Status::TooManyRequests)
 92 |         .header(mime)
 93 |         .finalize();
 94 |     match text {
 95 |         Cow::Owned(v) => r.set_sized_body(v.len(), Cursor::new(v)),
 96 |         Cow::Borrowed(v) => r.set_sized_body(v.len(), Cursor::new(v)),
 97 |     };
 98 | 
 99 |     r
100 | }
101 | 


--------------------------------------------------------------------------------
/src/limit_header_gen.rs:
--------------------------------------------------------------------------------
 1 | //! Provides [`Fairing`](rocket::fairing::Fairing) in the implementation
 2 | //! [`LimitHeaderGen`] which is [attachable](rocket::Rocket::attach()) to
 3 | //! [`Rocket`](rocket::Rocket)-instance.
 4 | 
 5 | use crate::{header::Header, ReqState};
 6 | use rocket::{
 7 |     fairing::{Fairing, Info, Kind},
 8 |     Request, Response,
 9 | };
10 | 
11 | /// Provides [`Fairing`](rocket::fairing::Fairing) implementation
12 | /// which is [attachable](rocket::Rocket::attach()) to
13 | /// [`Rocket`](rocket::Rocket)-instance.
14 | ///
15 | /// `LimitHeaderGen` generates [HTTP headers](crate::header) in
16 | /// [Reponses](rocket::Response) to [Requests](rocket::Request) at specific
17 | /// [mounted routes](rocket::Rocket::mount()).
18 | ///
19 | /// The [`Route`](rocket::Route) needs a [RocketGovernor guard](crate).
20 | /// The guard is responsible for managing the governor rate limit.
21 | ///
22 | /// Depending on the implementation of
23 | /// [`RocketGovernable::limit_info_allow()`](crate::RocketGovernable::limit_info_allow()),
24 | /// `LimitHeaderGen` fairing sets [HTTP headers](crate::header) like
25 | ///
26 | /// * [X_RATELIMIT_LIMIT](crate::header::X_RATELIMIT_LIMIT)
27 | /// * [X_RATELIMIT_REMAINING](crate::header::X_RATELIMIT_REMAINING)
28 | ///
29 | /// which can be used by HTTP clients to adjust service requests.
30 | ///
31 | /// ## Example usage
32 | ///
33 | /// ```rust
34 | /// use rocket;
35 | /// use rocket_governor;
36 | ///
37 | /// #[rocket::launch]
38 | /// fn launch_rocket() -> _ {
39 | ///     rocket::build().attach(rocket_governor::LimitHeaderGen::default())
40 | /// }
41 | /// ```
42 | ///
43 | pub struct LimitHeaderGen;
44 | 
45 | impl Default for LimitHeaderGen {
46 |     fn default() -> Self {
47 |         Self
48 |     }
49 | }
50 | 
51 | #[rocket::async_trait]
52 | impl Fairing for LimitHeaderGen {
53 |     fn info(&self) -> Info {
54 |         Info {
55 |             name: "RateLimit Header Generator",
56 |             kind: Kind::Response,
57 |         }
58 |     }
59 | 
60 |     /// Set rate limit headers if
61 |     /// [`RocketGovernable::limit_info_allow()`](crate::RocketGovernable::limit_info_allow())
62 |     /// returns true and [`ReqState`] has been cached in this case in
63 |     /// [`Request::local_cache`].
64 |     async fn on_response<'r>(&self, request: &'r Request<'_>, response: &mut Response<'r>) {
65 |         let state_opt = ReqState::get_or_default(request);
66 | 
67 |         if let Some(state) = state_opt {
68 |             response.set_header(Header::XRateLimitLimit(
69 |                 state.quota.burst_size().get().into(),
70 |             ));
71 |             response.set_header(Header::XRateLimitRemaining(state.request_capacity.into()));
72 |         }
73 |     }
74 | }
75 | 


--------------------------------------------------------------------------------
/src/logger.rs:
--------------------------------------------------------------------------------
 1 | #![allow(clippy::single_component_path_imports)] // this is how the feature works
 2 | 
 3 | #[cfg(feature = "logger")]
 4 | pub(crate) use rocket::{debug, error, info, trace};
 5 | 
 6 | #[cfg(not(feature = "logger"))]
 7 | macro_rules! debug {
 8 |     (target: $target:expr, $($arg:tt)+) => {};
 9 |     ($($arg:tt)+) => {};
10 | }
11 | #[cfg(not(feature = "logger"))]
12 | pub(crate) use debug;
13 | 
14 | #[cfg(not(feature = "logger"))]
15 | macro_rules! error {
16 |     (target: $target:expr, $($arg:tt)+) => {};
17 |     ($($arg:tt)+) => {};
18 | }
19 | #[cfg(not(feature = "logger"))]
20 | pub(crate) use error;
21 | 
22 | #[cfg(not(feature = "logger"))]
23 | macro_rules! info {
24 |     (target: $target:expr, $($arg:tt)+) => {};
25 |     ($($arg:tt)+) => {};
26 | }
27 | #[cfg(not(feature = "logger"))]
28 | pub(crate) use info;
29 | 
30 | #[cfg(not(feature = "logger"))]
31 | macro_rules! trace {
32 |     (target: $target:expr, $($arg:tt)+) => {};
33 |     ($($arg:tt)+) => {};
34 | }
35 | #[cfg(not(feature = "logger"))]
36 | pub(crate) use trace;
37 | 


--------------------------------------------------------------------------------
/src/registry.rs:
--------------------------------------------------------------------------------
 1 | use crate::logger::debug;
 2 | use governor::{
 3 |     clock::DefaultClock, middleware::StateInformationMiddleware,
 4 |     state::keyed::DefaultKeyedStateStore, Quota, RateLimiter,
 5 | };
 6 | use lazy_static::lazy_static;
 7 | use rocket::http::Method;
 8 | use std::{
 9 |     any::type_name,
10 |     collections::HashMap,
11 |     net::IpAddr,
12 |     sync::{Arc, RwLock},
13 | };
14 | 
15 | pub(crate) type RegisteredRateLimiter = Arc<
16 |     RateLimiter<IpAddr, DefaultKeyedStateStore<IpAddr>, DefaultClock, StateInformationMiddleware>,
17 | >;
18 | 
19 | #[derive(Debug)]
20 | pub(crate) struct Registry {
21 |     limiter: RwLock<HashMap<Method, HashMap<String, RegisteredRateLimiter>>>,
22 | }
23 | 
24 | impl Registry {
25 |     pub(crate) fn get_or_insert<T>(
26 |         method: Method,
27 |         route_name: &str,
28 |         quota: Quota,
29 |     ) -> RegisteredRateLimiter {
30 |         let route_name = type_name::<T>().to_string() + "::" + route_name;
31 | 
32 |         // check if exist with readlock
33 |         let limiter = if let Ok(rlock) = REG.limiter.read() {
34 |             if let Some(meth_found) = rlock.get(&method) {
35 |                 if let Some(limiter) = meth_found.get(&route_name) {
36 |                     debug!("limiter found method {} route {}", &method, &route_name);
37 |                     return Arc::clone(limiter);
38 |                 }
39 |             }
40 |             None
41 |         } else {
42 |             None
43 |         };
44 | 
45 |         // get the existing limiter or create the new one with writelock and return the created
46 |         let limiter = if let Some(limiter) = limiter {
47 |             limiter
48 |         } else {
49 |             let mut wlock_meth_map = REG.limiter.write().unwrap();
50 |             if let Some(meth_found) = wlock_meth_map.get_mut(&method) {
51 |                 if let Some(limiter) = meth_found.get(&route_name) {
52 |                     debug!("limiter found method {} route {}", &method, &route_name);
53 |                     Arc::clone(limiter)
54 |                 } else {
55 |                     debug!("new limiter method {} route {}", &method, &route_name);
56 |                     let limiter = Arc::new(
57 |                         RateLimiter::keyed(quota).with_middleware::<StateInformationMiddleware>(),
58 |                     );
59 |                     meth_found.insert(route_name, Arc::clone(&limiter));
60 |                     limiter
61 |                 }
62 |             } else {
63 |                 debug!("new limiter method {} route {}", &method, &route_name);
64 |                 let mut lim_map = HashMap::new();
65 |                 let limiter = Arc::new(
66 |                     RateLimiter::keyed(quota).with_middleware::<StateInformationMiddleware>(),
67 |                 );
68 |                 lim_map.insert(route_name, Arc::clone(&limiter));
69 |                 wlock_meth_map.insert(method, lim_map);
70 |                 limiter
71 |             }
72 |         };
73 | 
74 |         // simple cleanup
75 |         limiter.retain_recent();
76 |         limiter.shrink_to_fit();
77 | 
78 |         limiter
79 |     }
80 | }
81 | 
82 | impl Default for Registry {
83 |     fn default() -> Self {
84 |         Self {
85 |             limiter: RwLock::new(HashMap::new()),
86 |         }
87 |     }
88 | }
89 | 
90 | lazy_static! {
91 |     static ref REG: Registry = Registry::default();
92 | }
93 | 


--------------------------------------------------------------------------------
/src/req_state.rs:
--------------------------------------------------------------------------------
  1 | //! This module provides the [ReqState] which you may get in contact
  2 | //! during implementation of [RocketGovernable](super::RocketGovernable).
  3 | //!
  4 | //! __Available__ only with __feature limit_info__!
  5 | //!
  6 | //! There might be no need to create an object by yourself.
  7 | //!
  8 | //! [ReqState::default()] provides an object which might be only useful in the
  9 | //! inner library usage.
 10 | 
 11 | use super::{NonZeroU32, Quota};
 12 | use rocket::Request;
 13 | 
 14 | /// `ReqState` is the data struct to handle information about [`Quota`] and
 15 | /// limits in the [`Request`] state.
 16 | ///
 17 | /// This info is provided to the client by [HTTP headers](crate::header).
 18 | ///
 19 | /// When implementing
 20 | /// [RocketGovernable::limit_info_allow()](super::RocketGovernable::limit_info_allow())
 21 | /// there is passed an object `ReqState` as parameter to provide the
 22 | /// possibility for individual decisions about setting
 23 | /// [HTTP headers](crate::header).
 24 | #[derive(Debug)]
 25 | pub struct ReqState {
 26 |     /// The [`Request::local_cache()`] is initialized on first call of [`ReqState::get_or_default()`].  
 27 |     /// This is detected by `is_default`.
 28 |     is_default: bool,
 29 | 
 30 |     /// The [Quota](super::Quota) of the current [Request].
 31 |     pub(crate) quota: Quota,
 32 | 
 33 |     /// Number of [Requests](Request) which can be done, before the end point is
 34 |     /// limited.
 35 |     ///
 36 |     /// Limitation starts below `0`.
 37 |     pub(crate) request_capacity: u32,
 38 | }
 39 | 
 40 | impl ReqState {
 41 |     /// Create new [`ReqState`] with provided values
 42 |     pub(crate) fn new(quota: Quota, request_capacity: u32) -> Self {
 43 |         Self {
 44 |             is_default: false,
 45 |             quota,
 46 |             request_capacity,
 47 |         }
 48 |     }
 49 | 
 50 |     /// Retrieve cached [`ReqState`] from [`Request::local_cache()`]
 51 |     ///
 52 |     /// Be careful, you'll set cache value to [`ReqState::default()`] if there
 53 |     /// has been cached no other value before.
 54 |     pub(crate) fn get_or_default<'r>(request: &'r Request) -> Option<&'r Self> {
 55 |         let state: &ReqState = request.local_cache(|| {
 56 |             ReqState::default() // dummy, because we ignore it if not set
 57 |         });
 58 | 
 59 |         if !state.is_default {
 60 |             Some(state)
 61 |         } else {
 62 |             None
 63 |         }
 64 |     }
 65 | 
 66 |     /// The [Quota](super::Quota) of the current [Request].
 67 |     pub fn quota(&self) -> &Quota {
 68 |         &self.quota
 69 |     }
 70 | 
 71 |     /// Number of [Requests](Request) which can be done, before the end point is
 72 |     /// limited.
 73 |     ///
 74 |     /// Limitation starts below `0`.
 75 |     pub fn request_capacity(&self) -> u32 {
 76 |         self.request_capacity
 77 |     }
 78 | }
 79 | 
 80 | impl Default for ReqState {
 81 |     /// No meaningful implementation.  
 82 |     /// Just to provide some default to get `ReqState::get_or_default()` to
 83 |     /// work.
 84 |     fn default() -> Self {
 85 |         Self {
 86 |             is_default: true,
 87 |             quota: Quota::per_second(NonZeroU32::new(1).unwrap()),
 88 |             request_capacity: 0,
 89 |         }
 90 |     }
 91 | }
 92 | 
 93 | #[cfg(test)]
 94 | mod tests {
 95 |     use super::*;
 96 |     use rocket::{
 97 |         get,
 98 |         http::{Header, Status},
 99 |         local::blocking::Client,
100 |         routes, Build, Rocket,
101 |     };
102 | 
103 |     #[get("/")]
104 |     fn route_test() -> Status {
105 |         Status::Ok
106 |     }
107 | 
108 |     fn launch_rocket() -> Rocket<Build> {
109 |         rocket::build().mount("/", routes![route_test])
110 |     }
111 | 
112 |     #[test]
113 |     fn test_req_state() {
114 |         let client = Client::untracked(launch_rocket()).expect("no rocket instance");
115 |         let mut req = client.get("/");
116 |         req.add_header(Header::new("X-Real-IP", "127.1.1.1"));
117 |         // req.dispatch();
118 |         let request = req.inner_mut();
119 |         let _ = request
120 |             .local_cache(|| ReqState::new(Quota::per_second(NonZeroU32::new(1).unwrap()), 10));
121 |         let _ = request.real_ip();
122 |         let req_state = ReqState::get_or_default(request);
123 | 
124 |         assert!(req_state.is_some());
125 |         assert_eq!(req_state.unwrap().request_capacity, 10);
126 | 
127 |         // 2nd time
128 |         let req_state = ReqState::get_or_default(request);
129 | 
130 |         assert!(req_state.is_some());
131 |         assert_eq!(req_state.unwrap().request_capacity, 10);
132 | 
133 |         // default is none
134 |         let mut req = client.get("/");
135 |         req.add_header(Header::new("X-Real-IP", "127.1.1.2"));
136 |         let request = req.inner_mut();
137 |         let req_state = ReqState::get_or_default(request);
138 | 
139 |         assert!(req_state.is_none());
140 |     }
141 | }
142 | 


--------------------------------------------------------------------------------
/src/rocket_governable.rs:
--------------------------------------------------------------------------------
  1 | //! Module for trait [RocketGovernable].
  2 | 
  3 | #[cfg(feature = "limit_info")]
  4 | use super::ReqState;
  5 | use super::{Method, NonZeroU32, Quota};
  6 | use rocket::async_trait;
  7 | 
  8 | /// The [RocketGovernable] guard trait.
  9 | ///
 10 | /// [rocket-governor](crate) is a [rocket] guard implementation of the
 11 | /// [governor] rate limiter.
 12 | ///
 13 | /// Declare a struct and use it with the generic
 14 | /// [RocketGovernor](super::RocketGovernor) guard.
 15 | /// This requires to implement [RocketGovernable] for your struct.
 16 | ///
 17 | /// See the top level [crate] documentation.
 18 | ///
 19 | /// [governor]: https://docs.rs/governor/
 20 | /// [rocket]: https://docs.rs/rocket/
 21 | ///
 22 | #[async_trait]
 23 | pub trait RocketGovernable<'r> {
 24 |     /// Returns the [Quota] of the [RocketGovernable].
 25 |     ///
 26 |     /// This is called only once per method/route_name combination.
 27 |     /// So it makes only sense to return always the same [Quota] for
 28 |     /// equal parameter combinations and no dynamic calculation.
 29 |     ///
 30 |     /// This is also the requirement to have correct information set
 31 |     /// in HTTP headers by registered [`rocket_governor_catcher()`](crate::rocket_governor_catcher()).
 32 |     ///
 33 |     /// [Quota]: https://docs.rs/governor/latest/governor/struct.Quota.html
 34 |     #[must_use]
 35 |     fn quota(method: Method, route_name: &str) -> Quota;
 36 | 
 37 |     /// Returns `true` if HTTP rate limit info [headers](crate::header)
 38 |     /// should be set in requests.
 39 |     ///
 40 |     /// Implement [`limit_info_allow()`](RocketGovernable::limit_info_allow())
 41 |     /// to change to your preference.  
 42 |     ///
 43 |     /// The trait implementation enables info headers only just the request
 44 |     /// before any further request would be rate limited.
 45 |     /// This is because of **speed**, **bandwidth** and **safety**.
 46 |     ///
 47 |     /// In
 48 |     /// [draft-ietf-httpapi-ratelimit-headers#section-6.2](https://datatracker.ietf.org/doc/html/draft-ietf-httpapi-ratelimit-headers#section-6.2)
 49 |     /// you can read the following information about...
 50 |     ///
 51 |     /// ## Information disclosure
 52 |     ///
 53 |     /// Servers should not disclose to untrusted parties operational capacity
 54 |     /// information that can be used to saturate its infrastructural
 55 |     /// resources.
 56 |     ///
 57 |     /// While this specification does not mandate whether non 2xx responses
 58 |     /// consume quota, if 401 and 403 responses count on quota a malicious
 59 |     /// client could probe the endpoint to get traffic information of another
 60 |     /// user.
 61 |     ///
 62 |     /// As intermediaries might retransmit requests and consume quota-units
 63 |     /// without prior knowledge of the User Agent, RateLimit fields might
 64 |     /// reveal the existence of an intermediary to the User Agent.
 65 |     ///
 66 |     /// ## Feature availability
 67 |     ///
 68 |     /// [`limit_info_allow()`](RocketGovernable::limit_info_allow()) is only
 69 |     /// available when feature __limit_info__ is enabled.
 70 |     ///
 71 |     /// ## Example
 72 |     ///
 73 |     /// ```rust
 74 |     /// use rocket_governor::{Method, Quota, ReqState, RocketGovernable};
 75 |     ///
 76 |     /// pub struct RateLimitGuard;
 77 |     ///
 78 |     /// impl<'r> RocketGovernable<'r> for RateLimitGuard {
 79 |     ///     fn quota(_method: Method, _route_name: &str) -> Quota {
 80 |     ///         Quota::per_second(Self::nonzero(1u32))
 81 |     ///     }
 82 |     ///
 83 |     ///     fn limit_info_allow(
 84 |     ///         method: Option<Method>,
 85 |     ///         route_name: Option<&str>,
 86 |     ///         state: &ReqState,
 87 |     ///     ) -> bool {
 88 |     ///         let mut cap = 1;
 89 |     ///         if let Some(m) = method {
 90 |     ///             if m == Method::Post {
 91 |     ///                 if let Some(route) = route_name {
 92 |     ///                     if route == "admin_action" {
 93 |     ///                         cap = 20;
 94 |     ///                     }
 95 |     ///                 }
 96 |     ///             }
 97 |     ///         }
 98 |     ///         state.request_capacity() <= cap
 99 |     ///     }
100 |     /// }
101 |     ///
102 |     /// ```
103 |     ///
104 |     // TODO: update RFC link when published
105 |     #[cfg(feature = "limit_info")]
106 |     fn limit_info_allow(
107 |         method: Option<Method>,
108 |         route_name: Option<&str>,
109 |         state: &ReqState,
110 |     ) -> bool {
111 |         let (_, _) = (method, route_name); // unused warning
112 | 
113 |         state.request_capacity <= 1
114 |     }
115 | 
116 |     /// Converts a non-zero number [u32] to [NonZeroU32](std::num::NonZeroU32).
117 |     ///
118 |     /// Number zero/0 becomes 1.
119 |     ///
120 |     #[inline]
121 |     fn nonzero(n: u32) -> NonZeroU32 {
122 |         NonZeroU32::new(n).unwrap_or_else(|| NonZeroU32::new(1u32).unwrap())
123 |     }
124 | }
125 | 
126 | #[cfg(feature = "limit_info")]
127 | #[cfg(test)]
128 | mod tests {
129 |     use super::*;
130 |     use rocket::{
131 |         get,
132 |         http::{Header, Status},
133 |         local::blocking::Client,
134 |         routes, Build, Rocket,
135 |     };
136 | 
137 |     struct RateLimitGuard;
138 | 
139 |     impl<'r> RocketGovernable<'r> for RateLimitGuard {
140 |         fn quota(_method: Method, _route_name: &str) -> Quota {
141 |             Quota::per_second(Self::nonzero(1u32))
142 |         }
143 |     }
144 | 
145 |     #[get("/")]
146 |     fn route_test() -> Status {
147 |         Status::Ok
148 |     }
149 | 
150 |     fn launch_rocket() -> Rocket<Build> {
151 |         rocket::build().mount("/", routes![route_test])
152 |     }
153 | 
154 |     #[test]
155 |     fn test_limit_info_allow() {
156 |         let client = Client::untracked(launch_rocket()).expect("no rocket instance");
157 |         let mut req = client.get("/");
158 |         req.add_header(Header::new("X-Real-IP", "127.2.1.1"));
159 |         // req.dispatch();
160 |         let request = req.inner_mut();
161 |         let state = request
162 |             .local_cache(|| ReqState::new(Quota::per_second(NonZeroU32::new(1).unwrap()), 2));
163 |         let _ = request.real_ip();
164 | 
165 |         assert!(!RateLimitGuard::limit_info_allow(None, None, state));
166 | 
167 |         let mut req = client.get("/");
168 |         req.add_header(Header::new("X-Real-IP", "127.2.1.2"));
169 |         // req.dispatch();
170 |         let request = req.inner_mut();
171 |         let state = request
172 |             .local_cache(|| ReqState::new(Quota::per_second(NonZeroU32::new(1).unwrap()), 1));
173 |         let _ = request.real_ip();
174 | 
175 |         assert!(RateLimitGuard::limit_info_allow(None, None, state));
176 | 
177 |         let mut req = client.get("/");
178 |         req.add_header(Header::new("X-Real-IP", "127.2.1.3"));
179 |         // req.dispatch();
180 |         let request = req.inner_mut();
181 |         let state = request
182 |             .local_cache(|| ReqState::new(Quota::per_second(NonZeroU32::new(1).unwrap()), 0));
183 |         let _ = request.real_ip();
184 | 
185 |         assert!(RateLimitGuard::limit_info_allow(None, None, state));
186 |     }
187 | }
188 | 


--------------------------------------------------------------------------------
/tests/tests.rs:
--------------------------------------------------------------------------------
  1 | #![deny(unsafe_code)]
  2 | #![deny(warnings)]
  3 | #![deny(clippy::all)]
  4 | 
  5 | use rocket::{
  6 |     catchers, get,
  7 |     http::{Accept, ContentType, Header, Status},
  8 |     launch,
  9 |     local::blocking::Client,
 10 |     routes,
 11 | };
 12 | use rocket_governor::header as rg_header;
 13 | use rocket_governor::{rocket_governor_catcher, Method, Quota, RocketGovernable, RocketGovernor};
 14 | use std::{str::FromStr, thread, time::Duration};
 15 | 
 16 | pub struct RateLimitGuard;
 17 | 
 18 | impl<'r> RocketGovernable<'r> for RateLimitGuard {
 19 |     fn quota(_method: Method, _route_name: &str) -> Quota {
 20 |         Quota::per_second(Self::nonzero(1u32))
 21 |     }
 22 | }
 23 | 
 24 | pub struct RateLimitGuardWithMember {
 25 |     pub member: u8,
 26 | }
 27 | 
 28 | impl<'r> RocketGovernable<'r> for RateLimitGuardWithMember {
 29 |     fn quota(_method: Method, _route_name: &str) -> Quota {
 30 |         Quota::with_period(Duration::from_secs(2u64)).unwrap()
 31 |     }
 32 | }
 33 | 
 34 | pub struct RateLimitGGuard;
 35 | impl<'r> RocketGovernable<'r> for RateLimitGGuard {
 36 |     fn quota(_method: Method, _route_name: &str) -> Quota {
 37 |         Quota::per_second(Self::nonzero(1u32))
 38 |     }
 39 | }
 40 | 
 41 | #[get("/")]
 42 | fn route_test(_limitguard: RocketGovernor<RateLimitGuard>) -> Status {
 43 |     Status::Ok
 44 | }
 45 | 
 46 | #[get("/member")]
 47 | fn route_member(_limitguard: RocketGovernor<RateLimitGuardWithMember>) -> Status {
 48 |     Status::Ok
 49 | }
 50 | 
 51 | mod guard2 {
 52 |     use rocket::{get, http::Status};
 53 |     use rocket_governor::{Method, Quota, RocketGovernable, RocketGovernor};
 54 | 
 55 |     pub struct RateLimitGuard;
 56 | 
 57 |     impl<'r> RocketGovernable<'r> for RateLimitGuard {
 58 |         fn quota(_method: Method, route_name: &str) -> Quota {
 59 |             match route_name {
 60 |                 "route_hour" => Quota::per_hour(Self::nonzero(1)),
 61 |                 "route_multi" => Quota::per_hour(Self::nonzero(4)),
 62 |                 _ => Quota::per_second(Self::nonzero(1u32)),
 63 |             }
 64 |         }
 65 |     }
 66 | 
 67 |     #[get("/")]
 68 |     pub fn route_test(_limitguard: RocketGovernor<RateLimitGuard>) -> Status {
 69 |         Status::Ok
 70 |     }
 71 | 
 72 |     #[get("/hour")]
 73 |     pub fn route_hour(_limitguard: RocketGovernor<RateLimitGuard>) -> Status {
 74 |         Status::Ok
 75 |     }
 76 | 
 77 |     #[get("/multi")]
 78 |     pub fn route_multi(_limitguard: RocketGovernor<RateLimitGuard>) -> Status {
 79 |         Status::Ok
 80 |     }
 81 | }
 82 | 
 83 | #[launch]
 84 | fn launch_rocket() -> _ {
 85 |     #[allow(unused_mut)] // attach fairing only on feature limit_info
 86 |     let mut r = rocket::build()
 87 |         .mount("/", routes![route_test, route_member])
 88 |         .register("/", catchers!(rocket_governor_catcher))
 89 |         .mount(
 90 |             "/guard2",
 91 |             routes![guard2::route_test, guard2::route_hour, guard2::route_multi],
 92 |         )
 93 |         .register("/guard2", catchers!(rocket_governor_catcher));
 94 | 
 95 |     #[cfg(feature = "limit_info")]
 96 |     {
 97 |         r = r.attach(rocket_governor::LimitHeaderGen);
 98 |     }
 99 | 
100 |     r
101 | }
102 | 
103 | #[test]
104 | fn test_ratelimit() {
105 |     let client = Client::untracked(launch_rocket()).expect("no rocket instance");
106 |     let mut req = client.get("/");
107 |     req.add_header(Header::new("X-Real-IP", "127.0.1.1"));
108 |     let res = req.dispatch();
109 | 
110 |     assert_eq!(Status::Ok, res.status());
111 | 
112 |     let mut req = client.get("/");
113 |     req.add_header(Header::new("X-Real-IP", "127.0.1.1"));
114 |     let res = req.dispatch();
115 | 
116 |     assert_eq!(Status::TooManyRequests, res.status());
117 | 
118 |     let mut req = client.get("/");
119 |     req.add_header(Header::new("X-Real-IP", "127.0.1.1"));
120 |     let res = req.dispatch();
121 | 
122 |     assert_eq!(Status::TooManyRequests, res.status());
123 | 
124 |     let mut req = client.get("/");
125 |     req.add_header(Header::new("X-Real-IP", "127.0.1.2"));
126 |     let res = req.dispatch();
127 | 
128 |     assert_eq!(Status::Ok, res.status());
129 | 
130 |     thread::sleep(Duration::from_millis(1100u64));
131 | 
132 |     let mut req = client.get("/");
133 |     req.add_header(Header::new("X-Real-IP", "127.0.1.1"));
134 |     let res = req.dispatch();
135 | 
136 |     assert_eq!(Status::Ok, res.status());
137 | 
138 |     thread::sleep(Duration::from_millis(300u64));
139 | 
140 |     let mut req = client.get("/");
141 |     req.add_header(Header::new("X-Real-IP", "127.0.1.1"));
142 |     let res = req.dispatch();
143 | 
144 |     assert_eq!(Status::TooManyRequests, res.status());
145 | }
146 | 
147 | #[test]
148 | fn test_ratelimit_with_member() {
149 |     let client = Client::untracked(launch_rocket()).expect("no rocket instance");
150 |     let mut req = client.get("/member");
151 |     req.add_header(Header::new("X-Real-IP", "127.0.2.1"));
152 |     let res = req.dispatch();
153 | 
154 |     assert_eq!(Status::Ok, res.status());
155 | 
156 |     let mut req = client.get("/member");
157 |     req.add_header(Header::new("X-Real-IP", "127.0.2.2"));
158 |     let res = req.dispatch();
159 | 
160 |     assert_eq!(Status::Ok, res.status());
161 | 
162 |     let mut req = client.get("/member");
163 |     req.add_header(Header::new("X-Real-IP", "127.0.2.1"));
164 |     let res = req.dispatch();
165 | 
166 |     assert_eq!(Status::TooManyRequests, res.status());
167 | 
168 |     thread::sleep(Duration::from_millis(1100u64));
169 | 
170 |     let mut req = client.get("/member");
171 |     req.add_header(Header::new("X-Real-IP", "127.0.2.1"));
172 |     let res = req.dispatch();
173 | 
174 |     assert_eq!(Status::TooManyRequests, res.status());
175 | 
176 |     thread::sleep(Duration::from_millis(900u64));
177 | 
178 |     let mut req = client.get("/member");
179 |     req.add_header(Header::new("X-Real-IP", "127.0.2.1"));
180 |     let res = req.dispatch();
181 | 
182 |     assert_eq!(Status::Ok, res.status());
183 | }
184 | 
185 | #[test]
186 | fn test_ratelimit_header() {
187 |     let client = Client::untracked(launch_rocket()).expect("no rocket instance");
188 |     let mut req = client.get("/");
189 |     req.add_header(Header::new("X-Real-IP", "127.0.3.1"));
190 |     let res = req.dispatch();
191 | 
192 |     assert_eq!(Status::Ok, res.status());
193 | 
194 |     let mut req = client.get("/");
195 |     req.add_header(Header::new("X-Real-IP", "127.0.3.1"));
196 |     let res = req.dispatch();
197 | 
198 |     assert_eq!(Status::TooManyRequests, res.status());
199 | 
200 |     let ratelimit_header = res.headers().get_one(rg_header::X_RATELIMIT_ERROR);
201 |     assert_eq!(None, ratelimit_header);
202 | 
203 |     let reset_header = res.headers().get_one(rg_header::X_RATELIMIT_RESET);
204 |     assert_ne!(None, reset_header);
205 |     let reset_header = reset_header.unwrap();
206 |     assert!(!reset_header.is_empty());
207 |     u64::from_str(reset_header).unwrap();
208 | 
209 |     let mut req = client.get("/guard2/hour");
210 |     req.add_header(Header::new("X-Real-IP", "127.0.3.2"));
211 |     req.dispatch();
212 |     let mut req = client.get("/guard2/hour");
213 |     req.add_header(Header::new("X-Real-IP", "127.0.3.2"));
214 |     let res = req.dispatch();
215 | 
216 |     assert_eq!(Status::TooManyRequests, res.status());
217 | 
218 |     let retry_header = res.headers().get_one(rg_header::RETRY_AFTER);
219 |     assert_ne!(None, retry_header);
220 |     let retry_header = retry_header.unwrap();
221 |     assert!(!retry_header.is_empty());
222 |     assert!(u64::from_str(retry_header).unwrap() > 59 * 60);
223 | }
224 | 
225 | #[test]
226 | fn test_ratelimit_body() {
227 |     let client = Client::untracked(launch_rocket()).expect("no rocket instance");
228 |     let mut req = client.get("/");
229 |     req.add_header(Header::new("X-Real-IP", "127.0.4.1"));
230 |     let res = req.dispatch();
231 | 
232 |     assert_eq!(Status::Ok, res.status());
233 | 
234 |     let mut req = client.get("/");
235 |     req.add_header(Header::new("X-Real-IP", "127.0.4.1"));
236 |     let res = req.dispatch();
237 | 
238 |     assert_eq!(Status::TooManyRequests, res.status());
239 |     assert_eq!(ContentType::HTML, res.content_type().unwrap());
240 | 
241 |     let body_string = res.into_string().unwrap();
242 | 
243 |     assert!(body_string.starts_with("<!DOCTYPE html>"));
244 |     assert!(body_string.contains("429"));
245 | 
246 |     let req = client.get("/");
247 |     let res = req
248 |         .header(Accept::JSON)
249 |         .header(Header::new("X-Real-IP", "127.0.4.1"))
250 |         .dispatch();
251 | 
252 |     assert_eq!(Status::TooManyRequests, res.status());
253 |     assert_eq!(ContentType::JSON, res.content_type().unwrap());
254 | 
255 |     let body_string = res.into_string().unwrap();
256 | 
257 |     assert!(body_string.starts_with('{'));
258 |     assert!(body_string.contains("\"code\": 429"));
259 | }
260 | 
261 | #[test]
262 | fn test_ratelimit_guards_are_separated() {
263 |     let client = Client::untracked(launch_rocket()).expect("no rocket instance");
264 |     let mut req = client.get("/");
265 |     req.add_header(Header::new("X-Real-IP", "127.0.5.1"));
266 |     let res = req.dispatch();
267 | 
268 |     assert_eq!(Status::Ok, res.status());
269 | 
270 |     let mut req = client.get("/guard2");
271 |     req.add_header(Header::new("X-Real-IP", "127.0.5.1"));
272 |     let res = req.dispatch();
273 | 
274 |     assert_eq!(Status::Ok, res.status());
275 | 
276 |     let mut req = client.get("/");
277 |     req.add_header(Header::new("X-Real-IP", "127.0.5.1"));
278 |     let res = req.dispatch();
279 | 
280 |     assert_eq!(Status::TooManyRequests, res.status());
281 | 
282 |     let mut req = client.get("/guard2");
283 |     req.add_header(Header::new("X-Real-IP", "127.0.5.1"));
284 |     let res = req.dispatch();
285 | 
286 |     assert_eq!(Status::TooManyRequests, res.status());
287 | }
288 | 
289 | #[cfg(feature = "limit_info")]
290 | #[test]
291 | fn test_ratelimit_info_header() {
292 |     let client = Client::untracked(launch_rocket()).expect("no rocket instance");
293 |     let mut req = client.get("/guard2/multi");
294 |     req.add_header(Header::new("X-Real-IP", "127.0.6.1"));
295 |     let res = req.dispatch();
296 |     assert_eq!(Status::Ok, res.status());
297 |     let mut req = client.get("/guard2/multi");
298 |     req.add_header(Header::new("X-Real-IP", "127.0.6.1"));
299 |     let res = req.dispatch();
300 |     assert_eq!(Status::Ok, res.status());
301 | 
302 |     let retry_header = res.headers().get_one(rg_header::RETRY_AFTER);
303 |     assert_eq!(None, retry_header);
304 | 
305 |     let remain_header = res.headers().get_one(rg_header::X_RATELIMIT_REMAINING);
306 |     assert_eq!(None, remain_header); // only set when <= 1 remains
307 | 
308 |     let mut req = client.get("/guard2/multi");
309 |     req.add_header(Header::new("X-Real-IP", "127.0.6.1"));
310 |     let res = req.dispatch();
311 |     assert_eq!(Status::Ok, res.status());
312 | 
313 |     let remain_header = res.headers().get_one(rg_header::X_RATELIMIT_REMAINING);
314 |     assert_ne!(None, remain_header); // remains == 1
315 |     let remain_header = remain_header.unwrap();
316 |     assert!(!remain_header.is_empty());
317 |     assert_eq!(u32::from_str(remain_header).unwrap(), 1);
318 | 
319 |     let mut req = client.get("/guard2/multi");
320 |     req.add_header(Header::new("X-Real-IP", "127.0.6.1"));
321 |     let res = req.dispatch();
322 |     assert_eq!(Status::Ok, res.status());
323 | 
324 |     let remain_header = res.headers().get_one(rg_header::X_RATELIMIT_REMAINING);
325 |     assert_ne!(None, remain_header); // remains == 0
326 |     let remain_header = remain_header.unwrap();
327 |     assert!(!remain_header.is_empty());
328 |     assert_eq!(u32::from_str(remain_header).unwrap(), 0);
329 | 
330 |     let mut req = client.get("/guard2/multi");
331 |     req.add_header(Header::new("X-Real-IP", "127.0.6.1"));
332 |     let res = req.dispatch();
333 |     assert_eq!(Status::TooManyRequests, res.status());
334 | 
335 |     let limit_header = res.headers().get_one(rg_header::X_RATELIMIT_LIMIT);
336 |     assert_eq!(Some("4"), limit_header);
337 |     let remain_header = res.headers().get_one(rg_header::X_RATELIMIT_REMAINING);
338 |     assert_eq!(None, remain_header); // only set on not limited
339 |     let retry_header = res.headers().get_one(rg_header::RETRY_AFTER);
340 |     assert_ne!(None, retry_header);
341 | }
342 | 


--------------------------------------------------------------------------------