limiter/
lib.rs

1extern crate iron;
2
3use iron::BeforeMiddleware;
4use iron::headers::ContentLength;
5use iron::prelude::*;
6use iron::status;
7use iron::url;
8
9use std::default::Default;
10use std::error::Error;
11use std::fmt;
12use std::io::Read;
13
14const MAX_BODY_DEFAULT: u64 = 5e6 as u64;
15const MAX_URL_DEFAULT: usize = 256;
16
17/// The error thrown when the request body size is larger than the limit set
18/// by the middleware.
19#[derive(Debug)]
20pub struct RequestTooLarge;
21
22/// `RequestLimit` configures the maximum payload size and URL length for Iron's middleware system.
23pub struct RequestLimit {
24    /// The maximum size of a payload.
25    max_body: u64,
26
27    /// The maximum size of a URL.
28    max_url_length: usize,
29}
30
31impl fmt::Display for RequestTooLarge {
32    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
33        f.write_str(self.description())
34    }
35}
36
37impl Error for RequestTooLarge {
38    fn description(&self) -> &str {
39        "Request too large"
40    }
41}
42
43impl Default for RequestLimit {
44    fn default() -> RequestLimit {
45        RequestLimit {
46            max_body: MAX_BODY_DEFAULT,
47            max_url_length: MAX_URL_DEFAULT,
48        }
49    }
50}
51
52impl RequestLimit {
53    /// set_max_body_size overrides the maximum body request size of 5 MB.
54    pub fn set_max_body_size(&mut self, max_body: u64) {
55        self.max_body = max_body;
56    }
57
58    /// set_max_url_length overrides the default maximum URL length of 256.
59    pub fn set_max_url_length(&mut self, max_url_length: usize) {
60        self.max_url_length = max_url_length;
61    }
62
63    // Set the proper error response if the request body is too big or carry on
64    // down the chain.
65    fn check_payload(&self, total: u64) -> IronResult<()> {
66        if total > self.max_body {
67            Err(IronError::new(RequestTooLarge, status::PayloadTooLarge))
68        } else {
69            Ok(())
70        }
71    }
72
73    // Ensure that the URL length doesn't exceed the maximum.
74    fn check_url_length(&self, u: iron::Url) -> bool {
75        let real_url: url::Url = u.into();
76
77        real_url.as_str().len() <= self.max_url_length
78    }
79}
80
81impl BeforeMiddleware for RequestLimit {
82    // This middleware tries to read the content length of the request first
83    // before reading the request body.
84    fn before(&self, req: &mut Request) -> IronResult<()> {
85        if !self.check_url_length(req.url.clone()) {
86            return Err(IronError::new(RequestTooLarge, status::PayloadTooLarge));
87        }
88
89        match req.headers.get::<ContentLength>() {
90            Some(l) => self.check_payload(l.0),
91            None => self.check_payload(req.body.by_ref().bytes().count() as u64),
92        }
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::RequestLimit;
99
100    static GOOGLE: &'static str = "https://google.com";
101
102    #[test]
103    fn check_ok_response() {
104        let b = RequestLimit::default();
105        assert!(b.check_payload(5).is_ok());
106    }
107
108    #[test]
109    fn check_err_response() {
110        let mut b = RequestLimit::default();
111        b.set_max_body_size(1);
112        assert!(b.check_payload(2).is_err());
113    }
114
115    #[test]
116    fn test_lengthy_url() {
117        let mut b = RequestLimit::default();
118        b.set_max_url_length(5);
119        assert_eq!(false,
120                   b.check_url_length(::iron::Url::parse(GOOGLE).unwrap()));
121    }
122
123    #[test]
124    fn test_valid_url() {
125        let b = RequestLimit::default();
126        assert!(b.check_url_length(::iron::Url::parse(GOOGLE).unwrap()));
127    }
128}
limiter/lib.rs

limiter/
lib.rs
