Compare commits

..

33 Commits

Author SHA1 Message Date
Daniel Abramov 57d9e23939
Merge pull request #360 from doylemark/master 2 years ago
Mark Doyle 9533c02280 return correct protocol error when missing http version 2 years ago
Daniel Abramov d7b559d724
Merge pull request #359 from mtbc/allow-IPv6-SocketAddr 2 years ago
Mark T. B. Carroll 8901dcc535
remove [] enclosing IPv6 host address 2 years ago
Alexey Galakhov 8f23e1765e
Merge pull request #358 from alexheretic/buffer-writes 2 years ago
Alex Butler f6a610f925 Add write_flush_behaviour test 2 years ago
Alex Butler dea67d6cca Fix doc typo 2 years ago
Alex Butler 41818166cf refactor WebSocketContext new 2 years ago
Alex Butler 0cada00fb5 Refactor write_one_frame -> buffer_frame 2 years ago
Alex Butler f33bb2cb97 Ensure out_buffer written when !can_read 2 years ago
Alex Butler 1b47964f18 split write and write_out_buffer internals 2 years ago
Alex Butler 2ef5b9a5e2 Buffer writes before writing to the underlying stream 2 years ago
Alex Butler 2cf7cfef04 Rework write 100k bench to have a slow writes & even slower flushes 2 years ago
Alexey Galakhov 5a3115c09b
Merge pull request #357 from alexheretic/flush-writes-less 2 years ago
Alex Butler 06e55a4ef2 Refactor additional_send writing 2 years ago
Alex Butler 84a54b76e6 Rename methods to `read`, `send`, `write` & `flush` 2 years ago
Daniel Abramov 79b39eb146
Merge pull request #356 from snapview/dependabot/cargo/criterion-0.5.0 2 years ago
Alex Butler 0203a1849b Remove send_queue, use out_buffer instead 2 years ago
Alex Butler 483d229707 Remove implicit write flushing 2 years ago
Alex Butler d298089bf3 Add write 100k micro-bench 2 years ago
dependabot[bot] 7242a22b91
Update criterion requirement from 0.4.0 to 0.5.0 2 years ago
Daniel Abramov 371f823044
Merge pull request #354 from CBenoit/fix-error-on-bad-root-cert 2 years ago
Benoît CORTIER ee3ffc9e9d
Gracefully handle invalid native root certificates 2 years ago
Daniel Abramov e5efe537b8
Merge pull request #351 from nickelc/deps/webpki 2 years ago
Constantin Nickel 8a436e7550 Remove unused `TlsError::Webpki` error variant 2 years ago
Daniel Abramov 314feea305
Merge pull request #348 from atouchet/trv 2 years ago
Alex Touchet 50d5a37bdc
Switch build status badge to GitHub Actions 2 years ago
Daniel Abramov 79fa37888f
Merge pull request #347 from snapview/github-actions 2 years ago
Daniel Abramov 746d938412 Use `cargo fmt` from nightly 2 years ago
Daniel Abramov a4863d3f10 Make code compile with any feature set 2 years ago
Daniel Abramov 7e4a15446d Properly activate features for examples and tests 2 years ago
Daniel Abramov 87e9f576af Make `cargo fmt` happy 2 years ago
Daniel Abramov e758f7dc2a Exchange Travis CI for GitHub Actions 2 years ago
  1. 70
      .github/workflows/ci.yml
  2. 1
      .gitignore
  3. 14
      .travis.yml
  4. 17
      CHANGELOG.md
  5. 43
      Cargo.toml
  6. 6
      README.md
  7. 3
      benches/buffer.rs
  8. 75
      benches/write.rs
  9. 6
      examples/autobahn-client.rs
  10. 4
      examples/autobahn-server.rs
  11. 4
      examples/client.rs
  12. 4
      examples/server.rs
  13. 6
      examples/srv_accept_unmasked_frames.rs
  14. 2
      fuzz/fuzz_targets/read_message_client.rs
  15. 2
      fuzz/fuzz_targets/read_message_server.rs
  16. 1
      src/client.rs
  17. 10
      src/error.rs
  18. 6
      src/handshake/client.rs
  19. 43
      src/handshake/server.rs
  20. 4
      src/lib.rs
  21. 122
      src/protocol/frame/mod.rs
  22. 2
      src/protocol/message.rs
  23. 400
      src/protocol/mod.rs
  24. 12
      src/tls.rs
  25. 41
      tests/connection_reset.rs
  26. 9
      tests/no_send_after_close.rs
  27. 15
      tests/receive_after_init_close.rs
  28. 68
      tests/write.rs

@ -0,0 +1,70 @@
name: CI
on: [push, pull_request]
jobs:
fmt:
name: Format
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v3
- uses: dtolnay/rust-toolchain@nightly
with:
components: rustfmt
- run: cargo fmt --all --check
test:
name: Test
runs-on: ubuntu-latest
strategy:
matrix:
rust:
- stable
steps:
- name: Checkout sources
uses: actions/checkout@v3
- name: Install toolchain
uses: dtolnay/rust-toolchain@master
with:
toolchain: ${{ matrix.rust }}
- name: Install dependencies
run: sudo apt-get install libssl-dev
- name: Install cargo-hack
uses: taiki-e/install-action@cargo-hack
- name: Check
run: cargo hack check --feature-powerset --all-targets
- name: Test
run: cargo test --release
autobahn:
name: Autobahn tests
runs-on: ubuntu-latest
strategy:
matrix:
rust:
- stable
- beta
- nightly
steps:
- name: Checkout sources
uses: actions/checkout@v3
- name: Install toolchain
uses: dtolnay/rust-toolchain@master
with:
toolchain: ${{ matrix.rust }}
- name: Running Autobahn TestSuite for client
run: ./scripts/autobahn-client.sh
- name: Running Autobahn TestSuite for server
run: ./scripts/autobahn-server.sh

1
.gitignore vendored

@ -1,2 +1,3 @@
target
Cargo.lock
.vscode

@ -1,14 +0,0 @@
language: rust
rust:
- stable
services:
- docker
before_script:
- export PATH="$PATH:$HOME/.cargo/bin"
script:
- cargo test --release
- echo "Running Autobahn TestSuite for client" && ./scripts/autobahn-client.sh
- echo "Running Autobahn TestSuite for server" && ./scripts/autobahn-server.sh

@ -1,3 +1,20 @@
# Unreleased (0.20.0)
- Remove many implicit flushing behaviours. In general reading and writing messages will no
longer flush until calling `flush`. An exception is automatic responses (e.g. pongs)
which will continue to be written and flushed when reading and writing.
This allows writing a batch of messages and flushing once, improving performance.
- Add `WebSocket::read`, `write`, `send`, `flush`. Deprecate `read_message`, `write_message`, `write_pending`.
- Add `FrameSocket::read`, `write`, `send`, `flush`. Remove `read_frame`, `write_frame`, `write_pending`.
Note: Previous use of `write_frame` may be replaced with `send`.
- Add `WebSocketContext::read`, `write`, `flush`. Remove `read_message`, `write_message`, `write_pending`.
Note: Previous use of `write_message` may be replaced with `write` + `flush`.
- Remove `send_queue`, replaced with using the frame write buffer to achieve similar results.
* Add `WebSocketConfig::max_write_buffer_size`. Deprecate `max_send_queue`.
* Add `Error::WriteBufferFull`. Remove `Error::SendQueueFull`.
Note: `WriteBufferFull` returns the message that could not be written as a `Message::Frame`.
- Add ability to buffer multiple writes before writing to the underlying stream, controlled by
`WebSocketConfig::write_buffer_size` (default 128 KiB). Improves batch message write performance.
# 0.19.0
- Update TLS dependencies.

@ -1,7 +1,7 @@
[package]
name = "ng-tungstenite"
description = "fork of tungstenite for Nextgraph.org"
categories = []
name = "tungstenite"
description = "Lightweight stream-based WebSocket implementation"
categories = ["web-programming::websocket", "network-programming"]
keywords = ["websocket", "io", "web"]
authors = ["Alexey Galakhov", "Daniel Abramov"]
license = "MIT OR Apache-2.0"
@ -24,7 +24,7 @@ native-tls = ["native-tls-crate"]
native-tls-vendored = ["native-tls", "native-tls-crate/vendored"]
rustls-tls-native-roots = ["__rustls-tls", "rustls-native-certs"]
rustls-tls-webpki-roots = ["__rustls-tls", "webpki-roots"]
__rustls-tls = ["rustls", "webpki"]
__rustls-tls = ["rustls"]
[dependencies]
data-encoding = { version = "2", optional = true }
@ -52,17 +52,12 @@ version = "0.21.0"
optional = true
version = "0.6.0"
[dependencies.webpki]
optional = true
version = "0.22"
features = ["std"]
[dependencies.webpki-roots]
optional = true
version = "0.23"
[dev-dependencies]
criterion = "0.4.0"
criterion = "0.5.0"
env_logger = "0.10.0"
input_buffer = "0.5.0"
net2 = "0.2.37"
@ -71,3 +66,31 @@ rand = "0.8.4"
[[bench]]
name = "buffer"
harness = false
[[bench]]
name = "write"
harness = false
[[example]]
name = "client"
required-features = ["handshake"]
[[example]]
name = "server"
required-features = ["handshake"]
[[example]]
name = "autobahn-client"
required-features = ["handshake"]
[[example]]
name = "autobahn-server"
required-features = ["handshake"]
[[example]]
name = "callback-error"
required-features = ["handshake"]
[[example]]
name = "srv_accept_unmasked_frames"
required-features = ["handshake"]

@ -14,11 +14,11 @@ fn main () {
spawn (move || {
let mut websocket = accept(stream.unwrap()).unwrap();
loop {
let msg = websocket.read_message().unwrap();
let msg = websocket.read().unwrap();
// We do not want to send back ping/pong messages.
if msg.is_binary() || msg.is_text() {
websocket.write_message(msg).unwrap();
websocket.send(msg).unwrap();
}
}
});
@ -36,7 +36,7 @@ take a look at [`tokio-tungstenite`](https://github.com/snapview/tokio-tungsteni
[![MIT licensed](https://img.shields.io/badge/License-MIT-blue.svg)](./LICENSE-MIT)
[![Apache-2.0 licensed](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](./LICENSE-APACHE)
[![Crates.io](https://img.shields.io/crates/v/tungstenite.svg?maxAge=2592000)](https://crates.io/crates/tungstenite)
[![Build Status](https://travis-ci.org/snapview/tungstenite-rs.svg?branch=master)](https://travis-ci.org/snapview/tungstenite-rs)
[![Build Status](https://github.com/snapview/tungstenite-rs/actions/workflows/ci.yml/badge.svg)](https://github.com/snapview/tungstenite-rs/actions)
[Documentation](https://docs.rs/tungstenite)

@ -1,5 +1,4 @@
use std::io::Result as IoResult;
use std::io::{Cursor, Read};
use std::io::{Cursor, Read, Result as IoResult};
use bytes::Buf;
use criterion::*;

@ -0,0 +1,75 @@
//! Benchmarks for write performance.
use criterion::{BatchSize, Criterion};
use std::{
hint,
io::{self, Read, Write},
time::{Duration, Instant},
};
use tungstenite::{Message, WebSocket};
const MOCK_WRITE_LEN: usize = 8 * 1024 * 1024;
/// `Write` impl that simulates slowish writes and slow flushes.
///
/// Each `write` can buffer up to 8 MiB before flushing but takes an additional **~80ns**
/// to simulate stuff going on in the underlying stream.
/// Each `flush` takes **~8µs** to simulate flush io.
struct MockWrite(Vec<u8>);
impl Read for MockWrite {
fn read(&mut self, _: &mut [u8]) -> io::Result<usize> {
Err(io::Error::new(io::ErrorKind::WouldBlock, "reads not supported"))
}
}
impl Write for MockWrite {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if self.0.len() + buf.len() > MOCK_WRITE_LEN {
self.flush()?;
}
// simulate io
spin(Duration::from_nanos(80));
self.0.extend(buf);
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
if !self.0.is_empty() {
// simulate io
spin(Duration::from_micros(8));
self.0.clear();
}
Ok(())
}
}
fn spin(duration: Duration) {
let a = Instant::now();
while a.elapsed() < duration {
hint::spin_loop();
}
}
fn benchmark(c: &mut Criterion) {
// Writes 100k small json text messages then flushes
c.bench_function("write 100k small texts then flush", |b| {
let mut ws = WebSocket::from_raw_socket(
MockWrite(Vec::with_capacity(MOCK_WRITE_LEN)),
tungstenite::protocol::Role::Server,
None,
);
b.iter_batched(
|| (0..100_000).map(|i| Message::Text(format!("{{\"id\":{i}}}"))),
|batch| {
for msg in batch {
ws.write(msg).unwrap();
}
ws.flush().unwrap();
},
BatchSize::SmallInput,
)
});
}
criterion::criterion_group!(write_benches, benchmark);
criterion::criterion_main!(write_benches);

@ -7,7 +7,7 @@ const AGENT: &str = "Tungstenite";
fn get_case_count() -> Result<u32> {
let (mut socket, _) = connect(Url::parse("ws://localhost:9001/getCaseCount").unwrap())?;
let msg = socket.read_message()?;
let msg = socket.read()?;
socket.close(None)?;
Ok(msg.into_text()?.parse::<u32>().unwrap())
}
@ -26,9 +26,9 @@ fn run_test(case: u32) -> Result<()> {
Url::parse(&format!("ws://localhost:9001/runCase?case={}&agent={}", case, AGENT)).unwrap();
let (mut socket, _) = connect(case_url)?;
loop {
match socket.read_message()? {
match socket.read()? {
msg @ Message::Text(_) | msg @ Message::Binary(_) => {
socket.write_message(msg)?;
socket.send(msg)?;
}
Message::Ping(_) | Message::Pong(_) | Message::Close(_) | Message::Frame(_) => {}
}

@ -17,9 +17,9 @@ fn handle_client(stream: TcpStream) -> Result<()> {
let mut socket = accept(stream).map_err(must_not_block)?;
info!("Running test");
loop {
match socket.read_message()? {
match socket.read()? {
msg @ Message::Text(_) | msg @ Message::Binary(_) => {
socket.write_message(msg)?;
socket.send(msg)?;
}
Message::Ping(_) | Message::Pong(_) | Message::Close(_) | Message::Frame(_) => {}
}

@ -14,9 +14,9 @@ fn main() {
println!("* {}", header);
}
socket.write_message(Message::Text("Hello WebSocket".into())).unwrap();
socket.send(Message::Text("Hello WebSocket".into())).unwrap();
loop {
let msg = socket.read_message().expect("Error reading message");
let msg = socket.read().expect("Error reading message");
println!("Received: {}", msg);
}
// socket.close(None);

@ -28,9 +28,9 @@ fn main() {
let mut websocket = accept_hdr(stream.unwrap(), callback).unwrap();
loop {
let msg = websocket.read_message().unwrap();
let msg = websocket.read().unwrap();
if msg.is_binary() || msg.is_text() {
websocket.write_message(msg).unwrap();
websocket.send(msg).unwrap();
}
}
});

@ -27,20 +27,18 @@ fn main() {
};
let config = Some(WebSocketConfig {
max_send_queue: None,
max_message_size: None,
max_frame_size: None,
// This setting allows to accept client frames which are not masked
// This is not in compliance with RFC 6455 but might be handy in some
// rare cases where it is necessary to integrate with existing/legacy
// clients which are sending unmasked frames
accept_unmasked_frames: true,
..<_>::default()
});
let mut websocket = accept_hdr_with_config(stream.unwrap(), callback, config).unwrap();
loop {
let msg = websocket.read_message().unwrap();
let msg = websocket.read().unwrap();
if msg.is_binary() || msg.is_text() {
println!("received message {}", msg);
}

@ -33,5 +33,5 @@ fuzz_target!(|data: &[u8]| {
//let vector: Vec<u8> = data.into();
let cursor = Cursor::new(data);
let mut socket = WebSocket::from_raw_socket(WriteMoc(cursor), Role::Client, None);
socket.read_message().ok();
socket.read().ok();
});

@ -33,5 +33,5 @@ fuzz_target!(|data: &[u8]| {
//let vector: Vec<u8> = data.into();
let cursor = Cursor::new(data);
let mut socket = WebSocket::from_raw_socket(WriteMoc(cursor), Role::Server, None);
socket.read_message().ok();
socket.read().ok();
});

@ -54,6 +54,7 @@ pub fn connect_with_config<Req: IntoClientRequest>(
let uri = request.uri();
let mode = uri_mode(uri)?;
let host = request.uri().host().ok_or(Error::Url(UrlError::NoHostName))?;
let host = if host.starts_with('[') { &host[1..host.len() - 1] } else { host };
let port = uri.port_u16().unwrap_or(match mode {
Mode::Plain => 80,
Mode::Tls => 443,

@ -53,9 +53,9 @@ pub enum Error {
/// Protocol violation.
#[error("WebSocket protocol error: {0}")]
Protocol(#[from] ProtocolError),
/// Message send queue full.
#[error("Send queue is full")]
SendQueueFull(Message),
/// Message write buffer is full.
#[error("Write buffer is full")]
WriteBufferFull(Message),
/// UTF coding error.
#[error("UTF-8 encoding error")]
Utf8,
@ -271,10 +271,6 @@ pub enum TlsError {
#[cfg(feature = "__rustls-tls")]
#[error("rustls error: {0}")]
Rustls(#[from] rustls::Error),
/// Webpki error.
#[cfg(feature = "__rustls-tls")]
#[error("webpki error: {0}")]
Webpki(#[from] webpki::Error),
/// DNS name resolution error.
#[cfg(feature = "__rustls-tls")]
#[error("Invalid DNS name")]

@ -87,8 +87,8 @@ impl<S: Read + Write> HandshakeRole for ClientHandshake<S> {
Ok(r) => r,
Err(Error::Http(mut e)) => {
*e.body_mut() = Some(tail);
return Err(Error::Http(e))
},
return Err(Error::Http(e));
}
Err(e) => return Err(e),
};
@ -258,7 +258,7 @@ impl TryParse for Response {
impl<'h, 'b: 'h> FromHttparse<httparse::Response<'h, 'b>> for Response {
fn from_httparse(raw: httparse::Response<'h, 'b>) -> Result<Self> {
if raw.version.expect("Bug: no HTTP version") < /*1.*/1 {
return Err(Error::Protocol(ProtocolError::WrongHttpMethod));
return Err(Error::Protocol(ProtocolError::WrongHttpVersion));
}
let headers = HeaderMap::from_httparse(raw.headers)?;

@ -30,7 +30,7 @@ pub type Request = HttpRequest<()>;
pub type Response = HttpResponse<()>;
/// Server error response type.
pub type ErrorResponse = HttpResponse<Option<Vec<u8>>>;
pub type ErrorResponse = HttpResponse<Option<String>>;
fn create_parts<T>(request: &HttpRequest<T>) -> Result<Builder> {
if request.method() != http::Method::GET {
@ -156,15 +156,23 @@ pub trait Callback: Sized {
/// Called whenever the server read the request from the client and is ready to reply to it.
/// May return additional reply headers.
/// Returning an error resulting in rejecting the incoming connection.
fn on_request(self, request: &Request) -> StdResult<(), ErrorResponse>;
fn on_request(
self,
request: &Request,
response: Response,
) -> StdResult<Response, ErrorResponse>;
}
impl<F> Callback for F
where
F: FnOnce(&Request) -> StdResult<(), ErrorResponse>,
F: FnOnce(&Request, Response) -> StdResult<Response, ErrorResponse>,
{
fn on_request(self, request: &Request) -> StdResult<(), ErrorResponse> {
self(request)
fn on_request(
self,
request: &Request,
response: Response,
) -> StdResult<Response, ErrorResponse> {
self(request, response)
}
}
@ -173,8 +181,12 @@ where
pub struct NoCallback;
impl Callback for NoCallback {
fn on_request(self, _request: &Request) -> StdResult<(), ErrorResponse> {
Ok(())
fn on_request(
self,
_request: &Request,
response: Response,
) -> StdResult<Response, ErrorResponse> {
Ok(response)
}
}
@ -228,24 +240,24 @@ impl<S: Read + Write, C: Callback> HandshakeRole for ServerHandshake<S, C> {
return Err(Error::Protocol(ProtocolError::JunkAfterRequest));
}
let response = create_response(&result)?;
let callback_result = if let Some(callback) = self.callback.take() {
callback.on_request(&result)
callback.on_request(&result, response)
} else {
Ok(())
Ok(response)
};
match callback_result {
Ok(_) => {
let response = create_response(&result)?;
Ok(response) => {
let mut output = vec![];
write_response(&mut output, &response)?;
ProcessingResult::Continue(HandshakeMachine::start_write(stream, output))
}
Err(resp) => {
// if resp.status().is_success() {
// return Err(Error::Protocol(ProtocolError::CustomResponseSuccessful));
// }
if resp.status().is_success() {
return Err(Error::Protocol(ProtocolError::CustomResponseSuccessful));
}
self.error_response = Some(resp);
let resp = self.error_response.as_ref().unwrap();
@ -254,7 +266,7 @@ impl<S: Read + Write, C: Callback> HandshakeRole for ServerHandshake<S, C> {
write_response(&mut output, resp)?;
if let Some(body) = resp.body() {
output.extend(body);
output.extend_from_slice(body.as_bytes());
}
ProcessingResult::Continue(HandshakeMachine::start_write(stream, output))
@ -267,6 +279,7 @@ impl<S: Read + Write, C: Callback> HandshakeRole for ServerHandshake<S, C> {
debug!("Server handshake failed.");
let (parts, body) = err.into_parts();
let body = body.map(|b| b.as_bytes().to_vec());
return Err(Error::Http(http::Response::from_parts(parts, body)));
} else {
debug!("Server handshake done.");

@ -25,7 +25,7 @@ pub mod protocol;
#[cfg(feature = "handshake")]
mod server;
pub mod stream;
#[cfg(any(feature = "native-tls", feature = "__rustls-tls"))]
#[cfg(all(any(feature = "native-tls", feature = "__rustls-tls"), feature = "handshake"))]
mod tls;
pub mod util;
@ -44,5 +44,5 @@ pub use crate::{
server::{accept, accept_hdr, accept_hdr_with_config, accept_with_config},
};
#[cfg(any(feature = "native-tls", feature = "__rustls-tls"))]
#[cfg(all(any(feature = "native-tls", feature = "__rustls-tls"), feature = "handshake"))]
pub use tls::{client_tls, client_tls_with_config, Connector};

@ -6,15 +6,14 @@ pub mod coding;
mod frame;
mod mask;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Read, Write};
use log::*;
pub use self::frame::{CloseFrame, Frame, FrameHeader};
use crate::{
error::{CapacityError, Error, Result},
ReadBuffer,
Message, ReadBuffer,
};
use log::*;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Read, Write};
pub use self::frame::{CloseFrame, Frame, FrameHeader};
/// A reader and writer for WebSocket frames.
#[derive(Debug)]
@ -57,7 +56,7 @@ where
Stream: Read,
{
/// Read a frame from stream.
pub fn read_frame(&mut self, max_size: Option<usize>) -> Result<Option<Frame>> {
pub fn read(&mut self, max_size: Option<usize>) -> Result<Option<Frame>> {
self.codec.read_frame(&mut self.stream, max_size)
}
}
@ -66,18 +65,28 @@ impl<Stream> FrameSocket<Stream>
where
Stream: Write,
{
/// Writes and immediately flushes a frame.
/// Equivalent to calling [`write`](Self::write) then [`flush`](Self::flush).
pub fn send(&mut self, frame: Frame) -> Result<()> {
self.write(frame)?;
self.flush()
}
/// Write a frame to stream.
///
/// This function guarantees that the frame is queued regardless of any errors.
/// There is no need to resend the frame. In order to handle WouldBlock or Incomplete,
/// call write_pending() afterwards.
pub fn write_frame(&mut self, frame: Frame) -> Result<()> {
self.codec.write_frame(&mut self.stream, frame)
/// A subsequent call should be made to [`flush`](Self::flush) to flush writes.
///
/// This function guarantees that the frame is queued unless [`Error::WriteBufferFull`]
/// is returned.
/// In order to handle WouldBlock or Incomplete, call [`flush`](Self::flush) afterwards.
pub fn write(&mut self, frame: Frame) -> Result<()> {
self.codec.buffer_frame(&mut self.stream, frame)
}
/// Complete pending write, if any.
pub fn write_pending(&mut self) -> Result<()> {
self.codec.write_pending(&mut self.stream)
/// Flush writes.
pub fn flush(&mut self) -> Result<()> {
self.codec.write_out_buffer(&mut self.stream)?;
Ok(self.stream.flush()?)
}
}
@ -88,6 +97,14 @@ pub(super) struct FrameCodec {
in_buffer: ReadBuffer,
/// Buffer to send packets to the network.
out_buffer: Vec<u8>,
/// Capacity limit for `out_buffer`.
max_out_buffer_len: usize,
/// Buffer target length to reach before writing to the stream
/// on calls to `buffer_frame`.
///
/// Setting this to non-zero will buffer small writes from hitting
/// the stream.
out_buffer_write_len: usize,
/// Header and remaining size of the incoming packet being processed.
header: Option<(FrameHeader, u64)>,
}
@ -95,7 +112,13 @@ pub(super) struct FrameCodec {
impl FrameCodec {
/// Create a new frame codec.
pub(super) fn new() -> Self {
Self { in_buffer: ReadBuffer::new(), out_buffer: Vec::new(), header: None }
Self {
in_buffer: ReadBuffer::new(),
out_buffer: Vec::new(),
max_out_buffer_len: usize::MAX,
out_buffer_write_len: 0,
header: None,
}
}
/// Create a new frame codec from partially read data.
@ -103,10 +126,23 @@ impl FrameCodec {
Self {
in_buffer: ReadBuffer::from_partially_read(part),
out_buffer: Vec::new(),
max_out_buffer_len: usize::MAX,
out_buffer_write_len: 0,
header: None,
}
}
/// Sets a maximum size for the out buffer.
pub(super) fn set_max_out_buffer_len(&mut self, max: usize) {
self.max_out_buffer_len = max;
}
/// Sets [`Self::buffer_frame`] buffer target length to reach before
/// writing to the stream.
pub(super) fn set_out_buffer_write_len(&mut self, len: usize) {
self.out_buffer_write_len = len;
}
/// Read a frame from the provided stream.
pub(super) fn read_frame<Stream>(
&mut self,
@ -165,19 +201,37 @@ impl FrameCodec {
Ok(Some(frame))
}
/// Write a frame to the provided stream.
pub(super) fn write_frame<Stream>(&mut self, stream: &mut Stream, frame: Frame) -> Result<()>
/// Writes a frame into the `out_buffer`.
/// If the out buffer size is over the `out_buffer_write_len` will also write
/// the out buffer into the provided `stream`.
///
/// To ensure buffered frames are written call [`Self::write_out_buffer`].
///
/// May write to the stream, will **not** flush.
pub(super) fn buffer_frame<Stream>(&mut self, stream: &mut Stream, frame: Frame) -> Result<()>
where
Stream: Write,
{
if frame.len() + self.out_buffer.len() > self.max_out_buffer_len {
return Err(Error::WriteBufferFull(Message::Frame(frame)));
}
trace!("writing frame {}", frame);
self.out_buffer.reserve(frame.len());
frame.format(&mut self.out_buffer).expect("Bug: can't write to vector");
self.write_pending(stream)
if self.out_buffer.len() > self.out_buffer_write_len {
self.write_out_buffer(stream)
} else {
Ok(())
}
}
/// Complete pending write, if any.
pub(super) fn write_pending<Stream>(&mut self, stream: &mut Stream) -> Result<()>
/// Writes the out_buffer to the provided stream.
///
/// Does **not** flush.
pub(super) fn write_out_buffer<Stream>(&mut self, stream: &mut Stream) -> Result<()>
where
Stream: Write,
{
@ -193,16 +247,8 @@ impl FrameCodec {
}
self.out_buffer.drain(0..len);
}
stream.flush()?;
Ok(())
}
}
#[cfg(test)]
impl FrameCodec {
/// Returns the size of the output buffer.
pub(super) fn output_buffer_len(&self) -> usize {
self.out_buffer.len()
Ok(())
}
}
@ -224,11 +270,11 @@ mod tests {
let mut sock = FrameSocket::new(raw);
assert_eq!(
sock.read_frame(None).unwrap().unwrap().into_data(),
sock.read(None).unwrap().unwrap().into_data(),
vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]
);
assert_eq!(sock.read_frame(None).unwrap().unwrap().into_data(), vec![0x03, 0x02, 0x01]);
assert!(sock.read_frame(None).unwrap().is_none());
assert_eq!(sock.read(None).unwrap().unwrap().into_data(), vec![0x03, 0x02, 0x01]);
assert!(sock.read(None).unwrap().is_none());
let (_, rest) = sock.into_inner();
assert_eq!(rest, vec![0x99]);
@ -239,7 +285,7 @@ mod tests {
let raw = Cursor::new(vec![0x02, 0x03, 0x04, 0x05, 0x06, 0x07]);
let mut sock = FrameSocket::from_partially_read(raw, vec![0x82, 0x07, 0x01]);
assert_eq!(
sock.read_frame(None).unwrap().unwrap().into_data(),
sock.read(None).unwrap().unwrap().into_data(),
vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]
);
}
@ -249,10 +295,10 @@ mod tests {
let mut sock = FrameSocket::new(Vec::new());
let frame = Frame::ping(vec![0x04, 0x05]);
sock.write_frame(frame).unwrap();
sock.send(frame).unwrap();
let frame = Frame::pong(vec![0x01]);
sock.write_frame(frame).unwrap();
sock.send(frame).unwrap();
let (buf, _) = sock.into_inner();
assert_eq!(buf, vec![0x89, 0x02, 0x04, 0x05, 0x8a, 0x01, 0x01]);
@ -264,7 +310,7 @@ mod tests {
0x83, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00,
]);
let mut sock = FrameSocket::new(raw);
let _ = sock.read_frame(None); // should not crash
let _ = sock.read(None); // should not crash
}
#[test]
@ -272,7 +318,7 @@ mod tests {
let raw = Cursor::new(vec![0x82, 0x07, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]);
let mut sock = FrameSocket::new(raw);
assert!(matches!(
sock.read_frame(Some(5)),
sock.read(Some(5)),
Err(Error::Capacity(CapacityError::MessageTooLong { size: 7, max_size: 5 }))
));
}

@ -185,7 +185,7 @@ impl Message {
Message::Text(string.into())
}
/// Create a new binary WebSocket message by converting to Vec<u8>.
/// Create a new binary WebSocket message by converting to `Vec<u8>`.
pub fn binary<B>(bin: B) -> Message
where
B: Into<Vec<u8>>,

@ -6,13 +6,6 @@ mod message;
pub use self::{frame::CloseFrame, message::Message};
use log::*;
use std::{
collections::VecDeque,
io::{ErrorKind as IoErrorKind, Read, Write},
mem::replace,
};
use self::{
frame::{
coding::{CloseCode, Control as OpCtl, Data as OpData, OpCode},
@ -24,6 +17,11 @@ use crate::{
error::{Error, ProtocolError, Result},
util::NonBlockingResult,
};
use log::*;
use std::{
io::{ErrorKind as IoErrorKind, Read, Write},
mem::replace,
};
/// Indicates a Client or Server role of the websocket
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@ -37,10 +35,21 @@ pub enum Role {
/// The configuration for WebSocket connection.
#[derive(Debug, Clone, Copy)]
pub struct WebSocketConfig {
/// The size of the send queue. You can use it to turn on/off the backpressure features. `None`
/// means here that the size of the queue is unlimited. The default value is the unlimited
/// queue.
/// Does nothing, instead use `max_write_buffer_size`.
#[deprecated]
pub max_send_queue: Option<usize>,
/// The target minimum size of the write buffer to reach before writing the data
/// to the underlying stream.
/// The default value is 128 KiB.
///
/// Note: [`flush`](WebSocket::flush) will always fully write the buffer regardless.
pub write_buffer_size: usize,
/// The max size of the write buffer in bytes. Setting this can provide backpressure
/// in the case the write buffer is filling up due to write errors.
/// The default value is unlimited.
///
/// Note: Should always be set higher than [`write_buffer_size`](Self::write_buffer_size).
pub max_write_buffer_size: usize,
/// The maximum size of a message. `None` means no size limit. The default value is 64 MiB
/// which should be reasonably big for all normal use-cases but small enough to prevent
/// memory eating by a malicious user.
@ -60,8 +69,11 @@ pub struct WebSocketConfig {
impl Default for WebSocketConfig {
fn default() -> Self {
#[allow(deprecated)]
WebSocketConfig {
max_send_queue: None,
write_buffer_size: 128 * 1024,
max_write_buffer_size: usize::MAX,
max_message_size: Some(64 << 20),
max_frame_size: Some(16 << 20),
accept_unmasked_frames: false,
@ -73,6 +85,8 @@ impl Default for WebSocketConfig {
///
/// This is THE structure you want to create to be able to speak the WebSocket protocol.
/// It may be created by calling `connect`, `accept` or `client` functions.
///
/// Use [`WebSocket::read`], [`WebSocket::send`] to received and send messages.
#[derive(Debug)]
pub struct WebSocket<Stream> {
/// The underlying socket.
@ -146,82 +160,116 @@ impl<Stream> WebSocket<Stream> {
impl<Stream: Read + Write> WebSocket<Stream> {
/// Read a message from stream, if possible.
///
/// This will queue responses to ping and close messages to be sent. It will call
/// `write_pending` before trying to read in order to make sure that those responses
/// make progress even if you never call `write_pending`. That does mean that they
/// get sent out earliest on the next call to `read_message`, `write_message` or `write_pending`.
/// This will also queue responses to ping and close messages. These responses
/// will be written and flushed on the next call to [`read`](Self::read),
/// [`write`](Self::write) or [`flush`](Self::flush).
///
/// ## Closing the connection
/// # Closing the connection
/// When the remote endpoint decides to close the connection this will return
/// the close message with an optional close frame.
///
/// You should continue calling `read_message`, `write_message` or `write_pending` to drive
/// the reply to the close frame until [Error::ConnectionClosed] is returned. Once that happens
/// it is safe to drop the underlying connection.
pub fn read_message(&mut self) -> Result<Message> {
self.context.read_message(&mut self.socket)
/// You should continue calling [`read`](Self::read), [`write`](Self::write) or
/// [`flush`](Self::flush) to drive the reply to the close frame until [`Error::ConnectionClosed`]
/// is returned. Once that happens it is safe to drop the underlying connection.
pub fn read(&mut self) -> Result<Message> {
self.context.read(&mut self.socket)
}
/// Send a message to stream, if possible.
/// Writes and immediately flushes a message.
/// Equivalent to calling [`write`](Self::write) then [`flush`](Self::flush).
pub fn send(&mut self, message: Message) -> Result<()> {
self.write(message)?;
self.flush()
}
/// Write a message to the provided stream, if possible.
///
/// A subsequent call should be made to [`flush`](Self::flush) to flush writes.
///
/// WebSocket will buffer a configurable number of messages at a time, except to reply to Ping
/// requests. A Pong reply will jump the queue because the
/// [websocket RFC](https://tools.ietf.org/html/rfc6455#section-5.5.2) specifies it should be sent
/// as soon as is practical.
/// In the event of stream write failure the message frame will be stored
/// in the write buffer and will try again on the next call to [`write`](Self::write)
/// or [`flush`](Self::flush).
///
/// Note that upon receiving a ping message, tungstenite cues a pong reply automatically.
/// When you call either `read_message`, `write_message` or `write_pending` next it will try to send
/// that pong out if the underlying connection can take more data. This means you should not
/// respond to ping frames manually.
/// If the write buffer would exceed the configured [`WebSocketConfig::max_write_buffer_size`]
/// [`Err(WriteBufferFull(msg_frame))`](Error::WriteBufferFull) is returned.
///
/// This call will generally not flush. However, if there are queued automatic messages
/// they will be written and eagerly flushed.
///
/// For example, upon receiving ping messages tungstenite queues pong replies automatically.
/// The next call to [`read`](Self::read), [`write`](Self::write) or [`flush`](Self::flush)
/// will write & flush the pong reply. This means you should not respond to ping frames manually.
///
/// You can however send pong frames manually in order to indicate a unidirectional heartbeat
/// as described in [RFC 6455](https://tools.ietf.org/html/rfc6455#section-5.5.3). Note that
/// if `read_message` returns a ping, you should call `write_pending` until it doesn't return
/// WouldBlock before passing a pong to `write_message`, otherwise the response to the
/// ping will not be sent, but rather replaced by your custom pong message.
///
/// ## Errors
/// - If the WebSocket's send queue is full, `SendQueueFull` will be returned
/// along with the passed message. Otherwise, the message is queued and Ok(()) is returned.
/// - If the connection is closed and should be dropped, this will return [Error::ConnectionClosed].
/// - If you try again after [Error::ConnectionClosed] was returned either from here or from `read_message`,
/// [Error::AlreadyClosed] will be returned. This indicates a program error on your part.
/// - [Error::Io] is returned if the underlying connection returns an error
/// if [`read`](Self::read) returns a ping, you should [`flush`](Self::flush) before passing
/// a custom pong to [`write`](Self::write), otherwise the automatic queued response to the
/// ping will not be sent as it will be replaced by your custom pong message.
///
/// # Errors
/// - If the WebSocket's write buffer is full, [`Error::WriteBufferFull`] will be returned
/// along with the equivalent passed message frame.
/// - If the connection is closed and should be dropped, this will return [`Error::ConnectionClosed`].
/// - If you try again after [`Error::ConnectionClosed`] was returned either from here or from
/// [`read`](Self::read), [`Error::AlreadyClosed`] will be returned. This indicates a program
/// error on your part.
/// - [`Error::Io`] is returned if the underlying connection returns an error
/// (consider these fatal except for WouldBlock).
/// - [Error::Capacity] if your message size is bigger than the configured max message size.
pub fn write_message(&mut self, message: Message) -> Result<()> {
self.context.write_message(&mut self.socket, message)
/// - [`Error::Capacity`] if your message size is bigger than the configured max message size.
pub fn write(&mut self, message: Message) -> Result<()> {
self.context.write(&mut self.socket, message)
}
/// Flush the pending send queue.
pub fn write_pending(&mut self) -> Result<()> {
self.context.write_pending(&mut self.socket)
/// Flush writes.
///
/// Ensures all messages previously passed to [`write`](Self::write) and automatic
/// queued pong responses are written & flushed into the underlying stream.
pub fn flush(&mut self) -> Result<()> {
self.context.flush(&mut self.socket)
}
/// Close the connection.
///
/// This function guarantees that the close frame will be queued.
/// There is no need to call it again. Calling this function is
/// the same as calling `write_message(Message::Close(..))`.
/// the same as calling `write(Message::Close(..))`.
///
/// After queuing the close frame you should continue calling `read_message` or
/// `write_pending` to drive the close handshake to completion.
/// After queuing the close frame you should continue calling [`read`](Self::read) or
/// [`flush`](Self::flush) to drive the close handshake to completion.
///
/// The websocket RFC defines that the underlying connection should be closed
/// by the server. Tungstenite takes care of this asymmetry for you.
///
/// When the close handshake is finished (we have both sent and received
/// a close message), `read_message` or `write_pending` will return
/// a close message), [`read`](Self::read) or [`flush`](Self::flush) will return
/// [Error::ConnectionClosed] if this endpoint is the server.
///
/// If this endpoint is a client, [Error::ConnectionClosed] will only be
/// returned after the server has closed the underlying connection.
///
/// It is thus safe to drop the underlying connection as soon as [Error::ConnectionClosed]
/// is returned from `read_message` or `write_pending`.
/// is returned from [`read`](Self::read) or [`flush`](Self::flush).
pub fn close(&mut self, code: Option<CloseFrame>) -> Result<()> {
self.context.close(&mut self.socket, code)
}
/// Old name for [`read`](Self::read).
#[deprecated(note = "Use `read`")]
pub fn read_message(&mut self) -> Result<Message> {
self.read()
}
/// Old name for [`send`](Self::send).
#[deprecated(note = "Use `send`")]
pub fn write_message(&mut self, message: Message) -> Result<()> {
self.send(message)
}
/// Old name for [`flush`](Self::flush).
#[deprecated(note = "Use `flush`")]
pub fn write_pending(&mut self) -> Result<()> {
self.flush()
}
}
/// A context for managing WebSocket stream.
@ -235,10 +283,8 @@ pub struct WebSocketContext {
state: WebSocketState,
/// Receive: an incomplete message being processed.
incomplete: Option<IncompleteMessage>,
/// Send: a data send queue.
send_queue: VecDeque<Frame>,
/// Send: an OOB pong message.
pong: Option<Frame>,
/// Send in addition to regular messages E.g. "pong" or "close".
additional_send: Option<Frame>,
/// The configuration for the websocket session.
config: WebSocketConfig,
}
@ -246,28 +292,32 @@ pub struct WebSocketContext {
impl WebSocketContext {
/// Create a WebSocket context that manages a post-handshake stream.
pub fn new(role: Role, config: Option<WebSocketConfig>) -> Self {
WebSocketContext {
role,
frame: FrameCodec::new(),
state: WebSocketState::Active,
incomplete: None,
send_queue: VecDeque::new(),
pong: None,
config: config.unwrap_or_default(),
}
Self::_new(role, FrameCodec::new(), config.unwrap_or_default())
}
/// Create a WebSocket context that manages an post-handshake stream.
pub fn from_partially_read(part: Vec<u8>, role: Role, config: Option<WebSocketConfig>) -> Self {
WebSocketContext {
frame: FrameCodec::from_partially_read(part),
..WebSocketContext::new(role, config)
Self::_new(role, FrameCodec::from_partially_read(part), config.unwrap_or_default())
}
fn _new(role: Role, mut frame: FrameCodec, config: WebSocketConfig) -> Self {
frame.set_max_out_buffer_len(config.max_write_buffer_size);
frame.set_out_buffer_write_len(config.write_buffer_size);
Self {
role,
frame,
state: WebSocketState::Active,
incomplete: None,
additional_send: None,
config,
}
}
/// Change the configuration.
pub fn set_config(&mut self, set_func: impl FnOnce(&mut WebSocketConfig)) {
set_func(&mut self.config)
set_func(&mut self.config);
self.frame.set_max_out_buffer_len(self.config.max_write_buffer_size);
self.frame.set_out_buffer_write_len(self.config.write_buffer_size);
}
/// Read the configuration.
@ -294,17 +344,23 @@ impl WebSocketContext {
///
/// This function sends pong and close responses automatically.
/// However, it never blocks on write.
pub fn read_message<Stream>(&mut self, stream: &mut Stream) -> Result<Message>
pub fn read<Stream>(&mut self, stream: &mut Stream) -> Result<Message>
where
Stream: Read + Write,
{
// Do not read from already closed connections.
self.state.check_active()?;
self.state.check_not_terminated()?;
loop {
if self.additional_send.is_some() {
// Since we may get ping or close, we need to reply to the messages even during read.
// Thus we call write_pending() but ignore its blocking.
self.write_pending(stream).no_block()?;
// Thus we flush but ignore its blocking.
self.flush(stream).no_block()?;
} else if self.role == Role::Server && !self.state.can_read() {
self.state = WebSocketState::Terminated;
return Err(Error::ConnectionClosed);
}
// If we get here, either write blocks or we have nothing to write.
// Thus if read blocks, just let it return WouldBlock.
if let Some(message) = self.read_message_frame(stream)? {
@ -314,78 +370,94 @@ impl WebSocketContext {
}
}
/// Send a message to the provided stream, if possible.
/// Write a message to the provided stream.
///
/// A subsequent call should be made to [`flush`](Self::flush) to flush writes.
///
/// WebSocket will buffer a configurable number of messages at a time, except to reply to Ping
/// and Close requests. If the WebSocket's send queue is full, `SendQueueFull` will be returned
/// along with the passed message. Otherwise, the message is queued and Ok(()) is returned.
/// In the event of stream write failure the message frame will be stored
/// in the write buffer and will try again on the next call to [`write`](Self::write)
/// or [`flush`](Self::flush).
///
/// Note that only the last pong frame is stored to be sent, and only the
/// most recent pong frame is sent if multiple pong frames are queued.
pub fn write_message<Stream>(&mut self, stream: &mut Stream, message: Message) -> Result<()>
/// If the write buffer would exceed the configured [`WebSocketConfig::max_write_buffer_size`]
/// [`Err(WriteBufferFull(msg_frame))`](Error::WriteBufferFull) is returned.
pub fn write<Stream>(&mut self, stream: &mut Stream, message: Message) -> Result<()>
where
Stream: Read + Write,
{
// When terminated, return AlreadyClosed.
self.state.check_active()?;
self.state.check_not_terminated()?;
// Do not write after sending a close frame.
if !self.state.is_active() {
return Err(Error::Protocol(ProtocolError::SendAfterClosing));
}
if let Some(max_send_queue) = self.config.max_send_queue {
if self.send_queue.len() >= max_send_queue {
// Try to make some room for the new message.
// Do not return here if write would block, ignore WouldBlock silently
// since we must queue the message anyway.
self.write_pending(stream).no_block()?;
}
if self.send_queue.len() >= max_send_queue {
return Err(Error::SendQueueFull(message));
}
}
let frame = match message {
Message::Text(data) => Frame::message(data.into(), OpCode::Data(OpData::Text), true),
Message::Binary(data) => Frame::message(data, OpCode::Data(OpData::Binary), true),
Message::Ping(data) => Frame::ping(data),
Message::Pong(data) => {
self.pong = Some(Frame::pong(data));
return self.write_pending(stream);
self.set_additional(Frame::pong(data));
// Note: user pongs can be user flushed so no need to flush here
return self._write(stream, None).map(|_| ());
}
Message::Close(code) => return self.close(stream, code),
Message::Frame(f) => f,
};
self.send_queue.push_back(frame);
self.write_pending(stream)
let should_flush = self._write(stream, Some(frame))?;
if should_flush {
self.flush(stream)?;
}
Ok(())
}
/// Flush writes.
///
/// Ensures all messages previously passed to [`write`](Self::write) and automatically
/// queued pong responses are written & flushed into the `stream`.
#[inline]
pub fn flush<Stream>(&mut self, stream: &mut Stream) -> Result<()>
where
Stream: Read + Write,
{
self._write(stream, None)?;
self.frame.write_out_buffer(stream)?;
Ok(stream.flush()?)
}
/// Flush the pending send queue.
pub fn write_pending<Stream>(&mut self, stream: &mut Stream) -> Result<()>
/// Writes any data in the out_buffer, `additional_send` and given `data`.
///
/// Does **not** flush.
///
/// Returns true if the write contents indicate we should flush immediately.
fn _write<Stream>(&mut self, stream: &mut Stream, data: Option<Frame>) -> Result<bool>
where
Stream: Read + Write,
{
// First, make sure we have no pending frame sending.
self.frame.write_pending(stream)?;
if let Some(data) = data {
self.buffer_frame(stream, data)?;
}
// Upon receipt of a Ping frame, an endpoint MUST send a Pong frame in
// response, unless it already received a Close frame. It SHOULD
// respond with Pong frame as soon as is practical. (RFC 6455)
if let Some(pong) = self.pong.take() {
trace!("Sending pong reply");
self.send_one_frame(stream, pong)?;
let should_flush = if let Some(msg) = self.additional_send.take() {
trace!("Sending pong/close");
match self.buffer_frame(stream, msg) {
Err(Error::WriteBufferFull(Message::Frame(msg))) => {
// if an system message would exceed the buffer put it back in
// `additional_send` for retry. Otherwise returning this error
// may not make sense to the user, e.g. calling `flush`.
self.set_additional(msg);
false
}
Err(err) => return Err(err),
Ok(_) => true,
}
// If we have any unsent frames, send them.
trace!("Frames still in queue: {}", self.send_queue.len());
while let Some(data) = self.send_queue.pop_front() {
self.send_one_frame(stream, data)?;
}
// If we get to this point, the send queue is empty and the underlying socket is still
// willing to take more data.
} else {
false
};
// If we're closing and there is nothing to send anymore, we should close the connection.
if self.role == Role::Server && !self.state.can_read() {
@ -395,10 +467,11 @@ impl WebSocketContext {
// maximum segment lifetimes (2MSL), while there is no corresponding
// server impact as a TIME_WAIT connection is immediately reopened upon
// a new SYN with a higher seq number). (RFC 6455)
self.frame.write_out_buffer(stream)?;
self.state = WebSocketState::Terminated;
Err(Error::ConnectionClosed)
} else {
Ok(())
Ok(should_flush)
}
}
@ -406,7 +479,7 @@ impl WebSocketContext {
///
/// This function guarantees that the close frame will be queued.
/// There is no need to call it again. Calling this function is
/// the same as calling `write(Message::Close(..))`.
/// the same as calling `send(Message::Close(..))`.
pub fn close<Stream>(&mut self, stream: &mut Stream, code: Option<CloseFrame>) -> Result<()>
where
Stream: Read + Write,
@ -414,11 +487,9 @@ impl WebSocketContext {
if let WebSocketState::Active = self.state {
self.state = WebSocketState::ClosedByUs;
let frame = Frame::close(code);
self.send_queue.push_back(frame);
} else {
// Already closed, nothing to do.
self._write(stream, Some(frame))?;
}
self.write_pending(stream)
self.flush(stream)
}
/// Try to decode one message frame. May return None.
@ -487,7 +558,7 @@ impl WebSocketContext {
let data = frame.into_data();
// No ping processing after we sent a close frame.
if self.state.is_active() {
self.pong = Some(Frame::pong(data.clone()));
self.set_additional(Frame::pong(data.clone()));
}
Ok(Some(Message::Ping(data)))
}
@ -571,7 +642,7 @@ impl WebSocketContext {
let reply = Frame::close(close.clone());
debug!("Replying to close with {:?}", reply);
self.send_queue.push_back(reply);
self.set_additional(reply);
Some(close)
}
@ -588,8 +659,8 @@ impl WebSocketContext {
}
}
/// Send a single pending frame.
fn send_one_frame<Stream>(&mut self, stream: &mut Stream, mut frame: Frame) -> Result<()>
/// Write a single frame into the write-buffer.
fn buffer_frame<Stream>(&mut self, stream: &mut Stream, mut frame: Frame) -> Result<()>
where
Stream: Read + Write,
{
@ -603,7 +674,18 @@ impl WebSocketContext {
}
trace!("Sending frame: {:?}", frame);
self.frame.write_frame(stream, frame).check_connection_reset(self.state)
self.frame.buffer_frame(stream, frame).check_connection_reset(self.state)
}
/// Replace `additional_send` if it is currently a `Pong` message.
fn set_additional(&mut self, add: Frame) {
let empty_or_pong = self
.additional_send
.as_ref()
.map_or(true, |f| f.header().opcode == OpCode::Control(OpCtl::Pong));
if empty_or_pong {
self.additional_send.replace(add);
}
}
}
@ -636,7 +718,7 @@ impl WebSocketState {
}
/// Check if the state is active, return error if not.
fn check_active(self) -> Result<()> {
fn check_not_terminated(self) -> Result<()> {
match self {
WebSocketState::Terminated => Err(Error::AlreadyClosed),
_ => Ok(()),
@ -688,64 +770,6 @@ mod tests {
}
}
struct WouldBlockStreamMoc;
impl io::Write for WouldBlockStreamMoc {
fn write(&mut self, _: &[u8]) -> io::Result<usize> {
Err(io::Error::new(io::ErrorKind::WouldBlock, "would block"))
}
fn flush(&mut self) -> io::Result<()> {
Err(io::Error::new(io::ErrorKind::WouldBlock, "would block"))
}
}
impl io::Read for WouldBlockStreamMoc {
fn read(&mut self, _: &mut [u8]) -> io::Result<usize> {
Err(io::Error::new(io::ErrorKind::WouldBlock, "would block"))
}
}
#[test]
fn queue_logic() {
// Create a socket with the queue size of 1.
let mut socket = WebSocket::from_raw_socket(
WouldBlockStreamMoc,
Role::Client,
Some(WebSocketConfig { max_send_queue: Some(1), ..Default::default() }),
);
// Test message that we're going to send.
let message = Message::Binary(vec![0xFF; 1024]);
// Helper to check the error.
let assert_would_block = |error| {
if let Error::Io(io_error) = error {
assert_eq!(io_error.kind(), io::ErrorKind::WouldBlock);
} else {
panic!("Expected WouldBlock error");
}
};
// The first attempt of writing must not fail, since the queue is empty at start.
// But since the underlying mock object always returns `WouldBlock`, so is the result.
assert_would_block(dbg!(socket.write_message(message.clone()).unwrap_err()));
// Any subsequent attempts must return an error telling that the queue is full.
for _i in 0..100 {
assert!(matches!(
socket.write_message(message.clone()).unwrap_err(),
Error::SendQueueFull(..)
));
}
// The size of the output buffer must not be bigger than the size of that message
// that we managed to write to the output buffer at first. Since we could not make
// any progress (because of the logic of the moc buffer), the size remains unchanged.
if socket.context.frame.output_buffer_len() > message.len() {
panic!("Too many frames in the queue");
}
}
#[test]
fn receive_messages() {
let incoming = Cursor::new(vec![
@ -754,10 +778,10 @@ mod tests {
0x03,
]);
let mut socket = WebSocket::from_raw_socket(WriteMoc(incoming), Role::Client, None);
assert_eq!(socket.read_message().unwrap(), Message::Ping(vec![1, 2]));
assert_eq!(socket.read_message().unwrap(), Message::Pong(vec![3]));
assert_eq!(socket.read_message().unwrap(), Message::Text("Hello, World!".into()));
assert_eq!(socket.read_message().unwrap(), Message::Binary(vec![0x01, 0x02, 0x03]));
assert_eq!(socket.read().unwrap(), Message::Ping(vec![1, 2]));
assert_eq!(socket.read().unwrap(), Message::Pong(vec![3]));
assert_eq!(socket.read().unwrap(), Message::Text("Hello, World!".into()));
assert_eq!(socket.read().unwrap(), Message::Binary(vec![0x01, 0x02, 0x03]));
}
#[test]
@ -770,7 +794,7 @@ mod tests {
let mut socket = WebSocket::from_raw_socket(WriteMoc(incoming), Role::Client, Some(limit));
assert!(matches!(
socket.read_message(),
socket.read(),
Err(Error::Capacity(CapacityError::MessageTooLong { size: 13, max_size: 10 }))
));
}
@ -782,7 +806,7 @@ mod tests {
let mut socket = WebSocket::from_raw_socket(WriteMoc(incoming), Role::Client, Some(limit));
assert!(matches!(
socket.read_message(),
socket.read(),
Err(Error::Capacity(CapacityError::MessageTooLong { size: 3, max_size: 2 }))
));
}

@ -104,11 +104,13 @@ mod encryption {
#[cfg(feature = "rustls-tls-native-roots")]
{
for cert in rustls_native_certs::load_native_certs()? {
root_store
.add(&rustls::Certificate(cert.0))
.map_err(TlsError::Rustls)?;
}
let native_certs = rustls_native_certs::load_native_certs()?;
let der_certs: Vec<Vec<u8>> =
native_certs.into_iter().map(|cert| cert.0).collect();
let total_number = der_certs.len();
let (number_added, number_ignored) =
root_store.add_parsable_certificates(&der_certs);
log::debug!("Added {number_added}/{total_number} native root certificates (ignored {number_ignored})");
}
#[cfg(feature = "rustls-tls-webpki-roots")]
{

@ -1,6 +1,7 @@
//! Verifies that the server returns a `ConnectionClosed` error when the connection
//! is closed from the server's point of view and drop the underlying tcp socket.
#![cfg(any(feature = "native-tls", feature = "__rustls-tls"))]
#![cfg(all(any(feature = "native-tls", feature = "__rustls-tls"), feature = "handshake"))]
use std::{
net::{TcpListener, TcpStream},
@ -51,27 +52,27 @@ fn test_server_close() {
do_test(
3012,
|mut cli_sock| {
cli_sock.write_message(Message::Text("Hello WebSocket".into())).unwrap();
cli_sock.send(Message::Text("Hello WebSocket".into())).unwrap();
let message = cli_sock.read_message().unwrap(); // receive close from server
let message = cli_sock.read().unwrap(); // receive close from server
assert!(message.is_close());
let err = cli_sock.read_message().unwrap_err(); // now we should get ConnectionClosed
let err = cli_sock.read().unwrap_err(); // now we should get ConnectionClosed
match err {
Error::ConnectionClosed => {}
_ => panic!("unexpected error: {:?}", err),
}
},
|mut srv_sock| {
let message = srv_sock.read_message().unwrap();
let message = srv_sock.read().unwrap();
assert_eq!(message.into_data(), b"Hello WebSocket");
srv_sock.close(None).unwrap(); // send close to client
let message = srv_sock.read_message().unwrap(); // receive acknowledgement
let message = srv_sock.read().unwrap(); // receive acknowledgement
assert!(message.is_close());
let err = srv_sock.read_message().unwrap_err(); // now we should get ConnectionClosed
let err = srv_sock.read().unwrap_err(); // now we should get ConnectionClosed
match err {
Error::ConnectionClosed => {}
_ => panic!("unexpected error: {:?}", err),
@ -85,26 +86,26 @@ fn test_evil_server_close() {
do_test(
3013,
|mut cli_sock| {
cli_sock.write_message(Message::Text("Hello WebSocket".into())).unwrap();
cli_sock.send(Message::Text("Hello WebSocket".into())).unwrap();
sleep(Duration::from_secs(1));
let message = cli_sock.read_message().unwrap(); // receive close from server
let message = cli_sock.read().unwrap(); // receive close from server
assert!(message.is_close());
let err = cli_sock.read_message().unwrap_err(); // now we should get ConnectionClosed
let err = cli_sock.read().unwrap_err(); // now we should get ConnectionClosed
match err {
Error::ConnectionClosed => {}
_ => panic!("unexpected error: {:?}", err),
}
},
|mut srv_sock| {
let message = srv_sock.read_message().unwrap();
let message = srv_sock.read().unwrap();
assert_eq!(message.into_data(), b"Hello WebSocket");
srv_sock.close(None).unwrap(); // send close to client
let message = srv_sock.read_message().unwrap(); // receive acknowledgement
let message = srv_sock.read().unwrap(); // receive acknowledgement
assert!(message.is_close());
// and now just drop the connection without waiting for `ConnectionClosed`
srv_sock.get_mut().set_linger(Some(Duration::from_secs(0))).unwrap();
@ -118,32 +119,32 @@ fn test_client_close() {
do_test(
3014,
|mut cli_sock| {
cli_sock.write_message(Message::Text("Hello WebSocket".into())).unwrap();
cli_sock.send(Message::Text("Hello WebSocket".into())).unwrap();
let message = cli_sock.read_message().unwrap(); // receive answer from server
let message = cli_sock.read().unwrap(); // receive answer from server
assert_eq!(message.into_data(), b"From Server");
cli_sock.close(None).unwrap(); // send close to server
let message = cli_sock.read_message().unwrap(); // receive acknowledgement from server
let message = cli_sock.read().unwrap(); // receive acknowledgement from server
assert!(message.is_close());
let err = cli_sock.read_message().unwrap_err(); // now we should get ConnectionClosed
let err = cli_sock.read().unwrap_err(); // now we should get ConnectionClosed
match err {
Error::ConnectionClosed => {}
_ => panic!("unexpected error: {:?}", err),
}
},
|mut srv_sock| {
let message = srv_sock.read_message().unwrap();
let message = srv_sock.read().unwrap();
assert_eq!(message.into_data(), b"Hello WebSocket");
srv_sock.write_message(Message::Text("From Server".into())).unwrap();
srv_sock.send(Message::Text("From Server".into())).unwrap();
let message = srv_sock.read_message().unwrap(); // receive close from client
let message = srv_sock.read().unwrap(); // receive close from client
assert!(message.is_close());
let err = srv_sock.read_message().unwrap_err(); // now we should get ConnectionClosed
let err = srv_sock.read().unwrap_err(); // now we should get ConnectionClosed
match err {
Error::ConnectionClosed => {}
_ => panic!("unexpected error: {:?}", err),

@ -1,6 +1,8 @@
//! Verifies that we can read data messages even if we have initiated a close handshake,
//! but before we got confirmation.
#![cfg(feature = "handshake")]
use std::{
net::TcpListener,
process::exit,
@ -8,7 +10,6 @@ use std::{
time::Duration,
};
#[cfg(feature = "handshake")]
use tungstenite::{accept, connect, error::ProtocolError, Error, Message};
use url::Url;
@ -28,10 +29,10 @@ fn test_no_send_after_close() {
let client_thread = spawn(move || {
let (mut client, _) = connect(Url::parse("ws://localhost:3013/socket").unwrap()).unwrap();
let message = client.read_message().unwrap(); // receive close from server
let message = client.read().unwrap(); // receive close from server
assert!(message.is_close());
let err = client.read_message().unwrap_err(); // now we should get ConnectionClosed
let err = client.read().unwrap_err(); // now we should get ConnectionClosed
match err {
Error::ConnectionClosed => {}
_ => panic!("unexpected error: {:?}", err),
@ -43,7 +44,7 @@ fn test_no_send_after_close() {
client_handler.close(None).unwrap(); // send close to client
let err = client_handler.write_message(Message::Text("Hello WebSocket".into()));
let err = client_handler.send(Message::Text("Hello WebSocket".into()));
assert!(err.is_err());

@ -1,6 +1,8 @@
//! Verifies that we can read data messages even if we have initiated a close handshake,
//! but before we got confirmation.
#![cfg(feature = "handshake")]
use std::{
net::TcpListener,
process::exit,
@ -8,7 +10,6 @@ use std::{
time::Duration,
};
#[cfg(feature = "handshake")]
use tungstenite::{accept, connect, Error, Message};
use url::Url;
@ -28,12 +29,12 @@ fn test_receive_after_init_close() {
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();
client.send(Message::Text("Hello WebSocket".into())).unwrap();
let message = client.read_message().unwrap(); // receive close from server
let message = client.read().unwrap(); // receive close from server
assert!(message.is_close());
let err = client.read_message().unwrap_err(); // now we should get ConnectionClosed
let err = client.read().unwrap_err(); // now we should get ConnectionClosed
match err {
Error::ConnectionClosed => {}
_ => panic!("unexpected error: {:?}", err),
@ -46,12 +47,12 @@ fn test_receive_after_init_close() {
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();
let message = client_handler.read().unwrap();
assert_eq!(message.into_data(), b"Hello WebSocket");
assert!(client_handler.read_message().unwrap().is_close()); // receive acknowledgement
assert!(client_handler.read().unwrap().is_close()); // receive acknowledgement
let err = client_handler.read_message().unwrap_err(); // now we should get ConnectionClosed
let err = client_handler.read().unwrap_err(); // now we should get ConnectionClosed
match err {
Error::ConnectionClosed => {}
_ => panic!("unexpected error: {:?}", err),

@ -0,0 +1,68 @@
use std::io::{self, Read, Write};
use tungstenite::{protocol::WebSocketConfig, Message, WebSocket};
/// `Write` impl that records call stats and drops the data.
#[derive(Debug, Default)]
struct MockWrite {
written_bytes: usize,
write_count: usize,
flush_count: usize,
}
impl Read for MockWrite {
fn read(&mut self, _: &mut [u8]) -> io::Result<usize> {
Err(io::Error::new(io::ErrorKind::WouldBlock, "reads not supported"))
}
}
impl Write for MockWrite {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.written_bytes += buf.len();
self.write_count += 1;
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
self.flush_count += 1;
Ok(())
}
}
/// Test for write buffering and flushing behaviour.
#[test]
fn write_flush_behaviour() {
const SEND_ME_LEN: usize = 10;
const BATCH_ME_LEN: usize = 11;
const WRITE_BUFFER_SIZE: usize = 600;
let mut ws = WebSocket::from_raw_socket(
MockWrite::default(),
tungstenite::protocol::Role::Server,
Some(WebSocketConfig { write_buffer_size: WRITE_BUFFER_SIZE, ..<_>::default() }),
);
assert_eq!(ws.get_ref().written_bytes, 0);
assert_eq!(ws.get_ref().write_count, 0);
assert_eq!(ws.get_ref().flush_count, 0);
// `send` writes & flushes immediately
ws.send(Message::Text("Send me!".into())).unwrap();
assert_eq!(ws.get_ref().written_bytes, SEND_ME_LEN);
assert_eq!(ws.get_ref().write_count, 1);
assert_eq!(ws.get_ref().flush_count, 1);
// send a batch of messages
for msg in (0..100).map(|_| Message::Text("Batch me!".into())) {
ws.write(msg).unwrap();
}
// after 55 writes the out_buffer will exceed write_buffer_size=600
// and so do a single underlying write (not flushing).
assert_eq!(ws.get_ref().written_bytes, 55 * BATCH_ME_LEN + SEND_ME_LEN);
assert_eq!(ws.get_ref().write_count, 2);
assert_eq!(ws.get_ref().flush_count, 1);
// flushing will perform a single write for the remaining out_buffer & flush.
ws.flush().unwrap();
assert_eq!(ws.get_ref().written_bytes, 100 * BATCH_ME_LEN + SEND_ME_LEN);
assert_eq!(ws.get_ref().write_count, 3);
assert_eq!(ws.get_ref().flush_count, 2);
}
Loading…
Cancel
Save