Merge pull request #72 from najamelan/bugfix/receive_after_close

Keep processing incoming data even after we have initiated a close handshake.
pull/75/head
Daniel Abramov 5 years ago committed by GitHub
commit fc41b598ce
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 17
      src/protocol/mod.rs
  2. 4
      tests/connection_reset.rs
  3. 61
      tests/receive_after_init_close.rs

@ -396,7 +396,7 @@ impl WebSocketContext {
} }
} }
OpCode::Data(_) if !self.state.is_active() => { OpCode::Data(_) if !self.state.can_read() => {
// No data processing while closing. // No data processing while closing.
Ok(None) Ok(None)
} }
@ -533,7 +533,7 @@ impl WebSocketContext {
} }
/// The current connection state. /// The current connection state.
#[derive(Debug)] #[derive(Debug, PartialEq, Eq, Clone, Copy)]
enum WebSocketState { enum WebSocketState {
/// The connection is active. /// The connection is active.
Active, Active,
@ -550,12 +550,23 @@ enum WebSocketState {
impl WebSocketState { impl WebSocketState {
/// Tell if we're allowed to process normal messages. /// Tell if we're allowed to process normal messages.
fn is_active(&self) -> bool { fn is_active(&self) -> bool {
match *self { match self {
WebSocketState::Active => true, WebSocketState::Active => true,
_ => false, _ => false,
} }
} }
/// Tell if we should process incoming data. Note that if we send a close frame
/// but the remote hasn't confirmed, they might have sent data before they receive our
/// close frame, so we should still pass those to client code, hence ClosedByUs is valid.
fn can_read(&self) -> bool {
match self {
WebSocketState::Active |
WebSocketState::ClosedByUs => true,
_ => false,
}
}
/// Check if the state is active, return error if not. /// Check if the state is active, return error if not.
fn check_active(&self) -> Result<()> { fn check_active(&self) -> Result<()> {
match self { match self {

@ -34,7 +34,7 @@ fn test_close() {
let err = client.read_message().unwrap_err(); // now we should get ConnectionClosed let err = client.read_message().unwrap_err(); // now we should get ConnectionClosed
match err { match err {
Error::ConnectionClosed => {} Error::ConnectionClosed => {}
_ => panic!("unexpected error"), _ => panic!("unexpected error: {:?}", err),
} }
}); });
@ -51,7 +51,7 @@ fn test_close() {
let err = client_handler.read_message().unwrap_err(); // now we should get ConnectionClosed let err = client_handler.read_message().unwrap_err(); // now we should get ConnectionClosed
match err { match err {
Error::ConnectionClosed => {} Error::ConnectionClosed => {}
_ => panic!("unexpected error"), _ => panic!("unexpected error: {:?}", err),
} }
drop(client_handler); drop(client_handler);

@ -0,0 +1,61 @@
//! Verifies that we can read data messages even if we have initiated a close handshake,
//! but before we got confirmation.
use std::net::TcpListener;
use std::process::exit;
use std::thread::{sleep, spawn};
use std::time::Duration;
use tungstenite::{accept, connect, Error, Message};
use url::Url;
#[test]
fn test_receive_after_init_close() {
env_logger::init();
spawn(|| {
sleep(Duration::from_secs(5));
println!("Unit test executed too long, perhaps stuck on WOULDBLOCK...");
exit(1);
});
let server = TcpListener::bind("127.0.0.1:3013").unwrap();
let client_thread = spawn(move || {
let (mut client, _) = connect(Url::parse("ws://localhost:3013/socket").unwrap()).unwrap();
client
.write_message(Message::Text("Hello WebSocket".into()))
.unwrap();
let message = client.read_message().unwrap(); // receive close from server
assert!(message.is_close());
let err = client.read_message().unwrap_err(); // now we should get ConnectionClosed
match err {
Error::ConnectionClosed => {}
_ => panic!("unexpected error: {:?}", err),
}
});
let client_handler = server.incoming().next().unwrap();
let mut client_handler = accept(client_handler.unwrap()).unwrap();
client_handler.close(None).unwrap(); // send close to client
// This read should succeed even though we already initiated a close
let message = client_handler.read_message().unwrap();
assert_eq!(message.into_data(), b"Hello WebSocket");
assert!(client_handler.read_message().unwrap().is_close()); // receive acknowledgement
let err = client_handler.read_message().unwrap_err(); // now we should get ConnectionClosed
match err {
Error::ConnectionClosed => {}
_ => panic!("unexpected error: {:?}", err),
}
drop(client_handler);
client_thread.join().unwrap();
}
Loading…
Cancel
Save