actix/actix-extras
1
0
Fork 0
mirror of https://github.com/actix/actix-extras.git synced 2025-04-11 13:56:57 +02:00
Code Issues Releases Wiki Activity
actix-extras/src/actix_limitation/lib.rs.html

368 lines
18 KiB
HTML
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `actix-limitation/src/lib.rs`."><meta name="keywords" content="rust, rustlang, rust-lang"><title>lib.rs - source</title><link rel="preload" as="font" type="font/woff2" crossorigin href="../../SourceSerif4-Regular.ttf.woff2"><link rel="preload" as="font" type="font/woff2" crossorigin href="../../FiraSans-Regular.woff2"><link rel="preload" as="font" type="font/woff2" crossorigin href="../../FiraSans-Medium.woff2"><link rel="preload" as="font" type="font/woff2" crossorigin href="../../SourceCodePro-Regular.ttf.woff2"><link rel="preload" as="font" type="font/woff2" crossorigin href="../../SourceSerif4-Bold.ttf.woff2"><link rel="preload" as="font" type="font/woff2" crossorigin href="../../SourceCodePro-Semibold.ttf.woff2"><link rel="stylesheet" href="../../normalize.css"><link rel="stylesheet" href="../../rustdoc.css" id="mainThemeStyle"><link rel="stylesheet" href="../../ayu.css" disabled><link rel="stylesheet" href="../../dark.css" disabled><link rel="stylesheet" href="../../light.css" id="themeStyle"><script id="default-settings" ></script><script src="../../storage.js"></script><script defer src="../../source-script.js"></script><script defer src="../../source-files.js"></script><script defer src="../../main.js"></script><noscript><link rel="stylesheet" href="../../noscript.css"></noscript><link rel="icon" href="https://actix.rs/favicon.ico"></head><body class="rustdoc source"><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="mobile-topbar"><button class="sidebar-menu-toggle">&#9776;</button><a class="sidebar-logo" href="../../actix_limitation/index.html"><div class="logo-container"><img src="https://actix.rs/img/logo.png" alt="logo"></div></a><h2 class="location"></h2></nav><nav class="sidebar"><a class="sidebar-logo" href="../../actix_limitation/index.html"><div class="logo-container">
<img src="https://actix.rs/img/logo.png" alt="logo"></div></a></nav><main><div class="width-limiter"><div class="sub-container"><a class="sub-logo-container" href="../../actix_limitation/index.html">
<img src="https://actix.rs/img/logo.png" alt="logo"></a><nav class="sub"><form class="search-form"><div class="search-container"><span></span><input class="search-input" name="search" autocomplete="off" spellcheck="false" placeholder="Click or press ‘S’ to search, ‘?’ for more options…" type="search"><div id="help-button" title="help" tabindex="-1"><button type="button">?</button></div><div id="settings-menu" tabindex="-1"><a href="../../settings.html" title="settings"><img width="22" height="22" alt="Change settings" src="../../wheel.svg"></a></div></div></form></nav></div><section id="main-content" class="content"><div class="example-wrap"><pre class="line-numbers"><span id="1">1</span>
<span id="2">2</span>
<span id="3">3</span>
<span id="4">4</span>
<span id="5">5</span>
<span id="6">6</span>
<span id="7">7</span>
<span id="8">8</span>
<span id="9">9</span>
<span id="10">10</span>
<span id="11">11</span>
<span id="12">12</span>
<span id="13">13</span>
<span id="14">14</span>
<span id="15">15</span>
<span id="16">16</span>
<span id="17">17</span>
<span id="18">18</span>
<span id="19">19</span>
<span id="20">20</span>
<span id="21">21</span>
<span id="22">22</span>
<span id="23">23</span>
<span id="24">24</span>
<span id="25">25</span>
<span id="26">26</span>
<span id="27">27</span>
<span id="28">28</span>
<span id="29">29</span>
<span id="30">30</span>
<span id="31">31</span>
<span id="32">32</span>
<span id="33">33</span>
<span id="34">34</span>
<span id="35">35</span>
<span id="36">36</span>
<span id="37">37</span>
<span id="38">38</span>
<span id="39">39</span>
<span id="40">40</span>
<span id="41">41</span>
<span id="42">42</span>
<span id="43">43</span>
<span id="44">44</span>
<span id="45">45</span>
<span id="46">46</span>
<span id="47">47</span>
<span id="48">48</span>
<span id="49">49</span>
<span id="50">50</span>
<span id="51">51</span>
<span id="52">52</span>
<span id="53">53</span>
<span id="54">54</span>
<span id="55">55</span>
<span id="56">56</span>
<span id="57">57</span>
<span id="58">58</span>
<span id="59">59</span>
<span id="60">60</span>
<span id="61">61</span>
<span id="62">62</span>
<span id="63">63</span>
<span id="64">64</span>
<span id="65">65</span>
<span id="66">66</span>
<span id="67">67</span>
<span id="68">68</span>
<span id="69">69</span>
<span id="70">70</span>
<span id="71">71</span>
<span id="72">72</span>
<span id="73">73</span>
<span id="74">74</span>
<span id="75">75</span>
<span id="76">76</span>
<span id="77">77</span>
<span id="78">78</span>
<span id="79">79</span>
<span id="80">80</span>
<span id="81">81</span>
<span id="82">82</span>
<span id="83">83</span>
<span id="84">84</span>
<span id="85">85</span>
<span id="86">86</span>
<span id="87">87</span>
<span id="88">88</span>
<span id="89">89</span>
<span id="90">90</span>
<span id="91">91</span>
<span id="92">92</span>
<span id="93">93</span>
<span id="94">94</span>
<span id="95">95</span>
<span id="96">96</span>
<span id="97">97</span>
<span id="98">98</span>
<span id="99">99</span>
<span id="100">100</span>
<span id="101">101</span>
<span id="102">102</span>
<span id="103">103</span>
<span id="104">104</span>
<span id="105">105</span>
<span id="106">106</span>
<span id="107">107</span>
<span id="108">108</span>
<span id="109">109</span>
<span id="110">110</span>
<span id="111">111</span>
<span id="112">112</span>
<span id="113">113</span>
<span id="114">114</span>
<span id="115">115</span>
<span id="116">116</span>
<span id="117">117</span>
<span id="118">118</span>
<span id="119">119</span>
<span id="120">120</span>
<span id="121">121</span>
<span id="122">122</span>
<span id="123">123</span>
<span id="124">124</span>
<span id="125">125</span>
<span id="126">126</span>
<span id="127">127</span>
<span id="128">128</span>
<span id="129">129</span>
<span id="130">130</span>
<span id="131">131</span>
<span id="132">132</span>
<span id="133">133</span>
<span id="134">134</span>
<span id="135">135</span>
<span id="136">136</span>
<span id="137">137</span>
<span id="138">138</span>
<span id="139">139</span>
<span id="140">140</span>
<span id="141">141</span>
<span id="142">142</span>
<span id="143">143</span>
<span id="144">144</span>
<span id="145">145</span>
<span id="146">146</span>
<span id="147">147</span>
<span id="148">148</span>
<span id="149">149</span>
<span id="150">150</span>
<span id="151">151</span>
<span id="152">152</span>
<span id="153">153</span>
<span id="154">154</span>
<span id="155">155</span>
<span id="156">156</span>
<span id="157">157</span>
<span id="158">158</span>
<span id="159">159</span>
<span id="160">160</span>
<span id="161">161</span>
<span id="162">162</span>
<span id="163">163</span>
<span id="164">164</span>
<span id="165">165</span>
<span id="166">166</span>
<span id="167">167</span>
<span id="168">168</span>
<span id="169">169</span>
<span id="170">170</span>
<span id="171">171</span>
<span id="172">172</span>
<span id="173">173</span>
<span id="174">174</span>
<span id="175">175</span>
<span id="176">176</span>
<span id="177">177</span>
<span id="178">178</span>
<span id="179">179</span>
<span id="180">180</span>
<span id="181">181</span>
<span id="182">182</span>
</pre><pre class="rust"><code><span class="doccomment">//! Rate limiter using a fixed window counter for arbitrary keys, backed by Redis for Actix Web.
//!
//! ```toml
//! [dependencies]
//! actix-web = &quot;4&quot;
</span><span class="attribute">#![doc = <span class="macro">concat!</span>(<span class="string">&quot;actix-limitation = \&quot;&quot;</span>, <span class="macro">env!</span>(<span class="string">&quot;CARGO_PKG_VERSION_MAJOR&quot;</span>), <span class="string">&quot;.&quot;</span>, <span class="macro">env!</span>(<span class="string">&quot;CARGO_PKG_VERSION_MINOR&quot;</span>),<span class="string">&quot;\&quot;&quot;</span>)]
</span><span class="doccomment">//! ```
//!
//! ```no_run
//! use std::{sync::Arc, time::Duration};
//! use actix_web::{dev::ServiceRequest, get, web, App, HttpServer, Responder};
//! use actix_session::SessionExt as _;
//! use actix_limitation::{Limiter, RateLimiter};
//!
//! #[get(&quot;/{id}/{name}&quot;)]
//! async fn index(info: web::Path&lt;(u32, String)&gt;) -&gt; impl Responder {
//! format!(&quot;Hello {}! id:{}&quot;, info.1, info.0)
//! }
//!
//! #[actix_web::main]
//! async fn main() -&gt; std::io::Result&lt;()&gt; {
//! let limiter = web::Data::new(
//! Limiter::builder(&quot;redis://127.0.0.1&quot;)
//! .key_by(|req: &amp;ServiceRequest| {
//! req.get_session()
//! .get(&amp;&quot;session-id&quot;)
//! .unwrap_or_else(|_| req.cookie(&amp;&quot;rate-api-id&quot;).map(|c| c.to_string()))
//! })
//! .limit(5000)
//! .period(Duration::from_secs(3600)) // 60 minutes
//! .build()
//! .unwrap(),
//! );
//!
//! HttpServer::new(move || {
//! App::new()
//! .wrap(RateLimiter::default())
//! .app_data(limiter.clone())
//! .service(index)
//! })
//! .bind((&quot;127.0.0.1&quot;, 8080))?
//! .run()
//! .await
//! }
//! ```
</span><span class="attribute">#![forbid(unsafe_code)]
#![deny(rust_2018_idioms, nonstandard_style)]
#![warn(future_incompatible, missing_docs, missing_debug_implementations)]
#![doc(html_logo_url = <span class="string">&quot;https://actix.rs/img/logo.png&quot;</span>)]
#![doc(html_favicon_url = <span class="string">&quot;https://actix.rs/favicon.ico&quot;</span>)]
</span><span class="kw">use </span>std::{borrow::Cow, fmt, sync::Arc, time::Duration};
<span class="kw">use </span>actix_web::dev::ServiceRequest;
<span class="kw">use </span>redis::Client;
<span class="kw">mod </span>builder;
<span class="kw">mod </span>errors;
<span class="kw">mod </span>middleware;
<span class="kw">mod </span>status;
<span class="kw">pub use </span><span class="self">self</span>::builder::Builder;
<span class="kw">pub use </span><span class="self">self</span>::errors::Error;
<span class="kw">pub use </span><span class="self">self</span>::middleware::RateLimiter;
<span class="kw">pub use </span><span class="self">self</span>::status::Status;
<span class="doccomment">/// Default request limit.
</span><span class="kw">pub const </span>DEFAULT_REQUEST_LIMIT: usize = <span class="number">5000</span>;
<span class="doccomment">/// Default period (in seconds).
</span><span class="kw">pub const </span>DEFAULT_PERIOD_SECS: u64 = <span class="number">3600</span>;
<span class="doccomment">/// Default cookie name.
</span><span class="kw">pub const </span>DEFAULT_COOKIE_NAME: <span class="kw-2">&amp;</span>str = <span class="string">&quot;sid&quot;</span>;
<span class="doccomment">/// Default session key.
</span><span class="attribute">#[cfg(feature = <span class="string">&quot;session&quot;</span>)]
</span><span class="kw">pub const </span>DEFAULT_SESSION_KEY: <span class="kw-2">&amp;</span>str = <span class="string">&quot;rate-api-id&quot;</span>;
<span class="doccomment">/// Helper trait to impl Debug on GetKeyFn type
</span><span class="kw">trait </span>GetKeyFnT: Fn(<span class="kw-2">&amp;</span>ServiceRequest) -&gt; <span class="prelude-ty">Option</span>&lt;String&gt; {}
<span class="kw">impl</span>&lt;T&gt; GetKeyFnT <span class="kw">for </span>T <span class="kw">where </span>T: Fn(<span class="kw-2">&amp;</span>ServiceRequest) -&gt; <span class="prelude-ty">Option</span>&lt;String&gt; {}
<span class="doccomment">/// Get key function type with auto traits
</span><span class="kw">type </span>GetKeyFn = <span class="kw">dyn </span>GetKeyFnT + Send + Sync;
<span class="doccomment">/// Get key resolver function type
</span><span class="kw">impl </span>fmt::Debug <span class="kw">for </span>GetKeyFn {
<span class="kw">fn </span>fmt(<span class="kw-2">&amp;</span><span class="self">self</span>, f: <span class="kw-2">&amp;mut </span>fmt::Formatter&lt;<span class="lifetime">&#39;_</span>&gt;) -&gt; fmt::Result {
<span class="macro">write!</span>(f, <span class="string">&quot;GetKeyFn&quot;</span>)
}
}
<span class="doccomment">/// Wrapped Get key function Trait
</span><span class="kw">type </span>GetArcBoxKeyFn = Arc&lt;GetKeyFn&gt;;
<span class="doccomment">/// Rate limiter.
</span><span class="attribute">#[derive(Debug, Clone)]
</span><span class="kw">pub struct </span>Limiter {
client: Client,
limit: usize,
period: Duration,
get_key_fn: GetArcBoxKeyFn,
}
<span class="kw">impl </span>Limiter {
<span class="doccomment">/// Construct rate limiter builder with defaults.
///
/// See [`redis-rs` docs](https://docs.rs/redis/0.21/redis/#connection-parameters) on connection
/// parameters for how to set the Redis URL.
</span><span class="attribute">#[must_use]
</span><span class="kw">pub fn </span>builder(redis_url: <span class="kw">impl </span>Into&lt;String&gt;) -&gt; Builder {
Builder {
redis_url: redis_url.into(),
limit: DEFAULT_REQUEST_LIMIT,
period: Duration::from_secs(DEFAULT_PERIOD_SECS),
get_key_fn: <span class="prelude-val">None</span>,
cookie_name: Cow::Borrowed(DEFAULT_COOKIE_NAME),
<span class="attribute">#[cfg(feature = <span class="string">&quot;session&quot;</span>)]
</span>session_key: Cow::Borrowed(DEFAULT_SESSION_KEY),
}
}
<span class="doccomment">/// Consumes one rate limit unit, returning the status.
</span><span class="kw">pub async fn </span>count(<span class="kw-2">&amp;</span><span class="self">self</span>, key: <span class="kw">impl </span>Into&lt;String&gt;) -&gt; <span class="prelude-ty">Result</span>&lt;Status, Error&gt; {
<span class="kw">let </span>(count, reset) = <span class="self">self</span>.track(key).<span class="kw">await</span><span class="question-mark">?</span>;
<span class="kw">let </span>status = Status::new(count, <span class="self">self</span>.limit, reset);
<span class="kw">if </span>count &gt; <span class="self">self</span>.limit {
<span class="prelude-val">Err</span>(Error::LimitExceeded(status))
} <span class="kw">else </span>{
<span class="prelude-val">Ok</span>(status)
}
}
<span class="doccomment">/// Tracks the given key in a period and returns the count and TTL for the key in seconds.
</span><span class="kw">async fn </span>track(<span class="kw-2">&amp;</span><span class="self">self</span>, key: <span class="kw">impl </span>Into&lt;String&gt;) -&gt; <span class="prelude-ty">Result</span>&lt;(usize, usize), Error&gt; {
<span class="kw">let </span>key = key.into();
<span class="kw">let </span>expires = <span class="self">self</span>.period.as_secs();
<span class="kw">let </span><span class="kw-2">mut </span>connection = <span class="self">self</span>.client.get_tokio_connection().<span class="kw">await</span><span class="question-mark">?</span>;
<span class="comment">// The seed of this approach is outlined Atul R in a blog post about rate limiting using
// NodeJS and Redis. For more details, see https://blog.atulr.com/rate-limiter
</span><span class="kw">let </span><span class="kw-2">mut </span>pipe = redis::pipe();
pipe.atomic()
.cmd(<span class="string">&quot;SET&quot;</span>) <span class="comment">// Set key and value
</span>.arg(<span class="kw-2">&amp;</span>key)
.arg(<span class="number">0</span>)
.arg(<span class="string">&quot;EX&quot;</span>) <span class="comment">// Set the specified expire time, in seconds.
</span>.arg(expires)
.arg(<span class="string">&quot;NX&quot;</span>) <span class="comment">// Only set the key if it does not already exist.
</span>.ignore() <span class="comment">// --- ignore returned value of SET command ---
</span>.cmd(<span class="string">&quot;INCR&quot;</span>) <span class="comment">// Increment key
</span>.arg(<span class="kw-2">&amp;</span>key)
.cmd(<span class="string">&quot;TTL&quot;</span>) <span class="comment">// Return time-to-live of key
</span>.arg(<span class="kw-2">&amp;</span>key);
<span class="kw">let </span>(count, ttl) = pipe.query_async(<span class="kw-2">&amp;mut </span>connection).<span class="kw">await</span><span class="question-mark">?</span>;
<span class="kw">let </span>reset = Status::epoch_utc_plus(Duration::from_secs(ttl))<span class="question-mark">?</span>;
<span class="prelude-val">Ok</span>((count, reset))
}
}
<span class="attribute">#[cfg(test)]
</span><span class="kw">mod </span>tests {
<span class="kw">use super</span>::<span class="kw-2">*</span>;
<span class="attribute">#[test]
</span><span class="kw">fn </span>test_create_limiter() {
<span class="kw">let </span><span class="kw-2">mut </span>builder = Limiter::builder(<span class="string">&quot;redis://127.0.0.1:6379/1&quot;</span>);
<span class="kw">let </span>limiter = builder.build();
<span class="macro">assert!</span>(limiter.is_ok());
<span class="kw">let </span>limiter = limiter.unwrap();
<span class="macro">assert_eq!</span>(limiter.limit, <span class="number">5000</span>);
<span class="macro">assert_eq!</span>(limiter.period, Duration::from_secs(<span class="number">3600</span>));
}
}
</code></pre></div>
</section></div></main><div id="rustdoc-vars" data-root-path="../../" data-current-crate="actix_limitation" data-themes="ayu,dark,light" data-resource-suffix="" data-rustdoc-version="1.65.0-nightly (228710758 2022-09-10)" ></div></body></html>
Reference in New Issue View Git Blame Copy Permalink