Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitcbf6e99

Browse files
committed
Implement dataframe and message size limits
1 parentd0adc93 commitcbf6e99

File tree

2 files changed

+62
-2
lines changed

2 files changed

+62
-2
lines changed

‎src/receiver.rs

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,10 @@ use crate::ws;
1414
usecrate::ws::receiver::ReceiverasReceiverTrait;
1515
usecrate::ws::receiver::{DataFrameIterator,MessageIterator};
1616

17+
constDEFAULT_MAX_DATAFRAME_SIZE:usize =1024*1024*100;
18+
constDEFAULT_MAX_MESSAGE_SIZE:usize =1024*1024*200;
19+
constMAX_DATAFRAMES_IN_ONE_MESSAGE:usize =1024*1024;
20+
1721
/// This reader bundles an existing stream with a parsing algorithm.
1822
/// It is used by the client in its `.split()` function as the reading component.
1923
pubstructReader<R>
@@ -74,14 +78,33 @@ where
7478
pubstructReceiver{
7579
buffer:Vec<DataFrame>,
7680
mask:bool,
81+
// u32s instead uf usizes to economize used memory by this struct
82+
max_dataframe_size:u32,
83+
max_message_size:u32,
7784
}
7885

7986
implReceiver{
8087
/// Create a new Receiver using the specified Reader.
88+
///
89+
/// Uses built-in limits for dataframe and message sizes.
8190
pubfnnew(mask:bool) ->Receiver{
91+
Receiver::new_with_limits(mask,DEFAULT_MAX_DATAFRAME_SIZE,DEFAULT_MAX_MESSAGE_SIZE)
92+
}
93+
94+
/// Create a new Receiver using the specified Reader, with configurable limits
95+
///
96+
/// Sizes should not be larger than `u32::MAX`.
97+
///
98+
/// Note that `max_message_size` denotes message size where no new dataframes would be read,
99+
/// so actual maximum message size is larger.
100+
pubfnnew_with_limits(mask:bool,max_dataframe_size:usize,max_message_size:usize) ->Receiver{
101+
let max_dataframe_size:u32 = max_dataframe_size.min(u32::MAXasusize)asu32;
102+
let max_message_size:u32 = max_message_size.min(u32::MAXasusize)asu32;
82103
Receiver{
83104
buffer:Vec::new(),
84105
mask,
106+
max_dataframe_size,
107+
max_message_size,
85108
}
86109
}
87110
}
@@ -96,14 +119,15 @@ impl ws::Receiver for Receiver {
96119
where
97120
R:Read,
98121
{
99-
DataFrame::read_dataframe(reader,self.mask)
122+
DataFrame::read_dataframe_with_limit(reader,self.mask,self.max_dataframe_sizeasusize)
100123
}
101124

102125
/// Returns the data frames that constitute one message.
103126
fnrecv_message_dataframes<R>(&mutself,reader:&mutR) ->WebSocketResult<Vec<DataFrame>>
104127
where
105128
R:Read,
106129
{
130+
letmut current_message_length:usize =self.buffer.iter().map(|x|x.data.len()).sum();
107131
letmut finished =ifself.buffer.is_empty(){
108132
let first =self.recv_dataframe(reader)?;
109133

@@ -114,6 +138,7 @@ impl ws::Receiver for Receiver {
114138
}
115139

116140
let finished = first.finished;
141+
current_message_length += first.data.len();
117142
self.buffer.push(first);
118143
finished
119144
}else{
@@ -126,7 +151,10 @@ impl ws::Receiver for Receiver {
126151

127152
match next.opcodeasu8{
128153
// Continuation opcode
129-
0 =>self.buffer.push(next),
154+
0 =>{
155+
current_message_length += next.data.len();
156+
self.buffer.push(next)
157+
}
130158
// Control frame
131159
8..=15 =>{
132160
returnOk(vec![next]);
@@ -138,6 +166,19 @@ impl ws::Receiver for Receiver {
138166
));
139167
}
140168
}
169+
170+
if !finished{
171+
ifself.buffer.len() >=MAX_DATAFRAMES_IN_ONE_MESSAGE{
172+
returnErr(WebSocketError::ProtocolError(
173+
"Exceeded count of data frames in one WebSocket message",
174+
));
175+
}
176+
if current_message_length >=self.max_message_sizeasusize{
177+
returnErr(WebSocketError::ProtocolError(
178+
"Exceeded maximum WebSocket message size",
179+
));
180+
}
181+
}
141182
}
142183

143184
Ok(::std::mem::replace(&mutself.buffer,Vec::new()))

‎websocket-base/src/dataframe.rs

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,25 @@ impl DataFrame {
9696

9797
DataFrame::read_dataframe_body(header, data, should_be_masked)
9898
}
99+
100+
/// Reads a DataFrame from a Reader, or error out if header declares exceeding limit you specify
101+
pubfnread_dataframe_with_limit<R>(reader:&mutR,should_be_masked:bool,limit:usize) ->WebSocketResult<Self>
102+
where
103+
R:Read,
104+
{
105+
let header = dfh::read_header(reader)?;
106+
107+
if header.len > limitasu64{
108+
returnErr(io::Error::new(io::ErrorKind::InvalidData,"exceeded DataFrame length limit").into());
109+
}
110+
letmut data:Vec<u8> =Vec::with_capacity(header.lenasusize);
111+
let read = reader.take(header.len).read_to_end(&mut data)?;
112+
if(readasu64) < header.len{
113+
returnErr(io::Error::new(io::ErrorKind::UnexpectedEof,"incomplete payload").into());
114+
}
115+
116+
DataFrame::read_dataframe_body(header, data, should_be_masked)
117+
}
99118
}
100119

101120
implDataFrameableforDataFrame{

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp