Comunic/comunicapiv3
1
0
Fork 0
mirror of https://gitlab.com/comunic/comunicapiv3 synced 2024-11-29 16:56:28 +00:00
Code Issues Projects Releases Wiki Activity

Apply stream read limit

Browse Source
This commit is contained in:
Pierre HUBERT 2020-06-20 10:47:06 +02:00
parent 32b6c95e9f
commit bab1fe8272
2 changed files with 18 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/constants.rs
View File

@ -40,4 +40,4 @@ pub const DEFAULT_ACCOUNT_IMAGE: &str = "avatars/0Reverse.png";
pub const ERROR_ACCOUNT_IMAGE: &str = "avatars/0Red.png"; pub const ERROR_ACCOUNT_IMAGE: &str = "avatars/0Red.png";
/// Maximum requests size (50 Mo) /// Maximum requests size (50 Mo)
pub const MAX_REQUEST_SIZE: u64 = 50000000; pub const MAX_REQUEST_SIZE: usize = 50000000;

23
src/controllers/server.rs
View File

@ -1,5 +1,4 @@
use std::collections::HashMap; use std::collections::HashMap;
use std::io::Take;
use std::pin::Pin; use std::pin::Pin;
use actix_web::{App, FromRequest, http, HttpMessage, HttpRequest, HttpResponse, HttpServer, web}; use actix_web::{App, FromRequest, http, HttpMessage, HttpRequest, HttpResponse, HttpServer, web};
@ -23,9 +22,10 @@ use crate::data::http_request_handler::{HttpRequestHandler, RequestValue};
/// ///
/// @author Pierre Hubert /// @author Pierre Hubert
/// Custom stream to give it a limit /// Custom stream to give a limit to requests size
struct LimitedStream { struct LimitedStream {
stream: Box<dyn Stream<Item=Result<Bytes, PayloadError>> + Unpin + 'static> stream: Box<dyn Stream<Item=Result<Bytes, PayloadError>> + Unpin + 'static>,
already_read: usize,
} }
impl<'a> Stream for LimitedStream impl<'a> Stream for LimitedStream
@ -33,7 +33,17 @@ impl<'a> Stream for LimitedStream
type Item = Result<Bytes, PayloadError>; type Item = Result<Bytes, PayloadError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Pin::new(self.stream.as_mut()).poll_next(cx) if self.already_read >= MAX_REQUEST_SIZE { // TODO : check if works
return Poll::Ready(None);
}
let res = Pin::new(self.stream.as_mut()).poll_next(cx);
if let Poll::Ready(Some(Ok(d))) = &res {
self.already_read = self.already_read + d.len();
}
res
} }
} }
@ -54,7 +64,8 @@ impl FromRequest for CustomRequest {
let payload = Box::new(payload.take()); let payload = Box::new(payload.take());
let payload = LimitedStream { let payload = LimitedStream {
stream: payload stream: payload,
already_read: 0,
}; };
@ -63,7 +74,7 @@ impl FromRequest for CustomRequest {
// Check the size, if provided // Check the size, if provided
if req.headers().contains_key("Content-Length") { if req.headers().contains_key("Content-Length") {
if let Some(v) = req.headers().get("Content-Length") { if let Some(v) = req.headers().get("Content-Length") {
if String::from_utf8_lossy(v.as_bytes()).parse::<u64>().unwrap_or(0) > MAX_REQUEST_SIZE { if String::from_utf8_lossy(v.as_bytes()).parse::<usize>().unwrap_or(0) > MAX_REQUEST_SIZE {
return Err(actix_web::error::ErrorBadRequest("Request too big!")); return Err(actix_web::error::ErrorBadRequest("Request too big!"));
} }
} }