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

329 lines
22 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" type="text/css" href="../../normalize.css"><link rel="stylesheet" type="text/css" href="../../rustdoc.css" id="mainThemeStyle"><link rel="stylesheet" type="text/css" href="../../ayu.css" disabled><link rel="stylesheet" type="text/css" href="../../dark.css" disabled><link rel="stylesheet" type="text/css" 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"><button type="button" id="help-button" title="help">?</button><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>
</pre><pre class="rust"><code><span class="doccomment">//! Rate limiter using a fixed window counter for arbitrary keys, backed by Redis for Actix Web.</span>
<span class="doccomment">//!</span>
<span class="doccomment">//! ```toml</span>
<span class="doccomment">//! [dependencies]</span>
<span class="doccomment">//! actix-web = &quot;4&quot;</span>
<span class="doccomment">//! actix-limitation = &quot;0.1.4&quot;</span>
<span class="doccomment">//! ```</span>
<span class="doccomment">//!</span>
<span class="doccomment">//! ```no_run</span>
<span class="doccomment">//! use std::time::Duration;</span>
<span class="doccomment">//! use actix_web::{get, web, App, HttpServer, Responder};</span>
<span class="doccomment">//! use actix_limitation::{Limiter, RateLimiter};</span>
<span class="doccomment">//!</span>
<span class="doccomment">//! #[get(&quot;/{id}/{name}&quot;)]</span>
<span class="doccomment">//! async fn index(info: web::Path&lt;(u32, String)&gt;) -&gt; impl Responder {</span>
<span class="doccomment">//! format!(&quot;Hello {}! id:{}&quot;, info.1, info.0)</span>
<span class="doccomment">//! }</span>
<span class="doccomment">//!</span>
<span class="doccomment">//! #[actix_web::main]</span>
<span class="doccomment">//! async fn main() -&gt; std::io::Result&lt;()&gt; {</span>
<span class="doccomment">//! let limiter = web::Data::new(</span>
<span class="doccomment">//! Limiter::builder(&quot;redis://127.0.0.1&quot;)</span>
<span class="doccomment">//! .cookie_name(&quot;session-id&quot;.to_owned())</span>
<span class="doccomment">//! .session_key(&quot;rate-api-id&quot;.to_owned())</span>
<span class="doccomment">//! .limit(5000)</span>
<span class="doccomment">//! .period(Duration::from_secs(3600)) // 60 minutes</span>
<span class="doccomment">//! .build()</span>
<span class="doccomment">//! .unwrap(),</span>
<span class="doccomment">//! );</span>
<span class="doccomment">//!</span>
<span class="doccomment">//! HttpServer::new(move || {</span>
<span class="doccomment">//! App::new()</span>
<span class="doccomment">//! .wrap(RateLimiter)</span>
<span class="doccomment">//! .app_data(limiter.clone())</span>
<span class="doccomment">//! .service(index)</span>
<span class="doccomment">//! })</span>
<span class="doccomment">//! .bind(&quot;127.0.0.1:8080&quot;)?</span>
<span class="doccomment">//! .run()</span>
<span class="doccomment">//! .await</span>
<span class="doccomment">//! }</span>
<span class="doccomment">//! ```</span>
<span class="attribute">#![<span class="ident">forbid</span>(<span class="ident">unsafe_code</span>)]</span>
<span class="attribute">#![<span class="ident">deny</span>(<span class="ident">rust_2018_idioms</span>, <span class="ident">nonstandard_style</span>)]</span>
<span class="attribute">#![<span class="ident">warn</span>(<span class="ident">future_incompatible</span>, <span class="ident">missing_docs</span>, <span class="ident">missing_debug_implementations</span>)]</span>
<span class="attribute">#![<span class="ident">doc</span>(<span class="ident">html_logo_url</span> <span class="op">=</span> <span class="string">&quot;https://actix.rs/img/logo.png&quot;</span>)]</span>
<span class="attribute">#![<span class="ident">doc</span>(<span class="ident">html_favicon_url</span> <span class="op">=</span> <span class="string">&quot;https://actix.rs/favicon.ico&quot;</span>)]</span>
<span class="kw">use</span> <span class="ident">std</span>::{<span class="ident">borrow::Cow</span>, <span class="ident">time::Duration</span>};
<span class="kw">use</span> <span class="ident">redis::Client</span>;
<span class="kw">mod</span> <span class="ident">builder</span>;
<span class="kw">mod</span> <span class="ident">errors</span>;
<span class="kw">mod</span> <span class="ident">middleware</span>;
<span class="kw">mod</span> <span class="ident">status</span>;
<span class="kw">pub</span> <span class="kw">use</span> <span class="ident"><span class="self">self</span>::builder::Builder</span>;
<span class="kw">pub</span> <span class="kw">use</span> <span class="ident"><span class="self">self</span>::errors::Error</span>;
<span class="kw">pub</span> <span class="kw">use</span> <span class="ident"><span class="self">self</span>::middleware::RateLimiter</span>;
<span class="kw">pub</span> <span class="kw">use</span> <span class="ident"><span class="self">self</span>::status::Status</span>;
<span class="doccomment">/// Default request limit.</span>
<span class="kw">pub</span> <span class="kw">const</span> <span class="ident">DEFAULT_REQUEST_LIMIT</span>: <span class="ident">usize</span> <span class="op">=</span> <span class="number">5000</span>;
<span class="doccomment">/// Default period (in seconds).</span>
<span class="kw">pub</span> <span class="kw">const</span> <span class="ident">DEFAULT_PERIOD_SECS</span>: <span class="ident">u64</span> <span class="op">=</span> <span class="number">3600</span>;
<span class="doccomment">/// Default cookie name.</span>
<span class="kw">pub</span> <span class="kw">const</span> <span class="ident">DEFAULT_COOKIE_NAME</span>: <span class="kw-2">&amp;</span><span class="ident">str</span> <span class="op">=</span> <span class="string">&quot;sid&quot;</span>;
<span class="doccomment">/// Default session key.</span>
<span class="kw">pub</span> <span class="kw">const</span> <span class="ident">DEFAULT_SESSION_KEY</span>: <span class="kw-2">&amp;</span><span class="ident">str</span> <span class="op">=</span> <span class="string">&quot;rate-api-id&quot;</span>;
<span class="doccomment">/// Rate limiter.</span>
<span class="attribute">#[<span class="ident">derive</span>(<span class="ident">Clone</span>, <span class="ident">Debug</span>)]</span>
<span class="kw">pub</span> <span class="kw">struct</span> <span class="ident">Limiter</span> {
<span class="ident">client</span>: <span class="ident">Client</span>,
<span class="ident">limit</span>: <span class="ident">usize</span>,
<span class="ident">period</span>: <span class="ident">Duration</span>,
<span class="ident">cookie_name</span>: <span class="ident">Cow</span><span class="op">&lt;</span><span class="lifetime">&#39;static</span>, <span class="ident">str</span><span class="op">&gt;</span>,
<span class="ident">session_key</span>: <span class="ident">Cow</span><span class="op">&lt;</span><span class="lifetime">&#39;static</span>, <span class="ident">str</span><span class="op">&gt;</span>,
}
<span class="kw">impl</span> <span class="ident">Limiter</span> {
<span class="doccomment">/// Construct rate limiter builder with defaults.</span>
<span class="doccomment">///</span>
<span class="doccomment">/// See [`redis-rs` docs](https://docs.rs/redis/0.21/redis/#connection-parameters) on connection</span>
<span class="doccomment">/// parameters for how to set the Redis URL.</span>
<span class="attribute">#[<span class="ident">must_use</span>]</span>
<span class="kw">pub</span> <span class="kw">fn</span> <span class="ident">builder</span>(<span class="ident">redis_url</span>: <span class="kw-2">&amp;</span><span class="ident">str</span>) -&gt; <span class="ident">Builder</span><span class="op">&lt;</span><span class="lifetime">&#39;_</span><span class="op">&gt;</span> {
<span class="ident">Builder</span> {
<span class="ident">redis_url</span>,
<span class="ident">limit</span>: <span class="ident">DEFAULT_REQUEST_LIMIT</span>,
<span class="ident">period</span>: <span class="ident">Duration::from_secs</span>(<span class="ident">DEFAULT_PERIOD_SECS</span>),
<span class="ident">cookie_name</span>: <span class="ident">Cow::Borrowed</span>(<span class="ident">DEFAULT_COOKIE_NAME</span>),
<span class="ident">session_key</span>: <span class="ident">Cow::Borrowed</span>(<span class="ident">DEFAULT_SESSION_KEY</span>),
}
}
<span class="doccomment">/// Consumes one rate limit unit, returning the status.</span>
<span class="kw">pub</span> <span class="kw">async</span> <span class="kw">fn</span> <span class="ident">count</span>(<span class="kw-2">&amp;</span><span class="self">self</span>, <span class="ident">key</span>: <span class="kw">impl</span> <span class="ident">Into</span><span class="op">&lt;</span><span class="ident">String</span><span class="op">&gt;</span>) -&gt; <span class="prelude-ty">Result</span><span class="op">&lt;</span><span class="ident">Status</span>, <span class="ident">Error</span><span class="op">&gt;</span> {
<span class="kw">let</span> (<span class="ident">count</span>, <span class="ident">reset</span>) <span class="op">=</span> <span class="self">self</span>.<span class="ident">track</span>(<span class="ident">key</span>).<span class="kw">await</span><span class="question-mark">?</span>;
<span class="kw">let</span> <span class="ident">status</span> <span class="op">=</span> <span class="ident">Status::new</span>(<span class="ident">count</span>, <span class="self">self</span>.<span class="ident">limit</span>, <span class="ident">reset</span>);
<span class="kw">if</span> <span class="ident">count</span> <span class="op">&gt;</span> <span class="self">self</span>.<span class="ident">limit</span> {
<span class="prelude-val">Err</span>(<span class="ident">Error::LimitExceeded</span>(<span class="ident">status</span>))
} <span class="kw">else</span> {
<span class="prelude-val">Ok</span>(<span class="ident">status</span>)
}
}
<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</span> <span class="kw">fn</span> <span class="ident">track</span>(<span class="kw-2">&amp;</span><span class="self">self</span>, <span class="ident">key</span>: <span class="kw">impl</span> <span class="ident">Into</span><span class="op">&lt;</span><span class="ident">String</span><span class="op">&gt;</span>) -&gt; <span class="prelude-ty">Result</span><span class="op">&lt;</span>(<span class="ident">usize</span>, <span class="ident">usize</span>), <span class="ident">Error</span><span class="op">&gt;</span> {
<span class="kw">let</span> <span class="ident">key</span> <span class="op">=</span> <span class="ident">key</span>.<span class="ident">into</span>();
<span class="kw">let</span> <span class="ident">expires</span> <span class="op">=</span> <span class="self">self</span>.<span class="ident">period</span>.<span class="ident">as_secs</span>();
<span class="kw">let</span> <span class="kw-2">mut</span> <span class="ident">connection</span> <span class="op">=</span> <span class="self">self</span>.<span class="ident">client</span>.<span class="ident">get_tokio_connection</span>().<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</span>
<span class="comment">// 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> <span class="ident">pipe</span> <span class="op">=</span> <span class="ident">redis::pipe</span>();
<span class="ident">pipe</span>.<span class="ident">atomic</span>()
.<span class="ident">cmd</span>(<span class="string">&quot;SET&quot;</span>) <span class="comment">// Set key and value</span>
.<span class="ident">arg</span>(<span class="kw-2">&amp;</span><span class="ident">key</span>)
.<span class="ident">arg</span>(<span class="number">0</span>)
.<span class="ident">arg</span>(<span class="string">&quot;EX&quot;</span>) <span class="comment">// Set the specified expire time, in seconds.</span>
.<span class="ident">arg</span>(<span class="ident">expires</span>)
.<span class="ident">arg</span>(<span class="string">&quot;NX&quot;</span>) <span class="comment">// Only set the key if it does not already exist.</span>
.<span class="ident">ignore</span>() <span class="comment">// --- ignore returned value of SET command ---</span>
.<span class="ident">cmd</span>(<span class="string">&quot;INCR&quot;</span>) <span class="comment">// Increment key</span>
.<span class="ident">arg</span>(<span class="kw-2">&amp;</span><span class="ident">key</span>)
.<span class="ident">cmd</span>(<span class="string">&quot;TTL&quot;</span>) <span class="comment">// Return time-to-live of key</span>
.<span class="ident">arg</span>(<span class="kw-2">&amp;</span><span class="ident">key</span>);
<span class="kw">let</span> (<span class="ident">count</span>, <span class="ident">ttl</span>) <span class="op">=</span> <span class="ident">pipe</span>.<span class="ident">query_async</span>(<span class="kw-2">&amp;mut</span> <span class="ident">connection</span>).<span class="kw">await</span><span class="question-mark">?</span>;
<span class="kw">let</span> <span class="ident">reset</span> <span class="op">=</span> <span class="ident">Status::epoch_utc_plus</span>(<span class="ident">Duration::from_secs</span>(<span class="ident">ttl</span>))<span class="question-mark">?</span>;
<span class="prelude-val">Ok</span>((<span class="ident">count</span>, <span class="ident">reset</span>))
}
}
<span class="attribute">#[<span class="ident">cfg</span>(<span class="ident">test</span>)]</span>
<span class="kw">mod</span> <span class="ident">tests</span> {
<span class="kw">use</span> <span class="kw">super</span>::<span class="kw-2">*</span>;
<span class="attribute">#[<span class="ident">test</span>]</span>
<span class="kw">fn</span> <span class="ident">test_create_limiter</span>() {
<span class="kw">let</span> <span class="ident">builder</span> <span class="op">=</span> <span class="ident">Limiter::builder</span>(<span class="string">&quot;redis://127.0.0.1:6379/1&quot;</span>);
<span class="kw">let</span> <span class="ident">limiter</span> <span class="op">=</span> <span class="ident">builder</span>.<span class="ident">build</span>();
<span class="macro">assert!</span>(<span class="ident">limiter</span>.<span class="ident">is_ok</span>());
<span class="kw">let</span> <span class="ident">limiter</span> <span class="op">=</span> <span class="ident">limiter</span>.<span class="ident">unwrap</span>();
<span class="macro">assert_eq!</span>(<span class="ident">limiter</span>.<span class="ident">limit</span>, <span class="number">5000</span>);
<span class="macro">assert_eq!</span>(<span class="ident">limiter</span>.<span class="ident">period</span>, <span class="ident">Duration::from_secs</span>(<span class="number">3600</span>));
<span class="macro">assert_eq!</span>(<span class="ident">limiter</span>.<span class="ident">cookie_name</span>, <span class="ident">DEFAULT_COOKIE_NAME</span>);
<span class="macro">assert_eq!</span>(<span class="ident">limiter</span>.<span class="ident">session_key</span>, <span class="ident">DEFAULT_SESSION_KEY</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.63.0-nightly (fdca237d5 2022-06-24)" ></div>
</body></html>
Reference in New Issue View Git Blame Copy Permalink