actix_protobuf/
lib.rs

1//! Protobuf payload extractor for Actix Web.
2
3#![forbid(unsafe_code)]
4#![doc(html_logo_url = "https://actix.rs/img/logo.png")]
5#![doc(html_favicon_url = "https://actix.rs/favicon.ico")]
6#![cfg_attr(docsrs, feature(doc_cfg))]
7
8use std::{
9    fmt,
10    future::Future,
11    ops::{Deref, DerefMut},
12    pin::Pin,
13    task::{self, Poll},
14};
15
16use actix_web::{
17    body::BoxBody,
18    dev::Payload,
19    error::PayloadError,
20    http::header::{CONTENT_LENGTH, CONTENT_TYPE},
21    web::BytesMut,
22    Error, FromRequest, HttpMessage, HttpRequest, HttpResponse, HttpResponseBuilder, Responder,
23    ResponseError,
24};
25use derive_more::derive::Display;
26use futures_util::{
27    future::{FutureExt as _, LocalBoxFuture},
28    stream::StreamExt as _,
29};
30use prost::{DecodeError as ProtoBufDecodeError, EncodeError as ProtoBufEncodeError, Message};
31
32#[derive(Debug, Display)]
33pub enum ProtoBufPayloadError {
34    /// Payload size is bigger than 256k
35    #[display("Payload size is bigger than 256k")]
36    Overflow,
37
38    /// Content type error
39    #[display("Content type error")]
40    ContentType,
41
42    /// Serialize error
43    #[display("ProtoBuf serialize error: {_0}")]
44    Serialize(ProtoBufEncodeError),
45
46    /// Deserialize error
47    #[display("ProtoBuf deserialize error: {_0}")]
48    Deserialize(ProtoBufDecodeError),
49
50    /// Payload error
51    #[display("Error that occur during reading payload: {_0}")]
52    Payload(PayloadError),
53}
54
55// TODO: impl error for ProtoBufPayloadError
56
57impl ResponseError for ProtoBufPayloadError {
58    fn error_response(&self) -> HttpResponse {
59        match *self {
60            ProtoBufPayloadError::Overflow => HttpResponse::PayloadTooLarge().into(),
61            _ => HttpResponse::BadRequest().into(),
62        }
63    }
64}
65
66impl From<PayloadError> for ProtoBufPayloadError {
67    fn from(err: PayloadError) -> ProtoBufPayloadError {
68        ProtoBufPayloadError::Payload(err)
69    }
70}
71
72impl From<ProtoBufDecodeError> for ProtoBufPayloadError {
73    fn from(err: ProtoBufDecodeError) -> ProtoBufPayloadError {
74        ProtoBufPayloadError::Deserialize(err)
75    }
76}
77
78pub struct ProtoBuf<T: Message>(pub T);
79
80impl<T: Message> Deref for ProtoBuf<T> {
81    type Target = T;
82
83    fn deref(&self) -> &T {
84        &self.0
85    }
86}
87
88impl<T: Message> DerefMut for ProtoBuf<T> {
89    fn deref_mut(&mut self) -> &mut T {
90        &mut self.0
91    }
92}
93
94impl<T: Message> fmt::Debug for ProtoBuf<T>
95where
96    T: fmt::Debug,
97{
98    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
99        write!(f, "ProtoBuf: {:?}", self.0)
100    }
101}
102
103impl<T: Message> fmt::Display for ProtoBuf<T>
104where
105    T: fmt::Display,
106{
107    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
108        fmt::Display::fmt(&self.0, f)
109    }
110}
111
112pub struct ProtoBufConfig {
113    limit: usize,
114}
115
116impl ProtoBufConfig {
117    /// Change max size of payload. By default max size is 256Kb
118    pub fn limit(&mut self, limit: usize) -> &mut Self {
119        self.limit = limit;
120        self
121    }
122}
123
124impl Default for ProtoBufConfig {
125    fn default() -> Self {
126        ProtoBufConfig { limit: 262_144 }
127    }
128}
129
130impl<T> FromRequest for ProtoBuf<T>
131where
132    T: Message + Default + 'static,
133{
134    type Error = Error;
135    type Future = LocalBoxFuture<'static, Result<Self, Error>>;
136
137    #[inline]
138    fn from_request(req: &HttpRequest, payload: &mut Payload) -> Self::Future {
139        let limit = req
140            .app_data::<ProtoBufConfig>()
141            .map(|c| c.limit)
142            .unwrap_or(262_144);
143        ProtoBufMessage::new(req, payload)
144            .limit(limit)
145            .map(move |res| match res {
146                Ok(item) => Ok(ProtoBuf(item)),
147                Err(err) => Err(err.into()),
148            })
149            .boxed_local()
150    }
151}
152
153impl<T: Message + Default> Responder for ProtoBuf<T> {
154    type Body = BoxBody;
155
156    fn respond_to(self, _: &HttpRequest) -> HttpResponse {
157        let mut buf = Vec::new();
158        match self.0.encode(&mut buf) {
159            Ok(()) => HttpResponse::Ok()
160                .content_type("application/protobuf")
161                .body(buf),
162            Err(err) => HttpResponse::from_error(Error::from(ProtoBufPayloadError::Serialize(err))),
163        }
164    }
165}
166
167pub struct ProtoBufMessage<T: Message + Default> {
168    limit: usize,
169    length: Option<usize>,
170    stream: Option<Payload>,
171    err: Option<ProtoBufPayloadError>,
172    fut: Option<LocalBoxFuture<'static, Result<T, ProtoBufPayloadError>>>,
173}
174
175impl<T: Message + Default> ProtoBufMessage<T> {
176    /// Create `ProtoBufMessage` for request.
177    pub fn new(req: &HttpRequest, payload: &mut Payload) -> Self {
178        if req.content_type() != "application/protobuf"
179            && req.content_type() != "application/x-protobuf"
180        {
181            return ProtoBufMessage {
182                limit: 262_144,
183                length: None,
184                stream: None,
185                fut: None,
186                err: Some(ProtoBufPayloadError::ContentType),
187            };
188        }
189
190        let mut len = None;
191        if let Some(l) = req.headers().get(CONTENT_LENGTH) {
192            if let Ok(s) = l.to_str() {
193                if let Ok(l) = s.parse::<usize>() {
194                    len = Some(l)
195                }
196            }
197        }
198
199        ProtoBufMessage {
200            limit: 262_144,
201            length: len,
202            stream: Some(payload.take()),
203            fut: None,
204            err: None,
205        }
206    }
207
208    /// Change max size of payload. By default max size is 256Kb
209    pub fn limit(mut self, limit: usize) -> Self {
210        self.limit = limit;
211        self
212    }
213}
214
215impl<T: Message + Default + 'static> Future for ProtoBufMessage<T> {
216    type Output = Result<T, ProtoBufPayloadError>;
217
218    fn poll(mut self: Pin<&mut Self>, task: &mut task::Context<'_>) -> Poll<Self::Output> {
219        if let Some(ref mut fut) = self.fut {
220            return Pin::new(fut).poll(task);
221        }
222
223        if let Some(err) = self.err.take() {
224            return Poll::Ready(Err(err));
225        }
226
227        let limit = self.limit;
228        if let Some(len) = self.length.take() {
229            if len > limit {
230                return Poll::Ready(Err(ProtoBufPayloadError::Overflow));
231            }
232        }
233
234        let mut stream = self
235            .stream
236            .take()
237            .expect("ProtoBufMessage could not be used second time");
238
239        self.fut = Some(
240            async move {
241                let mut body = BytesMut::with_capacity(8192);
242
243                while let Some(item) = stream.next().await {
244                    let chunk = item?;
245                    if (body.len() + chunk.len()) > limit {
246                        return Err(ProtoBufPayloadError::Overflow);
247                    } else {
248                        body.extend_from_slice(&chunk);
249                    }
250                }
251
252                Ok(<T>::decode(&mut body)?)
253            }
254            .boxed_local(),
255        );
256        self.poll(task)
257    }
258}
259
260pub trait ProtoBufResponseBuilder {
261    fn protobuf<T: Message>(&mut self, value: T) -> Result<HttpResponse, Error>;
262}
263
264impl ProtoBufResponseBuilder for HttpResponseBuilder {
265    fn protobuf<T: Message>(&mut self, value: T) -> Result<HttpResponse, Error> {
266        self.insert_header((CONTENT_TYPE, "application/protobuf"));
267
268        let mut body = Vec::new();
269        value
270            .encode(&mut body)
271            .map_err(ProtoBufPayloadError::Serialize)?;
272
273        Ok(self.body(body))
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use actix_web::{http::header, test::TestRequest};
280
281    use super::*;
282
283    impl PartialEq for ProtoBufPayloadError {
284        fn eq(&self, other: &ProtoBufPayloadError) -> bool {
285            match *self {
286                ProtoBufPayloadError::Overflow => {
287                    matches!(*other, ProtoBufPayloadError::Overflow)
288                }
289                ProtoBufPayloadError::ContentType => {
290                    matches!(*other, ProtoBufPayloadError::ContentType)
291                }
292                _ => false,
293            }
294        }
295    }
296
297    #[derive(Clone, PartialEq, Eq, Message)]
298    pub struct MyObject {
299        #[prost(int32, tag = "1")]
300        pub number: i32,
301        #[prost(string, tag = "2")]
302        pub name: String,
303    }
304
305    #[actix_web::test]
306    async fn test_protobuf() {
307        let protobuf = ProtoBuf(MyObject {
308            number: 9,
309            name: "test".to_owned(),
310        });
311        let req = TestRequest::default().to_http_request();
312        let resp = protobuf.respond_to(&req);
313        let ct = resp.headers().get(header::CONTENT_TYPE).unwrap();
314        assert_eq!(ct, "application/protobuf");
315    }
316
317    #[actix_web::test]
318    async fn test_protobuf_message() {
319        let (req, mut pl) = TestRequest::default().to_http_parts();
320        let protobuf = ProtoBufMessage::<MyObject>::new(&req, &mut pl).await;
321        assert_eq!(protobuf.err().unwrap(), ProtoBufPayloadError::ContentType);
322
323        let (req, mut pl) = TestRequest::get()
324            .insert_header((header::CONTENT_TYPE, "application/text"))
325            .to_http_parts();
326        let protobuf = ProtoBufMessage::<MyObject>::new(&req, &mut pl).await;
327        assert_eq!(protobuf.err().unwrap(), ProtoBufPayloadError::ContentType);
328
329        let (req, mut pl) = TestRequest::get()
330            .insert_header((header::CONTENT_TYPE, "application/protobuf"))
331            .insert_header((header::CONTENT_LENGTH, "10000"))
332            .to_http_parts();
333        let protobuf = ProtoBufMessage::<MyObject>::new(&req, &mut pl)
334            .limit(100)
335            .await;
336        assert_eq!(protobuf.err().unwrap(), ProtoBufPayloadError::Overflow);
337    }
338}
actix_protobuf/lib.rs

actix_protobuf/
lib.rs
