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

Commit74c82be

Browse files
committed
Higher-level API for memory limits
1 parent3bde0f4 commit74c82be

File tree

4 files changed

+94
-7
lines changed

4 files changed

+94
-7
lines changed

‎src/client/builder.rs

Lines changed: 35 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@ use std::borrow::Cow;
1010
use std::convert::Into;
1111
pubuse url::{ParseError,Url};
1212

13+
constDEFAULT_MAX_DATAFRAME_SIZE:usize =1024*1024*100;
14+
constDEFAULT_MAX_MESSAGE_SIZE:usize =1024*1024*200;
15+
1316
#[cfg(any(feature ="sync", feature ="async"))]
1417
mod common_imports{
1518
pubusecrate::header::WebSocketAccept;
@@ -114,6 +117,8 @@ pub struct ClientBuilder<'u> {
114117
headers:Headers,
115118
version_set:bool,
116119
key_set:bool,
120+
max_dataframe_size:usize,
121+
max_message_size:usize,
117122
}
118123

119124
impl<'u>ClientBuilder<'u>{
@@ -161,6 +166,8 @@ impl<'u> ClientBuilder<'u> {
161166
version_set:false,
162167
key_set:false,
163168
headers:Headers::new(),
169+
max_dataframe_size:DEFAULT_MAX_DATAFRAME_SIZE,
170+
max_message_size:DEFAULT_MAX_MESSAGE_SIZE,
164171
}
165172
}
166173

@@ -289,6 +296,21 @@ impl<'u> ClientBuilder<'u> {
289296
self
290297
}
291298

299+
/// Set maximum dataframe size. Client will abort connection with error if it is exceed.
300+
/// Values larger than `u32::MAX` are not supported.
301+
pubfnmax_dataframe_size(mutself,value:usize) ->Self{
302+
self.max_dataframe_size = value;
303+
self
304+
}
305+
306+
/// Set maximum message size for which no more continuation dataframes are accepted.
307+
/// Client will abort connection with error if it is exceed.
308+
/// Values larger than `u32::MAX` are not supported.
309+
pubfnmax_message_size(mutself,value:usize) ->Self{
310+
self.max_message_size = value;
311+
self
312+
}
313+
292314
/// Add a custom `Sec-WebSocket-Key` header.
293315
/// Use this only if you know what you're doing, and this almost
294316
/// never has to be used.
@@ -493,7 +515,7 @@ impl<'u> ClientBuilder<'u> {
493515
// validate
494516
self.validate(&response)?;
495517

496-
Ok(Client::unchecked(reader, response.headers,true,false))
518+
Ok(Client::unchecked_with_limits(reader, response.headers,true,false,self.max_dataframe_size,self.max_dataframe_size))
497519
}
498520

499521
/// Connect to a websocket server asynchronously.
@@ -553,6 +575,8 @@ impl<'u> ClientBuilder<'u> {
553575
headers:self.headers,
554576
version_set:self.version_set,
555577
key_set:self.key_set,
578+
max_dataframe_size:self.max_dataframe_size,
579+
max_message_size:self.max_message_size,
556580
};
557581

558582
// check if we should connect over ssl or not
@@ -637,6 +661,8 @@ impl<'u> ClientBuilder<'u> {
637661
headers:self.headers,
638662
version_set:self.version_set,
639663
key_set:self.key_set,
664+
max_dataframe_size:self.max_dataframe_size,
665+
max_message_size:self.max_message_size,
640666
};
641667

642668
// put it all together
@@ -687,6 +713,8 @@ impl<'u> ClientBuilder<'u> {
687713
headers:self.headers,
688714
version_set:self.version_set,
689715
key_set:self.key_set,
716+
max_dataframe_size:self.max_dataframe_size,
717+
max_message_size:self.max_message_size,
690718
};
691719

692720
let future = tcp_stream.and_then(move |stream| builder.async_connect_on(stream));
@@ -746,6 +774,8 @@ impl<'u> ClientBuilder<'u> {
746774
headers:self.headers,
747775
version_set:self.version_set,
748776
key_set:self.key_set,
777+
max_dataframe_size:self.max_dataframe_size,
778+
max_message_size:self.max_message_size,
749779
};
750780
let resource = builder.build_request();
751781
let framed =crate::codec::http::HttpClientCodec.framed(stream);
@@ -755,6 +785,8 @@ impl<'u> ClientBuilder<'u> {
755785
subject:(Method::Get,RequestUri::AbsolutePath(resource)),
756786
};
757787

788+
let max_dataframe_size =self.max_dataframe_size;
789+
let max_message_size =self.max_message_size;
758790
let future = framed
759791
// send request
760792
.send(request)
@@ -770,8 +802,8 @@ impl<'u> ClientBuilder<'u> {
770802
.and_then(|message| builder.validate(&message).map(|()|(message, stream)))
771803
})
772804
// output the final client and metadata
773-
.map(|(message, stream)|{
774-
let codec =MessageCodec::default(Context::Client);
805+
.map(move|(message, stream)|{
806+
let codec =MessageCodec::new_with_limits(Context::Client, max_dataframe_size, max_message_size);
775807
let client =update_framed_codec(stream, codec);
776808
(client, message.headers)
777809
});

‎src/client/sync.rs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,23 @@ where
133133
}
134134
}
135135

136+
#[doc(hidden)]
137+
pubfnunchecked_with_limits(
138+
stream:BufReader<S>,
139+
headers:Headers,
140+
out_mask:bool,
141+
in_mask:bool,
142+
max_dataframe_size:usize,
143+
max_message_size:usize,
144+
) ->Self{
145+
Client{
146+
headers,
147+
stream,
148+
sender:Sender::new(out_mask),// true
149+
receiver:Receiver::new_with_limits(in_mask, max_dataframe_size, max_message_size),// false
150+
}
151+
}
152+
136153
/// Sends a single data frame to the remote endpoint.
137154
pubfnsend_dataframe<D>(&mutself,dataframe:&D) ->WebSocketResult<()>
138155
where

‎src/server/upgrade/async.rs

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ use hyper::status::StatusCode;
2121
use std::io::{self,ErrorKind};
2222
use tokio_codec::{Decoder,Framed,FramedParts};
2323

24+
constDEFAULT_MAX_DATAFRAME_SIZE:usize =1024*1024*100;
25+
constDEFAULT_MAX_MESSAGE_SIZE:usize =1024*1024*200;
26+
2427
/// An asynchronous websocket upgrade.
2528
///
2629
/// This struct is given when a connection is being upgraded to a websocket
@@ -84,7 +87,22 @@ where
8487
self.internal_accept(Some(custom_headers))
8588
}
8689

87-
fninternal_accept(mutself,custom_headers:Option<&Headers>) ->ClientNew<S>{
90+
/// Like `accept`, but also allows to set memory limits for incoming messages and dataframes
91+
pubfnaccept_with_limits(self,max_dataframe_size:usize,max_message_size:usize) ->ClientNew<S>{
92+
self.internal_accept_with_limits(None, max_dataframe_size, max_message_size)
93+
}
94+
95+
/// Like `accept_with`, but also allows to set memory limits for incoming messages and dataframes
96+
pubfnaccept_with_headers_and_limits(self,custom_headers:&Headers,max_dataframe_size:usize,max_message_size:usize) ->ClientNew<S>{
97+
self.internal_accept_with_limits(Some(custom_headers), max_dataframe_size, max_message_size)
98+
}
99+
100+
101+
fninternal_accept(self,custom_headers:Option<&Headers>) ->ClientNew<S>{
102+
self.internal_accept_with_limits(custom_headers,DEFAULT_MAX_DATAFRAME_SIZE,DEFAULT_MAX_MESSAGE_SIZE)
103+
}
104+
105+
fninternal_accept_with_limits(mutself,custom_headers:Option<&Headers>,max_dataframe_size:usize,max_message_size:usize) ->ClientNew<S>{
88106
let status =self.prepare_headers(custom_headers);
89107
letWsUpgrade{
90108
headers,
@@ -104,7 +122,7 @@ where
104122
headers: headers.clone(),
105123
})
106124
.map(move |s|{
107-
let codec =MessageCodec::default(Context::Server);
125+
let codec =MessageCodec::new_with_limits(Context::Server, max_dataframe_size, max_message_size);
108126
let client =update_framed_codec(s, codec);
109127
(client, headers)
110128
})

‎src/server/upgrade/sync.rs

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,9 @@ use hyper::http::h1::Incoming;
1313
use hyper::net::NetworkStream;
1414
use hyper::status::StatusCode;
1515

16+
constDEFAULT_MAX_DATAFRAME_SIZE:usize =1024*1024*100;
17+
constDEFAULT_MAX_MESSAGE_SIZE:usize =1024*1024*200;
18+
1619
/// This crate uses buffered readers to read in the handshake quickly, in order to
1720
/// interface with other use cases that don't use buffered readers the buffered readers
1821
/// is deconstructed when it is returned to the user and given as the underlying
@@ -61,7 +64,24 @@ where
6164
self.internal_accept(Some(custom_headers))
6265
}
6366

64-
fninternal_accept(mutself,headers:Option<&Headers>) ->Result<Client<S>,(S, io::Error)>{
67+
/// Accept the handshake request and send a response,
68+
/// if nothing goes wrong a client will be created.
69+
pubfnaccept_with_limits(self,max_dataframe_size:usize,max_message_size:usize) ->Result<Client<S>,(S, io::Error)>{
70+
self.internal_accept_with_limits(None, max_dataframe_size, max_message_size)
71+
}
72+
73+
/// Accept the handshake request and send a response while
74+
/// adding on a few headers. These headers are added before the required
75+
/// headers are, so some might be overwritten.
76+
pubfnaccept_with_headers_and_limits(self,custom_headers:&Headers,max_dataframe_size:usize,max_message_size:usize) ->Result<Client<S>,(S, io::Error)>{
77+
self.internal_accept_with_limits(Some(custom_headers), max_dataframe_size, max_message_size)
78+
}
79+
80+
fninternal_accept(self,headers:Option<&Headers>) ->Result<Client<S>,(S, io::Error)>{
81+
self.internal_accept_with_limits(headers,DEFAULT_MAX_DATAFRAME_SIZE,DEFAULT_MAX_MESSAGE_SIZE)
82+
}
83+
84+
fninternal_accept_with_limits(mutself,headers:Option<&Headers>,max_dataframe_size:usize,max_message_size:usize) ->Result<Client<S>,(S, io::Error)>{
6585
let status =self.prepare_headers(headers);
6686

6787
ifletErr(e) =self.send(status){
@@ -73,7 +93,7 @@ where
7393
None =>BufReader::new(self.stream),
7494
};
7595

76-
Ok(Client::unchecked(stream,self.headers,false,true))
96+
Ok(Client::unchecked_with_limits(stream,self.headers,false,true, max_dataframe_size, max_message_size))
7797
}
7898

7999
/// Reject the client's request to make a websocket connection.

0 commit comments

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