1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

extern crate rmp_serde;

use std::ops::{Deref, DerefMut};
use std::io::{Cursor, Read};

use rocket::outcome::{Outcome, IntoOutcome};
use rocket::request::Request;
use rocket::data::{self, Data, FromData};
use rocket::response::{self, Responder, Response};
use rocket::http::{ContentType, Status};

use serde::Serialize;
use serde::de::DeserializeOwned;

pub use self::rmp_serde::decode::Error as MsgPackError;

/// The `MsgPack` type: implements `FromData` and `Responder`, allowing you to
/// easily consume and respond with MessagePack data.
///
/// ## Receiving MessagePack
///
/// If you're receiving MessagePack data, simply add a `data` parameter to your
/// route arguments and ensure the type of the parameter is a `MsgPack<T>`,
/// where `T` is some type you'd like to parse from MessagePack. `T` must
/// implement `Deserialize` or `DeserializeOwned` from
/// [Serde](https://github.com/serde-rs/serde). The data is parsed from the HTTP
/// request body.
///
/// ```rust,ignore
/// #[post("/users/", format = "application/msgpack", data = "<user>")]
/// fn new_user(user: MsgPack<User>) {
///     ...
/// }
/// ```
///
/// You don't _need_ to use `format = "application/msgpack"`, but it _may_ be
/// what you want. Using `format = application/msgpack` means that any request
/// that doesn't specify "application/msgpack" as its first `Content-Type:`
/// header parameter will not be routed to this handler. By default, Rocket will
/// accept a Content Type of any of the following for MessagePack data:
/// `application/msgpack`, `application/x-msgpack`, `bin/msgpack`, or
/// `bin/x-msgpack`.
///
/// ## Sending MessagePack
///
/// If you're responding with MessagePack data, return a `MsgPack<T>` type,
/// where `T` implements `Serialize` from
/// [Serde](https://github.com/serde-rs/serde). The content type of the response
/// is set to `application/msgpack` automatically.
///
/// ```rust,ignore
/// #[get("/users/<id>")]
/// fn user(id: usize) -> MsgPack<User> {
///     let user_from_id = User::from(id);
///     ...
///     MsgPack(user_from_id)
/// }
/// ```
///
/// ## Incoming Data Limits
///
/// The default size limit for incoming MessagePack data is 1MiB. Setting a
/// limit protects your application from denial of service (DOS) attacks and
/// from resource exhaustion through high memory consumption. The limit can be
/// increased by setting the `limits.msgpack` configuration parameter. For
/// instance, to increase the MessagePack limit to 5MiB for all environments,
/// you may add the following to your `Rocket.toml`:
///
/// ```toml
/// [global.limits]
/// msgpack = 5242880
/// ```
#[derive(Debug)]
pub struct MsgPack<T>(pub T);

impl<T> MsgPack<T> {
    /// Consumes the `MsgPack` wrapper and returns the wrapped item.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use rocket_contrib::MsgPack;
    /// let string = "Hello".to_string();
    /// let my_msgpack = MsgPack(string);
    /// assert_eq!(my_msgpack.into_inner(), "Hello".to_string());
    /// ```
    #[inline(always)]
    pub fn into_inner(self) -> T {
        self.0
    }
}

/// Default limit for MessagePack is 1MB.
const LIMIT: u64 = 1 << 20;

/// Accepted content types are: `application/msgpack`, `application/x-msgpack`,
/// `bin/msgpack`, and `bin/x-msgpack`.
#[inline(always)]
fn is_msgpack_content_type(ct: &ContentType) -> bool {
    (ct.top() == "application" || ct.top() == "bin")
        && (ct.sub() == "msgpack" || ct.sub() == "x-msgpack")
}

impl<T: DeserializeOwned> FromData for MsgPack<T> {
    type Error = MsgPackError;

    fn from_data(request: &Request, data: Data) -> data::Outcome<Self, Self::Error> {
        if !request.content_type().map_or(false, |ct| is_msgpack_content_type(&ct)) {
            error_!("Content-Type is not MessagePack.");
            return Outcome::Forward(data);
        }

        let mut buf = Vec::new();
        let size_limit = request.limits().get("msgpack").unwrap_or(LIMIT);
        if let Err(e) = data.open().take(size_limit).read_to_end(&mut buf) {
            let e = MsgPackError::InvalidDataRead(e);
            error_!("Couldn't read request data: {:?}", e);
            return Outcome::Failure((Status::BadRequest, e));
        };

        rmp_serde::from_slice(&buf).map(|val| MsgPack(val))
            .map_err(|e| { error_!("Couldn't parse MessagePack body: {:?}", e); e })
            .into_outcome(Status::BadRequest)
    }
}

/// Serializes the wrapped value into MessagePack. Returns a response with
/// Content-Type `MsgPack` and a fixed-size body with the serialization. If
/// serialization fails, an `Err` of `Status::InternalServerError` is returned.
impl<T: Serialize> Responder<'static> for MsgPack<T> {
    fn respond_to(self, _: &Request) -> response::Result<'static> {
        rmp_serde::to_vec(&self.0).map_err(|e| {
            error_!("MsgPack failed to serialize: {:?}", e);
            Status::InternalServerError
        }).and_then(|buf| {
            Response::build()
                .sized_body(Cursor::new(buf))
                .ok()
        })
    }
}

impl<T> Deref for MsgPack<T> {
    type Target = T;

    #[inline(always)]
    fn deref<'a>(&'a self) -> &'a T {
        &self.0
    }
}

impl<T> DerefMut for MsgPack<T> {
    #[inline(always)]
    fn deref_mut<'a>(&'a mut self) -> &'a mut T {
        &mut self.0
    }
}