Limit buffer size for stream parser
This commit is contained in:
parent
0f3e941031
commit
982c5c2dcb
3 changed files with 21 additions and 1 deletions
|
@ -40,6 +40,8 @@ use webmetro::{
|
||||||
|
|
||||||
use super::to_hyper_error;
|
use super::to_hyper_error;
|
||||||
|
|
||||||
|
const BUFFER_LIMIT: usize = 2 * 1024 * 1024;
|
||||||
|
|
||||||
type BodyStream = Box<Stream<Item = Chunk, Error = HyperError>>;
|
type BodyStream = Box<Stream<Item = Chunk, Error = HyperError>>;
|
||||||
|
|
||||||
struct RelayServer(Arc<Mutex<Channel>>);
|
struct RelayServer(Arc<Mutex<Channel>>);
|
||||||
|
@ -62,7 +64,7 @@ impl RelayServer {
|
||||||
where S::Error: Error + Send {
|
where S::Error: Error + Send {
|
||||||
let source = stream
|
let source = stream
|
||||||
.map_err(WebmetroError::from_err)
|
.map_err(WebmetroError::from_err)
|
||||||
.parse_ebml().chunk_webm();
|
.parse_ebml().with_buffer_limit(BUFFER_LIMIT).chunk_webm();
|
||||||
let sink = Transmitter::new(self.get_channel());
|
let sink = Transmitter::new(self.get_channel());
|
||||||
|
|
||||||
Box::new(
|
Box::new(
|
||||||
|
|
|
@ -12,6 +12,7 @@ use ebml::EbmlError;
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
pub enum WebmetroError {
|
pub enum WebmetroError {
|
||||||
|
ResourcesExceeded,
|
||||||
EbmlError(EbmlError),
|
EbmlError(EbmlError),
|
||||||
IoError(IoError),
|
IoError(IoError),
|
||||||
Unknown(Box<Error + Send>)
|
Unknown(Box<Error + Send>)
|
||||||
|
@ -30,6 +31,7 @@ impl WebmetroError {
|
||||||
impl Display for WebmetroError {
|
impl Display for WebmetroError {
|
||||||
fn fmt(&self, f: &mut Formatter) -> FmtResult {
|
fn fmt(&self, f: &mut Formatter) -> FmtResult {
|
||||||
match self {
|
match self {
|
||||||
|
&WebmetroError::ResourcesExceeded => write!(f, "resources exceeded"),
|
||||||
&WebmetroError::EbmlError(ref err) => err.fmt(f),
|
&WebmetroError::EbmlError(ref err) => err.fmt(f),
|
||||||
&WebmetroError::IoError(ref err) => err.fmt(f),
|
&WebmetroError::IoError(ref err) => err.fmt(f),
|
||||||
&WebmetroError::Unknown(ref err) => err.fmt(f),
|
&WebmetroError::Unknown(ref err) => err.fmt(f),
|
||||||
|
@ -39,6 +41,7 @@ impl Display for WebmetroError {
|
||||||
impl Error for WebmetroError {
|
impl Error for WebmetroError {
|
||||||
fn description(&self) -> &str {
|
fn description(&self) -> &str {
|
||||||
match self {
|
match self {
|
||||||
|
&WebmetroError::ResourcesExceeded => "resources exceeded",
|
||||||
&WebmetroError::EbmlError(ref err) => err.description(),
|
&WebmetroError::EbmlError(ref err) => err.description(),
|
||||||
&WebmetroError::IoError(ref err) => err.description(),
|
&WebmetroError::IoError(ref err) => err.description(),
|
||||||
&WebmetroError::Unknown(ref err) => err.description(),
|
&WebmetroError::Unknown(ref err) => err.description(),
|
||||||
|
|
|
@ -10,14 +10,23 @@ use error::WebmetroError;
|
||||||
pub struct EbmlStreamingParser<S> {
|
pub struct EbmlStreamingParser<S> {
|
||||||
stream: S,
|
stream: S,
|
||||||
buffer: BytesMut,
|
buffer: BytesMut,
|
||||||
|
buffer_size_limit: Option<usize>,
|
||||||
last_read: usize
|
last_read: usize
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl<S> EbmlStreamingParser<S> {
|
||||||
|
pub fn with_buffer_limit(mut self, limit: usize) -> Self {
|
||||||
|
self.buffer_size_limit = Some(limit);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub trait StreamEbml where Self: Sized + Stream, Self::Item: AsRef<[u8]> {
|
pub trait StreamEbml where Self: Sized + Stream, Self::Item: AsRef<[u8]> {
|
||||||
fn parse_ebml(self) -> EbmlStreamingParser<Self> {
|
fn parse_ebml(self) -> EbmlStreamingParser<Self> {
|
||||||
EbmlStreamingParser {
|
EbmlStreamingParser {
|
||||||
stream: self,
|
stream: self,
|
||||||
buffer: BytesMut::new(),
|
buffer: BytesMut::new(),
|
||||||
|
buffer_size_limit: None,
|
||||||
last_read: 0
|
last_read: 0
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -48,6 +57,12 @@ impl<I: AsRef<[u8]>, S: Stream<Item = I, Error = WebmetroError>> EbmlStreamingPa
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if let Some(limit) = self.buffer_size_limit {
|
||||||
|
if limit <= self.buffer.len() {
|
||||||
|
return Err(WebmetroError::ResourcesExceeded);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
match self.stream.poll() {
|
match self.stream.poll() {
|
||||||
Ok(Async::Ready(Some(chunk))) => {
|
Ok(Async::Ready(Some(chunk))) => {
|
||||||
self.buffer.reserve(chunk.as_ref().len());
|
self.buffer.reserve(chunk.as_ref().len());
|
||||||
|
|
Loading…
Reference in a new issue