mercuryworkshop-mirror/epoxy-tls
1
0
Fork 0
mirror of https://github.com/MercuryWorkshop/epoxy-tls.git synced 2025-05-12 22:10:01 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #4 from MercuryWorkshop/wisp-v2

Browse source 
Wisp v2
This commit is contained in:
Toshit 2024-04-15 23:29:22 -07:00 • committed by GitHub
commit 1c22591817
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
26 changed files with 2403 additions and 459 deletions
Download patch file Download diff file Expand all files Collapse all files

759
Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

2
Cargo.toml
View file

@ -1,6 +1,6 @@
[workspace]
resolver = "2"
members = ["server", "client", "wisp", "simple-wisp-client"]
members = ["server", "client", "wisp", "simple-wisp-client", "certs-grabber"]
default-members = ["server"]
[profile.release]

13
certs-grabber/Cargo.toml Normal file
View file

@ -0,0 +1,13 @@
[package]
name = "certs-grabber"
version = "0.1.0"
edition = "2021"
[dependencies]
hex = "0.4.3"
ring = "0.17.8"
rustls-pki-types = "1.4.1"
rustls-webpki = "0.102.2"
tokio = { version = "1.37.0", features = ["full"] }
webpki-ccadb = "0.1.0"
x509-parser = "0.16.0"

64
certs-grabber/src/main.rs Normal file
View file

@ -0,0 +1,64 @@
use std::fmt::Write;
use ring::digest::{digest, SHA256};
use rustls_pki_types::{CertificateDer, TrustAnchor};
use webpki::anchor_from_trusted_cert;
use webpki_ccadb::fetch_ccadb_roots;
#[tokio::main]
async fn main() {
let tls_roots_map = fetch_ccadb_roots().await;
let mut code = String::with_capacity(256 * 1_024);
code.push_str("const ROOTS = [");
for (_, root) in tls_roots_map {
// Verify the DER FP matches the metadata FP.
let der = root.der();
let calculated_fp = digest(&SHA256, &der);
let metadata_fp = hex::decode(&root.sha256_fingerprint).expect("malformed fingerprint");
assert_eq!(calculated_fp.as_ref(), metadata_fp.as_slice());
let ta_der = CertificateDer::from(der.as_ref());
let TrustAnchor {
subject,
subject_public_key_info,
name_constraints,
} = anchor_from_trusted_cert(&ta_der).expect("malformed trust anchor der");
/*
let (_, parsed_cert) =
x509_parser::parse_x509_certificate(&der).expect("malformed x509 der");
let issuer = name_to_string(parsed_cert.issuer());
let subject_str = name_to_string(parsed_cert.subject());
let label = root.common_name_or_certificate_name.clone();
let serial = root.serial().to_string();
let sha256_fp = root.sha256_fp();
*/
code.write_fmt(format_args!(
"{{subject:new Uint8Array([{}]),subject_public_key_info:new Uint8Array([{}]),name_constraints:{}}},",
subject
.as_ref()
.iter()
.map(|x| x.to_string())
.collect::<Vec<String>>().join(","),
subject_public_key_info
.as_ref()
.iter()
.map(|x| x.to_string())
.collect::<Vec<String>>().join(","),
if let Some(constraints) = name_constraints {
format!("new Uint8Array([{}])",constraints
.as_ref()
.iter()
.map(|x| x.to_string())
.collect::<Vec<String>>().join(","))
} else {
"null".into()
}
))
.unwrap();
}
code.pop();
code.push_str("];");
println!("{}", code);
}

5
client/Cargo.toml
View file

@ -17,7 +17,6 @@ wasm-bindgen = { version = "0.2.91", features = ["enable-interning"] }
wasm-bindgen-futures = "0.4.39"
futures-util = "0.3.30"
js-sys = "0.3.66"
webpki-roots = "0.26.0"
tokio-rustls = "0.25.0"
web-sys = { version = "0.3.66", features = ["Request", "RequestInit", "Headers", "Response", "ResponseInit", "WebSocket", "BinaryType", "MessageEvent"] }
wasm-streams = "0.4.0"
@ -25,7 +24,7 @@ tokio-util = { version = "0.7.10", features = ["io"] }
async-compression = { version = "0.4.5", features = ["tokio", "gzip", "brotli"] }
fastwebsockets = { version = "0.6.0", features = ["unstable-split"] }
base64 = "0.21.7"
wisp-mux = { path = "../wisp", features = ["tokio_io"] }
wisp-mux = { path = "../wisp", features = ["tokio_io", "wasm"] }
async_io_stream = { version = "0.3.3", features = ["tokio_io"] }
getrandom = { version = "0.2.12", features = ["js"] }
hyper-util-wasm = { version = "0.1.3", features = ["client", "client-legacy", "http1", "http2"] }
@ -35,6 +34,7 @@ console_error_panic_hook = "0.1.7"
send_wrapper = "0.6.0"
event-listener = "5.2.0"
wasmtimer = "0.2.0"
async-trait = "0.1.80"
[dependencies.ring]
features = ["wasm32_unknown_unknown_js"]
@ -46,3 +46,4 @@ features = ["web"]
default-env = "0.1.1"
wasm-bindgen-test = "0.3.42"
web-sys = { version = "0.3.69", features = ["FormData", "UrlSearchParams"] }
webpki-roots = "0.26.0"

9
client/build.sh
View file

@ -41,5 +41,14 @@ echo "}\ndeclare function epoxy(maybe_memory?: WebAssembly.Memory): Promise<type
cp out/epoxy_client.d.ts pkg/epoxy.d.ts
cp out/epoxy_client_bg.wasm pkg/epoxy.wasm
echo "[epx] fetching certs"
(
cd ../certs-grabber
cargo run
) > pkg/certs.js
cat pkg/certs.js > pkg/certs-module.js
echo "export default ROOTS;" >> pkg/certs-module.js
echo "[epx] fetching certs finished"
rm -r out/
echo "[epx] done!"

13
client/demo.js
View file

@ -1,4 +1,5 @@
import epoxy from "./pkg/epoxy-module-bundled.js";
import CERTS from "./pkg/certs-module.js";
onmessage = async (msg) => {
console.debug("recieved demo:", msg);
@ -29,13 +30,13 @@ onmessage = async (msg) => {
postMessage(JSON.stringify(str, null, 4));
}
const { EpoxyClient, certs } = await epoxy();
const { EpoxyClient } = await epoxy();
console.log("certs:", certs());
console.log("certs:", CERTS);
const tconn0 = performance.now();
// args: websocket url, user agent, redirect limit
let epoxy_client = await new EpoxyClient("ws://localhost:4000", navigator.userAgent, 10);
// args: websocket url, user agent, redirect limit, certs
let epoxy_client = await new EpoxyClient("ws://localhost:4000", navigator.userAgent, 10, CERTS);
const tconn1 = performance.now();
log(`conn establish took ${tconn1 - tconn0} ms or ${(tconn1 - tconn0) / 1000} s`);
@ -237,9 +238,9 @@ onmessage = async (msg) => {
log(`total avg mux (${num_outer_tests} tests of ${num_inner_tests} reqs): ${total_mux_multi} ms or ${total_mux_multi / 1000} s`);
} else {
let resp = await epoxy_client.fetch("https://httpbin.org/get");
let resp = await epoxy_client.fetch("https://www.example.com/");
console.log(resp, Object.fromEntries(resp.headers));
plog(await resp.json());
log(await resp.text());
}
log("done");
};

46
client/src/lib.rs
View file

@ -10,6 +10,7 @@ mod wrappers;
use tls_stream::EpxTlsStream;
use tokioio::TokioIo;
use udp_stream::EpxUdpStream;
use utils::object_to_trustanchor;
pub use utils::{Boolinator, ReplaceErr, UriExt};
use websocket::EpxWebSocket;
use wrappers::{IncomingBody, ServiceWrapper, TlsWispService, WebSocketWrapper};
@ -70,42 +71,10 @@ fn init() {
intern("Content-Type");
}
fn cert_to_jval(cert: &TrustAnchor) -> Result<JsValue, JsValue> {
let val = Object::new();
Reflect::set(
&val,
&jval!("subject"),
&Uint8Array::from(cert.subject.as_ref()),
)?;
Reflect::set(
&val,
&jval!("subject_public_key_info"),
&Uint8Array::from(cert.subject_public_key_info.as_ref()),
)?;
Reflect::set(
&val,
&jval!("name_constraints"),
&jval!(cert
.name_constraints
.as_ref()
.map(|x| Uint8Array::from(x.as_ref()))),
)?;
Ok(val.into())
}
#[wasm_bindgen]
pub fn certs() -> Result<JsValue, JsValue> {
Ok(webpki_roots::TLS_SERVER_ROOTS
.iter()
.map(cert_to_jval)
.collect::<Result<Array, JsValue>>()?
.into())
}
#[wasm_bindgen(inspectable)]
pub struct EpoxyClient {
rustls_config: Arc<rustls::ClientConfig>,
mux: Arc<RwLock<ClientMux<WebSocketWrapper>>>,
mux: Arc<RwLock<ClientMux>>,
hyper_client: Client<TlsWispService, HttpBody>,
#[wasm_bindgen(getter_with_clone)]
pub useragent: String,
@ -120,6 +89,7 @@ impl EpoxyClient {
ws_url: String,
useragent: String,
redirect_limit: usize,
certs: Array,
) -> Result<EpoxyClient, JsError> {
let ws_uri = ws_url
.parse::<uri::Uri>()
@ -137,7 +107,13 @@ impl EpoxyClient {
utils::spawn_mux_fut(mux.clone(), fut, ws_url.clone());
let mut certstore = RootCertStore::empty();
certstore.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
let certs: Result<Vec<TrustAnchor>, JsValue> =
certs.iter().map(object_to_trustanchor).collect();
certstore.extend(
certs
.replace_err("Failed to get certificates from cert store")?
.into_iter(),
);
let rustls_config = Arc::new(
rustls::ClientConfig::builder()
@ -164,7 +140,7 @@ impl EpoxyClient {
async fn get_tls_io(&self, url_host: &str, url_port: u16) -> Result<EpxIoTlsStream, JsError> {
let channel = self
.mux
.read()
.write()
.await
.client_new_stream(StreamType::Tcp, url_host.to_string(), url_port)
.await

2
client/src/udp_stream.rs
View file

@ -33,7 +33,7 @@ impl EpxUdpStream {
let io = tcp
.mux
.read()
.write()
.await
.client_new_stream(StreamType::Udp, url_host.to_string(), url_port)
.await

42
client/src/utils.rs
View file

@ -1,12 +1,13 @@
use crate::*;
use rustls_pki_types::Der;
use wasm_bindgen::prelude::*;
use wasm_bindgen_futures::JsFuture;
use hyper::rt::Executor;
use js_sys::ArrayBuffer;
use std::future::Future;
use wisp_mux::WispError;
use wisp_mux::{extensions::udp::UdpProtocolExtensionBuilder, WispError};
#[wasm_bindgen]
extern "C" {
@ -194,26 +195,24 @@ pub async fn make_mux(
url: &str,
) -> Result<
(
ClientMux<WebSocketWrapper>,
impl Future<Output = Result<(), WispError>>,
ClientMux,
impl Future<Output = Result<(), WispError>> + Send,
),
WispError,
> {
let (wtx, wrx) = WebSocketWrapper::connect(url, vec![])
.await
.map_err(|_| WispError::WsImplSocketClosed)?;
let (wtx, wrx) =
WebSocketWrapper::connect(url, vec![]).map_err(|_| WispError::WsImplSocketClosed)?;
wtx.wait_for_open().await;
let mux = ClientMux::new(wrx, wtx).await?;
Ok(mux)
ClientMux::new(wrx, wtx, Some(&[Box::new(UdpProtocolExtensionBuilder())])).await
}
pub fn spawn_mux_fut(
mux: Arc<RwLock<ClientMux<WebSocketWrapper>>>,
fut: impl Future<Output = Result<(), WispError>> + 'static,
mux: Arc<RwLock<ClientMux>>,
fut: impl Future<Output = Result<(), WispError>> + Send + 'static,
url: String,
) {
wasm_bindgen_futures::spawn_local(async move {
debug!("epoxy: mux future started");
if let Err(e) = fut.await {
log!("epoxy: error in mux future, restarting: {:?}", e);
while let Err(e) = replace_mux(mux.clone(), &url).await {
@ -225,13 +224,10 @@ pub fn spawn_mux_fut(
});
}
pub async fn replace_mux(
mux: Arc<RwLock<ClientMux<WebSocketWrapper>>>,
url: &str,
) -> Result<(), WispError> {
pub async fn replace_mux(mux: Arc<RwLock<ClientMux>>, url: &str) -> Result<(), WispError> {
let (mux_replace, fut) = make_mux(url).await?;
let mut mux_write = mux.write().await;
mux_write.close().await?;
let _ = mux_write.close().await;
*mux_write = mux_replace;
drop(mux_write);
spawn_mux_fut(mux, fut, url.into());
@ -264,3 +260,17 @@ pub async fn jval_to_u8_array_req(val: JsValue) -> Result<(Uint8Array, web_sys::
req,
))
}
pub fn object_to_trustanchor(obj: JsValue) -> Result<TrustAnchor<'static>, JsValue> {
let subject: Uint8Array = Reflect::get(&obj, &jval!("subject"))?.dyn_into()?;
let pub_key_info: Uint8Array =
Reflect::get(&obj, &jval!("subject_public_key_info"))?.dyn_into()?;
let name_constraints: Option<Uint8Array> = Reflect::get(&obj, &jval!("name_constraints"))
.and_then(|x| x.dyn_into())
.ok();
Ok(TrustAnchor {
subject: Der::from(subject.to_vec()),
subject_public_key_info: Der::from(pub_key_info.to_vec()),
name_constraints: name_constraints.map(|x| Der::from(x.to_vec())),
})
}

10
client/src/websocket.rs
View file

@ -106,7 +106,7 @@ impl EpxWebSocket {
break;
}
// ping/pong/continue
_ => {},
_ => {}
}
}
});
@ -115,7 +115,13 @@ impl EpxWebSocket {
.call0(&Object::default())
.replace_err("Failed to call onopen")?;
Ok(Self { tx, onerror, origin, protocols, url: url.to_string() })
Ok(Self {
tx,
onerror,
origin,
protocols,
url: url.to_string(),
})
}
.await;
if let Err(ret) = ret {

24
client/src/wrappers.rs
View file

@ -53,7 +53,7 @@ impl Stream for IncomingBody {
}
#[derive(Clone)]
pub struct ServiceWrapper(pub Arc<RwLock<ClientMux<WebSocketWrapper>>>, pub String);
pub struct ServiceWrapper(pub Arc<RwLock<ClientMux>>, pub String);
impl tower_service::Service<hyper::Uri> for ServiceWrapper {
type Response = TokioIo<EpxIoUnencryptedStream>;
@ -69,7 +69,7 @@ impl tower_service::Service<hyper::Uri> for ServiceWrapper {
let mux_url = self.1.clone();
async move {
let stream = mux
.read()
.write()
.await
.client_new_stream(
StreamType::Tcp,
@ -143,6 +143,7 @@ impl tower_service::Service<hyper::Uri> for TlsWispService {
pub enum WebSocketError {
Unknown,
SendFailed,
CloseFailed,
}
impl std::fmt::Display for WebSocketError {
@ -151,6 +152,7 @@ impl std::fmt::Display for WebSocketError {
match self {
Unknown => write!(f, "Unknown error"),
SendFailed => write!(f, "Send failed"),
CloseFailed => write!(f, "Close failed"),
}
}
}
@ -193,11 +195,9 @@ pub struct WebSocketReader {
close_event: Arc<Event>,
}
#[async_trait::async_trait]
impl WebSocketRead for WebSocketReader {
async fn wisp_read_frame(
&mut self,
_: &LockedWebSocketWrite<impl WebSocketWrite>,
) -> Result<Frame, WispError> {
async fn wisp_read_frame(&mut self, _: &LockedWebSocketWrite) -> Result<Frame, WispError> {
use WebSocketMessage::*;
if self.closed.load(Ordering::Acquire) {
return Err(WispError::WsImplSocketClosed);
@ -215,10 +215,7 @@ impl WebSocketRead for WebSocketReader {
}
impl WebSocketWrapper {
pub async fn connect(
url: &str,
protocols: Vec<String>,
) -> Result<(Self, WebSocketReader), JsValue> {
pub fn connect(url: &str, protocols: Vec<String>) -> Result<(Self, WebSocketReader), JsValue> {
let (read_tx, read_rx) = mpsc::unbounded_channel();
let closed = Arc::new(AtomicBool::new(false));
@ -306,6 +303,7 @@ impl WebSocketWrapper {
}
}
#[async_trait::async_trait]
impl WebSocketWrite for WebSocketWrapper {
async fn wisp_write_frame(&mut self, frame: Frame) -> Result<(), WispError> {
use wisp_mux::ws::OpCode::*;
@ -328,6 +326,12 @@ impl WebSocketWrite for WebSocketWrapper {
_ => Err(WispError::WsImplNotSupported),
}
}
async fn wisp_close(&mut self) -> Result<(), WispError> {
self.inner
.close()
.map_err(|_| WebSocketError::CloseFailed.into())
}
}
impl Drop for WebSocketWrapper {

32
client/tests/fetch.rs
View file

@ -1,6 +1,7 @@
use default_env::default_env;
use epoxy_client::EpoxyClient;
use js_sys::{JsString, Object, Reflect, Uint8Array, JSON};
use js_sys::{Array, JsString, Object, Reflect, Uint8Array, JSON};
use rustls_pki_types::TrustAnchor;
use tokio::sync::OnceCell;
use wasm_bindgen::JsValue;
use wasm_bindgen_futures::JsFuture;
@ -12,11 +13,40 @@ wasm_bindgen_test_configure!(run_in_dedicated_worker);
static USER_AGENT: &str = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36";
static EPOXY_CLIENT: OnceCell<EpoxyClient> = OnceCell::const_new();
pub fn trustanchor_to_object(cert: &TrustAnchor) -> Result<JsValue, JsValue> {
let val = Object::new();
Reflect::set(
&val,
&JsValue::from("subject"),
&Uint8Array::from(cert.subject.as_ref()),
)?;
Reflect::set(
&val,
&JsValue::from("subject_public_key_info"),
&Uint8Array::from(cert.subject_public_key_info.as_ref()),
)?;
Reflect::set(
&val,
&JsValue::from("name_constraints"),
&JsValue::from(
cert.name_constraints
.as_ref()
.map(|x| Uint8Array::from(x.as_ref())),
),
)?;
Ok(val.into())
}
async fn get_client_w_ua(useragent: &str, redirect_limit: usize) -> EpoxyClient {
EpoxyClient::new(
"ws://localhost:4000".into(),
useragent.into(),
redirect_limit,
webpki_roots::TLS_SERVER_ROOTS
.iter()
.map(trustanchor_to_object)
.collect::<Result<Array, JsValue>>()
.expect("Failed to create certs"),
)
.await
.ok()

1
rustfmt.toml Normal file
View file

@ -0,0 +1 @@
imports_granularity = "Crate"

150
server/src/main.rs
View file

@ -1,26 +1,35 @@
#![feature(let_chains, ip)]
use std::io::Error;
use std::{collections::HashMap, io::Error, path::PathBuf, sync::Arc};
use bytes::Bytes;
use clap::Parser;
use fastwebsockets::{
upgrade::{self, UpgradeFut}, CloseCode, FragmentCollector, FragmentCollectorRead, Frame, OpCode, Payload,
WebSocketError,
upgrade::{self, UpgradeFut},
CloseCode, FragmentCollector, FragmentCollectorRead, Frame, OpCode, Payload, WebSocketError,
};
use futures_util::{SinkExt, StreamExt, TryFutureExt};
use hyper::{
body::Incoming, server::conn::http1, service::service_fn, Request, Response,
StatusCode,
body::Incoming, server::conn::http1, service::service_fn, Request, Response, StatusCode,
};
use hyper_util::rt::TokioIo;
use tokio::net::{lookup_host, TcpListener, TcpStream, UdpSocket};
#[cfg(unix)]
use tokio::net::{UnixListener, UnixStream};
use tokio::{
io::copy_bidirectional,
net::{lookup_host, TcpListener, TcpStream, UdpSocket},
};
use tokio_util::codec::{BytesCodec, Framed};
#[cfg(unix)]
use tokio_util::either::Either;
use wisp_mux::{CloseReason, ConnectPacket, MuxStream, ServerMux, StreamType, WispError};
use wisp_mux::{
extensions::{
password::{PasswordProtocolExtension, PasswordProtocolExtensionBuilder},
udp::UdpProtocolExtensionBuilder,
ProtocolExtensionBuilder,
},
CloseReason, ConnectPacket, MuxStream, ServerMux, StreamType, WispError,
};
type HttpBody = http_body_util::Full<hyper::body::Bytes>;
@ -54,6 +63,20 @@ struct Cli {
/// Whether the server should block ports other than 80 or 443
#[arg(long)]
block_non_http: bool,
/// Path to a file containing `user:password` separated by newlines. This is plaintext!!!
///
/// `user` cannot contain `:`. Whitespace will be trimmed.
#[arg(long)]
auth: Option<PathBuf>,
}
#[derive(Clone)]
struct MuxOptions {
pub block_local: bool,
pub block_udp: bool,
pub block_non_http: bool,
pub enforce_auth: bool,
pub auth: Arc<Vec<Box<(dyn ProtocolExtensionBuilder + Send + Sync)>>>,
}
#[cfg(not(unix))]
@ -136,19 +159,47 @@ async fn main() -> Result<(), Error> {
"/".to_string()
};
let mut auth = HashMap::new();
let enforce_auth = opt.auth.is_some();
if let Some(file) = opt.auth {
let file = std::fs::read_to_string(file)?;
for entry in file.split('\n').filter_map(|x| {
if x.contains(':') {
Some(x.trim())
} else {
None
}
}) {
let split: Vec<_> = entry.split(':').collect();
let username = split[0];
let password = split[1..].join(":");
println!(
"adding username {:?} password {:?} to allowed auth",
username, password
);
auth.insert(username.to_string(), password.to_string());
}
}
let pw_ext = PasswordProtocolExtensionBuilder::new_server(auth);
let mux_options = MuxOptions {
block_local: opt.block_local,
block_non_http: opt.block_non_http,
block_udp: opt.block_udp,
auth: Arc::new(vec![
Box::new(UdpProtocolExtensionBuilder()),
Box::new(pw_ext),
]),
enforce_auth,
};
println!("listening on `{}` with prefix `{}`", addr, prefix);
while let Ok((stream, addr)) = socket.accept().await {
let prefix = prefix.clone();
let mux_options = mux_options.clone();
tokio::spawn(async move {
let service = service_fn(move |res| {
accept_http(
res,
addr.clone(),
prefix.clone(),
opt.block_local,
opt.block_udp,
opt.block_non_http,
)
accept_http(res, addr.clone(), prefix.clone(), mux_options.clone())
});
let conn = http1::Builder::new()
.serve_connection(TokioIo::new(stream), service)
@ -166,9 +217,7 @@ async fn accept_http(
mut req: Request<Incoming>,
addr: String,
prefix: String,
block_local: bool,
block_udp: bool,
block_non_http: bool,
mux_options: MuxOptions,
) -> Result<Response<HttpBody>, WebSocketError> {
let uri = req.uri().path().to_string();
if upgrade::is_upgrade_request(&req)
@ -177,12 +226,17 @@ async fn accept_http(
let (res, fut) = upgrade::upgrade(&mut req)?;
if uri.is_empty() {
tokio::spawn(async move {
accept_ws(fut, addr.clone(), block_local, block_udp, block_non_http).await
});
tokio::spawn(async move { accept_ws(fut, addr.clone(), mux_options).await });
} else if let Some(uri) = uri.strip_prefix('/').map(|x| x.to_string()) {
tokio::spawn(async move {
accept_wsproxy(fut, uri, addr.clone(), block_local, block_non_http).await
accept_wsproxy(
fut,
uri,
addr.clone(),
mux_options.block_local,
mux_options.block_non_http,
)
.await
});
}
@ -210,7 +264,7 @@ async fn handle_mux(packet: ConnectPacket, mut stream: MuxStream) -> Result<bool
.await
.map_err(|x| WispError::Other(Box::new(x)))?;
let mut mux_stream = stream.into_io().into_asyncrw();
tokio::io::copy_bidirectional(&mut tcp_stream, &mut mux_stream)
copy_bidirectional(&mut mux_stream, &mut tcp_stream)
.await
.map_err(|x| WispError::Other(Box::new(x)))?;
}
@ -245,6 +299,10 @@ async fn handle_mux(packet: ConnectPacket, mut stream: MuxStream) -> Result<bool
}
}
}
StreamType::Unknown(_) => {
stream.close(CloseReason::ServerStreamInvalidInfo).await?;
return Ok(false);
}
}
Ok(true)
}
@ -252,16 +310,43 @@ async fn handle_mux(packet: ConnectPacket, mut stream: MuxStream) -> Result<bool
async fn accept_ws(
ws: UpgradeFut,
addr: String,
block_local: bool,
block_non_http: bool,
block_udp: bool,
mux_options: MuxOptions,
) -> Result<(), Box<dyn std::error::Error + Sync + Send>> {
let (rx, tx) = ws.await?.split(tokio::io::split);
let rx = FragmentCollectorRead::new(rx);
println!("{:?}: connected", addr);
// to prevent memory ""leaks"" because users are sending in packets way too fast the buffer
// size is set to 128
let (mut mux, fut) = if mux_options.enforce_auth {
let (mut mux, fut) = ServerMux::new(rx, tx, 128, Some(mux_options.auth.as_slice())).await?;
if !mux
.supported_extension_ids
.iter()
.any(|x| *x == PasswordProtocolExtension::ID)
{
println!(
"{:?}: client did not support auth or password was invalid",
addr
);
mux.close_extension_incompat().await?;
return Ok(());
}
(mux, fut)
} else {
ServerMux::new(
rx,
tx,
128,
Some(&[Box::new(UdpProtocolExtensionBuilder())]),
)
.await?
};
let (mut mux, fut) = ServerMux::new(rx, tx, u32::MAX);
println!(
"{:?}: downgraded: {} extensions supported: {:?}",
addr, mux.downgraded, mux.supported_extension_ids
);
tokio::spawn(async move {
if let Err(e) = fut.await {
@ -271,14 +356,14 @@ async fn accept_ws(
while let Some((packet, mut stream)) = mux.server_new_stream().await {
tokio::spawn(async move {
if (block_non_http
if (mux_options.block_non_http
&& !(packet.destination_port == 80 || packet.destination_port == 443))
|| (block_udp && packet.stream_type == StreamType::Udp)
|| (mux_options.block_udp && packet.stream_type == StreamType::Udp)
{
let _ = stream.close(CloseReason::ServerStreamBlockedAddress).await;
return;
}
if block_local {
if mux_options.block_local {
match lookup_host(format!(
"{}:{}",
packet.destination_hostname, packet.destination_port
@ -310,10 +395,9 @@ async fn accept_ws(
})
.and_then(|should_send| async move {
if should_send {
close_ok.close(CloseReason::Voluntary).await
} else {
Ok(())
let _ = close_ok.close(CloseReason::Voluntary).await;
}
Ok(())
})
.await;
});

92
simple-wisp-client/src/main.rs
View file

@ -11,7 +11,13 @@ use hyper::{
};
use simple_moving_average::{SingleSumSMA, SMA};
use std::{
error::Error, future::Future, io::{stdout, IsTerminal, Write}, net::SocketAddr, process::exit, sync::Arc, time::{Duration, Instant}, usize
error::Error,
future::Future,
io::{stdout, IsTerminal, Write},
net::SocketAddr,
process::exit,
sync::Arc,
time::{Duration, Instant},
};
use tokio::{
net::TcpStream,
@ -21,7 +27,14 @@ use tokio::{
};
use tokio_native_tls::{native_tls, TlsConnector};
use tokio_util::either::Either;
use wisp_mux::{ClientMux, StreamType, WispError};
use wisp_mux::{
extensions::{
password::{PasswordProtocolExtension, PasswordProtocolExtensionBuilder},
udp::{UdpProtocolExtension, UdpProtocolExtensionBuilder},
ProtocolExtensionBuilder,
},
ClientMux, StreamType, WispError,
};
#[derive(Debug)]
enum WispClientError {
@ -71,6 +84,17 @@ struct Cli {
/// Duration to run the test for
#[arg(short, long)]
duration: Option<humantime::Duration>,
/// Ask for UDP
#[arg(short, long)]
udp: bool,
/// Enable auth: format is `username:password`
///
/// Usernames and passwords are sent in plaintext!!
#[arg(long)]
auth: Option<String>,
/// Make a Wisp V1 connection
#[arg(long)]
wisp_v1: bool,
}
#[tokio::main(flavor = "multi_thread")]
@ -94,6 +118,13 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
let addr_dest = opts.tcp.ip().to_string();
let addr_dest_port = opts.tcp.port();
let auth = opts.auth.map(|auth| {
let split: Vec<_> = auth.split(':').collect();
let username = split[0].to_string();
let password = split[1..].join(":");
PasswordProtocolExtensionBuilder::new_client(username, password)
});
println!(
"connecting to {} and sending &[0; 1024 * {}] to {} with threads {}",
opts.wisp, opts.packet_size, opts.tcp, opts.streams,
@ -117,7 +148,6 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
fastwebsockets::handshake::generate_key(),
)
.header("Sec-WebSocket-Version", "13")
.header("Sec-WebSocket-Protocol", "wisp-v1")
.body(Empty::<Bytes>::new())?;
let (ws, _) = handshake::client(&SpawnExecutor, req, socket).await?;
@ -125,7 +155,53 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
let (rx, tx) = ws.split(tokio::io::split);
let rx = FragmentCollectorRead::new(rx);
let (mux, fut) = ClientMux::new(rx, tx).await?;
let mut extensions: Vec<Box<(dyn ProtocolExtensionBuilder + Send + Sync)>> = Vec::new();
if opts.udp {
extensions.push(Box::new(UdpProtocolExtensionBuilder()));
}
let enforce_auth = auth.is_some();
if let Some(auth) = auth {
extensions.push(Box::new(auth));
}
let (mut mux, fut) = if opts.wisp_v1 {
ClientMux::new(rx, tx, None).await?
} else {
ClientMux::new(rx, tx, Some(extensions.as_slice())).await?
};
if opts.udp
&& !mux
.supported_extension_ids
.iter()
.any(|x| *x == UdpProtocolExtension::ID)
{
println!(
"server did not support udp, was downgraded {}, extensions supported {:?}",
mux.downgraded, mux.supported_extension_ids
);
mux.close_extension_incompat().await?;
exit(1);
}
if enforce_auth
&& !mux
.supported_extension_ids
.iter()
.any(|x| *x == PasswordProtocolExtension::ID)
{
println!(
"server did not support passwords or password was incorrect, was downgraded {}, extensions supported {:?}",
mux.downgraded, mux.supported_extension_ids
);
mux.close_extension_incompat().await?;
exit(1);
}
println!(
"connected and created ClientMux, was downgraded {}, extensions supported {:?}",
mux.downgraded, mux.supported_extension_ids
);
let mut threads = Vec::with_capacity(opts.streams * 2 + 3);
threads.push(tokio::spawn(fut));
@ -177,7 +253,7 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
avg.get_average() * opts.packet_size,
);
if is_term {
print!("\x1b[2K{}\r", stat);
println!("\x1b[1A\x1b[2K{}\r", stat);
} else {
println!("{}", stat);
}
@ -208,6 +284,8 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
let out = select_all(threads.into_iter()).await;
let duration_since = Instant::now().duration_since(start_time);
if let Err(err) = out.0? {
println!("\n\nerr: {:?}", err);
exit(1);
@ -215,10 +293,10 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
out.2.into_iter().for_each(|x| x.abort());
let duration_since = Instant::now().duration_since(start_time);
mux.close().await?;
println!(
"\n\nresults: {} packets of &[0; 1024 * {}] ({} KiB) sent in {} ({} KiB/s)",
"\nresults: {} packets of &[0; 1024 * {}] ({} KiB) sent in {} ({} KiB/s)",
cnt.get(),
opts.packet_size,
cnt.get() * opts.packet_size,

6
wisp/Cargo.toml
View file

@ -1,6 +1,6 @@
[package]
name = "wisp-mux"
version = "3.0.0"
version = "4.0.0"
license = "LGPL-3.0-only"
description = "A library for easily creating Wisp servers and clients."
homepage = "https://github.com/MercuryWorkshop/epoxy-tls/tree/multiplexed/wisp"
@ -9,12 +9,15 @@ readme = "README.md"
edition = "2021"
[dependencies]
async-trait = "0.1.79"
async_io_stream = "0.3.3"
bytes = "1.5.0"
dashmap = { version = "5.5.3", features = ["inline"] }
event-listener = "5.0.0"
fastwebsockets = { version = "0.7.1", features = ["unstable-split"], optional = true }
flume = "0.11.0"
futures = "0.3.30"
futures-timer = "3.0.3"
futures-util = "0.3.30"
pin-project-lite = "0.2.13"
tokio = { version = "1.35.1", optional = true, default-features = false }
@ -22,6 +25,7 @@ tokio = { version = "1.35.1", optional = true, default-features = false }
[features]
fastwebsockets = ["dep:fastwebsockets", "dep:tokio"]
tokio_io = ["async_io_stream/tokio_io"]
wasm = ["futures-timer/wasm-bindgen"]
[package.metadata.docs.rs]
all-features = true

106
wisp/src/extensions/mod.rs Normal file
View file

@ -0,0 +1,106 @@
//! Wisp protocol extensions.
pub mod password;
pub mod udp;
use std::ops::{Deref, DerefMut};
use async_trait::async_trait;
use bytes::{BufMut, Bytes, BytesMut};
use crate::{
ws::{LockedWebSocketWrite, WebSocketRead},
Role, WispError,
};
/// Type-erased protocol extension that implements Clone.
#[derive(Debug)]
pub struct AnyProtocolExtension(Box<dyn ProtocolExtension + Sync + Send>);
impl AnyProtocolExtension {
/// Create a new type-erased protocol extension.
pub fn new<T: ProtocolExtension + Sync + Send + 'static>(extension: T) -> Self {
Self(Box::new(extension))
}
}
impl Deref for AnyProtocolExtension {
type Target = dyn ProtocolExtension;
fn deref(&self) -> &Self::Target {
self.0.deref()
}
}
impl DerefMut for AnyProtocolExtension {
fn deref_mut(&mut self) -> &mut Self::Target {
self.0.deref_mut()
}
}
impl Clone for AnyProtocolExtension {
fn clone(&self) -> Self {
Self(self.0.box_clone())
}
}
impl From<AnyProtocolExtension> for Bytes {
fn from(value: AnyProtocolExtension) -> Self {
let mut bytes = BytesMut::with_capacity(5);
let payload = value.encode();
bytes.put_u8(value.get_id());
bytes.put_u32_le(payload.len() as u32);
bytes.extend(payload);
bytes.freeze()
}
}
/// A Wisp protocol extension.
///
/// See [the
/// docs](https://github.com/MercuryWorkshop/wisp-protocol/blob/main/protocol.md#protocol-extensions).
#[async_trait]
pub trait ProtocolExtension: std::fmt::Debug {
/// Get the protocol extension ID.
fn get_id(&self) -> u8;
/// Get the protocol extension's supported packets.
///
/// Used to decide whether to call the protocol extension's packet handler.
fn get_supported_packets(&self) -> &'static [u8];
/// Encode self into Bytes.
fn encode(&self) -> Bytes;
/// Handle the handshake part of a Wisp connection.
///
/// This should be used to send or receive data before any streams are created.
async fn handle_handshake(
&mut self,
read: &mut dyn WebSocketRead,
write: &LockedWebSocketWrite,
) -> Result<(), WispError>;
/// Handle receiving a packet.
async fn handle_packet(
&mut self,
packet: Bytes,
read: &mut dyn WebSocketRead,
write: &LockedWebSocketWrite,
) -> Result<(), WispError>;
/// Clone the protocol extension.
fn box_clone(&self) -> Box<dyn ProtocolExtension + Sync + Send>;
}
/// Trait to build a Wisp protocol extension from a payload.
pub trait ProtocolExtensionBuilder {
/// Get the protocol extension ID.
///
/// Used to decide whether this builder should be used.
fn get_id(&self) -> u8;
/// Build a protocol extension from the extension's metadata.
fn build_from_bytes(&self, bytes: Bytes, role: Role)
-> Result<AnyProtocolExtension, WispError>;
/// Build a protocol extension to send to the other side.
fn build_to_extension(&self, role: Role) -> AnyProtocolExtension;
}

276
wisp/src/extensions/password.rs Normal file
View file

@ -0,0 +1,276 @@
//! Password protocol extension.
//!
//! Passwords are sent in plain text!!
//!
//! # Example
//! Server:
//! ```
//! let mut passwords = HashMap::new();
//! passwords.insert("user1".to_string(), "pw".to_string());
//! let (mux, fut) = ServerMux::new(
//! rx,
//! tx,
//! 128,
//! Some(&[Box::new(PasswordProtocolExtensionBuilder::new_server(passwords))])
//! );
//! ```
//!
//! Client:
//! ```
//! let (mux, fut) = ClientMux::new(
//! rx,
//! tx,
//! 128,
//! Some(&[
//! Box::new(PasswordProtocolExtensionBuilder::new_client(
//! "user1".to_string(),
//! "pw".to_string()
//! ))
//! ])
//! );
//! ```
//! See [the docs](https://github.com/MercuryWorkshop/wisp-protocol/blob/main/protocol.md#0x02---password-authentication)
use std::{collections::HashMap, error::Error, fmt::Display, string::FromUtf8Error};
use async_trait::async_trait;
use bytes::{Buf, BufMut, Bytes, BytesMut};
use crate::{
ws::{LockedWebSocketWrite, WebSocketRead},
Role, WispError,
};
use super::{AnyProtocolExtension, ProtocolExtension, ProtocolExtensionBuilder};
#[derive(Debug, Clone)]
/// Password protocol extension.
///
/// **Passwords are sent in plain text!!**
/// **This extension will panic when encoding if the username's length does not fit within a u8
/// or the password's length does not fit within a u16.**
pub struct PasswordProtocolExtension {
/// The username to log in with.
///
/// This string's length must fit within a u8.
pub username: String,
/// The password to log in with.
///
/// This string's length must fit within a u16.
pub password: String,
role: Role,
}
impl PasswordProtocolExtension {
/// Password protocol extension ID.
pub const ID: u8 = 0x02;
/// Create a new password protocol extension for the server.
///
/// This signifies that the server requires a password.
pub fn new_server() -> Self {
Self {
username: String::new(),
password: String::new(),
role: Role::Server,
}
}
/// Create a new password protocol extension for the client, with a username and password.
///
/// The username's length must fit within a u8. The password's length must fit within a
/// u16.
pub fn new_client(username: String, password: String) -> Self {
Self {
username,
password,
role: Role::Client,
}
}
}
#[async_trait]
impl ProtocolExtension for PasswordProtocolExtension {
fn get_id(&self) -> u8 {
Self::ID
}
fn get_supported_packets(&self) -> &'static [u8] {
&[]
}
fn encode(&self) -> Bytes {
match self.role {
Role::Server => Bytes::new(),
Role::Client => {
let username = Bytes::from(self.username.clone().into_bytes());
let password = Bytes::from(self.password.clone().into_bytes());
let username_len = u8::try_from(username.len()).expect("username was too long");
let password_len = u16::try_from(password.len()).expect("password was too long");
let mut bytes =
BytesMut::with_capacity(3 + username_len as usize + password_len as usize);
bytes.put_u8(username_len);
bytes.put_u16_le(password_len);
bytes.extend(username);
bytes.extend(password);
bytes.freeze()
}
}
}
async fn handle_handshake(
&mut self,
_: &mut dyn WebSocketRead,
_: &LockedWebSocketWrite,
) -> Result<(), WispError> {
Ok(())
}
async fn handle_packet(
&mut self,
_: Bytes,
_: &mut dyn WebSocketRead,
_: &LockedWebSocketWrite,
) -> Result<(), WispError> {
Ok(())
}
fn box_clone(&self) -> Box<dyn ProtocolExtension + Sync + Send> {
Box::new(self.clone())
}
}
#[derive(Debug)]
enum PasswordProtocolExtensionError {
Utf8Error(FromUtf8Error),
InvalidUsername,
InvalidPassword,
}
impl Display for PasswordProtocolExtensionError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
use PasswordProtocolExtensionError as E;
match self {
E::Utf8Error(e) => write!(f, "{}", e),
E::InvalidUsername => write!(f, "Invalid username"),
E::InvalidPassword => write!(f, "Invalid password"),
}
}
}
impl Error for PasswordProtocolExtensionError {}
impl From<PasswordProtocolExtensionError> for WispError {
fn from(value: PasswordProtocolExtensionError) -> Self {
WispError::ExtensionImplError(Box::new(value))
}
}
impl From<FromUtf8Error> for PasswordProtocolExtensionError {
fn from(value: FromUtf8Error) -> Self {
PasswordProtocolExtensionError::Utf8Error(value)
}
}
impl From<PasswordProtocolExtension> for AnyProtocolExtension {
fn from(value: PasswordProtocolExtension) -> Self {
AnyProtocolExtension(Box::new(value))
}
}
/// Password protocol extension builder.
///
/// **Passwords are sent in plain text!!**
pub struct PasswordProtocolExtensionBuilder {
/// Map of users and their passwords to allow. Only used on server.
pub users: HashMap<String, String>,
/// Username to authenticate with. Only used on client.
pub username: String,
/// Password to authenticate with. Only used on client.
pub password: String,
}
impl PasswordProtocolExtensionBuilder {
/// Create a new password protocol extension builder for the server, with a map of users
/// and passwords to allow.
pub fn new_server(users: HashMap<String, String>) -> Self {
Self {
users,
username: String::new(),
password: String::new(),
}
}
/// Create a new password protocol extension builder for the client, with a username and
/// password to authenticate with.
pub fn new_client(username: String, password: String) -> Self {
Self {
users: HashMap::new(),
username,
password,
}
}
}
impl ProtocolExtensionBuilder for PasswordProtocolExtensionBuilder {
fn get_id(&self) -> u8 {
PasswordProtocolExtension::ID
}
fn build_from_bytes(
&self,
mut payload: Bytes,
role: crate::Role,
) -> Result<AnyProtocolExtension, WispError> {
match role {
Role::Server => {
if payload.remaining() < 3 {
return Err(WispError::PacketTooSmall);
}
let username_len = payload.get_u8();
let password_len = payload.get_u16_le();
if payload.remaining() < (password_len + username_len as u16) as usize {
return Err(WispError::PacketTooSmall);
}
use PasswordProtocolExtensionError as EError;
let username =
String::from_utf8(payload.copy_to_bytes(username_len as usize).to_vec())
.map_err(|x| WispError::from(EError::from(x)))?;
let password =
String::from_utf8(payload.copy_to_bytes(password_len as usize).to_vec())
.map_err(|x| WispError::from(EError::from(x)))?;
let Some(user) = self.users.iter().find(|x| *x.0 == username) else {
return Err(EError::InvalidUsername.into());
};
if *user.1 != password {
return Err(EError::InvalidPassword.into());
}
Ok(PasswordProtocolExtension {
username,
password,
role,
}
.into())
}
Role::Client => {
Ok(PasswordProtocolExtension::new_client(String::new(), String::new()).into())
}
}
}
fn build_to_extension(&self, role: Role) -> AnyProtocolExtension {
match role {
Role::Server => PasswordProtocolExtension::new_server(),
Role::Client => {
PasswordProtocolExtension::new_client(self.username.clone(), self.password.clone())
}
}
.into()
}
}

93
wisp/src/extensions/udp.rs Normal file
View file

@ -0,0 +1,93 @@
//! UDP protocol extension.
//!
//! # Example
//! ```
//! let (mux, fut) = ServerMux::new(
//! rx,
//! tx,
//! 128,
//! Some(&[Box::new(UdpProtocolExtensionBuilder())])
//! );
//! ```
//! See [the docs](https://github.com/MercuryWorkshop/wisp-protocol/blob/main/protocol.md#0x01---udp)
use async_trait::async_trait;
use bytes::Bytes;
use crate::{
ws::{LockedWebSocketWrite, WebSocketRead},
WispError,
};
use super::{AnyProtocolExtension, ProtocolExtension, ProtocolExtensionBuilder};
#[derive(Debug)]
/// UDP protocol extension.
pub struct UdpProtocolExtension();
impl UdpProtocolExtension {
/// UDP protocol extension ID.
pub const ID: u8 = 0x01;
}
#[async_trait]
impl ProtocolExtension for UdpProtocolExtension {
fn get_id(&self) -> u8 {
Self::ID
}
fn get_supported_packets(&self) -> &'static [u8] {
&[]
}
fn encode(&self) -> Bytes {
Bytes::new()
}
async fn handle_handshake(
&mut self,
_: &mut dyn WebSocketRead,
_: &LockedWebSocketWrite,
) -> Result<(), WispError> {
Ok(())
}
async fn handle_packet(
&mut self,
_: Bytes,
_: &mut dyn WebSocketRead,
_: &LockedWebSocketWrite,
) -> Result<(), WispError> {
Ok(())
}
fn box_clone(&self) -> Box<dyn ProtocolExtension + Sync + Send> {
Box::new(Self())
}
}
impl From<UdpProtocolExtension> for AnyProtocolExtension {
fn from(value: UdpProtocolExtension) -> Self {
AnyProtocolExtension(Box::new(value))
}
}
/// UDP protocol extension builder.
pub struct UdpProtocolExtensionBuilder();
impl ProtocolExtensionBuilder for UdpProtocolExtensionBuilder {
fn get_id(&self) -> u8 {
UdpProtocolExtension::ID
}
fn build_from_bytes(
&self,
_: Bytes,
_: crate::Role,
) -> Result<AnyProtocolExtension, WispError> {
Ok(UdpProtocolExtension().into())
}
fn build_to_extension(&self, _: crate::Role) -> AnyProtocolExtension {
UdpProtocolExtension().into()
}
}

49
wisp/src/fastwebsockets.rs
View file

@ -1,9 +1,14 @@
use bytes::Bytes;
use std::ops::Deref;
use async_trait::async_trait;
use bytes::BytesMut;
use fastwebsockets::{
FragmentCollectorRead, Frame, OpCode, Payload, WebSocketError, WebSocketWrite,
CloseCode, FragmentCollectorRead, Frame, OpCode, Payload, WebSocketError, WebSocketWrite,
};
use tokio::io::{AsyncRead, AsyncWrite};
use crate::{ws::LockedWebSocketWrite, WispError};
impl From<OpCode> for crate::ws::OpCode {
fn from(opcode: OpCode) -> Self {
use OpCode::*;
@ -21,29 +26,25 @@ impl From<OpCode> for crate::ws::OpCode {
}
impl From<Frame<'_>> for crate::ws::Frame {
fn from(mut frame: Frame) -> Self {
fn from(frame: Frame) -> Self {
Self {
finished: frame.fin,
opcode: frame.opcode.into(),
payload: Bytes::copy_from_slice(frame.payload.to_mut()),
payload: BytesMut::from(frame.payload.deref()).freeze(),
}
}
}
impl From<crate::ws::Frame> for Frame<'_> {
impl<'a> From<crate::ws::Frame> for Frame<'a> {
fn from(frame: crate::ws::Frame) -> Self {
use crate::ws::OpCode::*;
let payload = Payload::Owned(frame.payload.into());
match frame.opcode {
Text => Self::text(Payload::Owned(frame.payload.to_vec())),
Binary => Self::binary(Payload::Owned(frame.payload.to_vec())),
Close => Self::close_raw(Payload::Owned(frame.payload.to_vec())),
Ping => Self::new(
true,
OpCode::Ping,
None,
Payload::Owned(frame.payload.to_vec()),
),
Pong => Self::pong(Payload::Owned(frame.payload.to_vec())),
Text => Self::text(payload),
Binary => Self::binary(payload),
Close => Self::close_raw(payload),
Ping => Self::new(true, OpCode::Ping, None, payload),
Pong => Self::pong(payload),
}
}
}
@ -58,11 +59,12 @@ impl From<WebSocketError> for crate::WispError {
}
}
impl<S: AsyncRead + Unpin> crate::ws::WebSocketRead for FragmentCollectorRead<S> {
#[async_trait]
impl<S: AsyncRead + Unpin + Send> crate::ws::WebSocketRead for FragmentCollectorRead<S> {
async fn wisp_read_frame(
&mut self,
tx: &crate::ws::LockedWebSocketWrite<impl crate::ws::WebSocketWrite>,
) -> Result<crate::ws::Frame, crate::WispError> {
tx: &LockedWebSocketWrite,
) -> Result<crate::ws::Frame, WispError> {
Ok(self
.read_frame(&mut |frame| async { tx.write_frame(frame.into()).await })
.await?
@ -70,8 +72,15 @@ impl<S: AsyncRead + Unpin> crate::ws::WebSocketRead for FragmentCollectorRead<S>
}
}
impl<S: AsyncWrite + Unpin> crate::ws::WebSocketWrite for WebSocketWrite<S> {
async fn wisp_write_frame(&mut self, frame: crate::ws::Frame) -> Result<(), crate::WispError> {
#[async_trait]
impl<S: AsyncWrite + Unpin + Send> crate::ws::WebSocketWrite for WebSocketWrite<S> {
async fn wisp_write_frame(&mut self, frame: crate::ws::Frame) -> Result<(), WispError> {
self.write_frame(frame.into()).await.map_err(|e| e.into())
}
async fn wisp_close(&mut self) -> Result<(), WispError> {
self.write_frame(Frame::close(CloseCode::Normal.into(), b""))
.await
.map_err(|e| e.into())
}
}

671
wisp/src/lib.rs
View file

@ -1,9 +1,10 @@
#![deny(missing_docs)]
#![deny(missing_docs, warnings)]
#![cfg_attr(docsrs, feature(doc_cfg))]
//! A library for easily creating [Wisp] clients and servers.
//!
//! [Wisp]: https://github.com/MercuryWorkshop/wisp-protocol
pub mod extensions;
#[cfg(feature = "fastwebsockets")]
#[cfg_attr(docsrs, doc(cfg(feature = "fastwebsockets")))]
mod fastwebsockets;
@ -12,18 +13,26 @@ mod sink_unfold;
mod stream;
pub mod ws;
pub use crate::packet::*;
pub use crate::stream::*;
pub use crate::{packet::*, stream::*};
use bytes::Bytes;
use dashmap::DashMap;
use event_listener::Event;
use futures::SinkExt;
use futures::{channel::mpsc, Future, FutureExt, StreamExt};
use std::sync::{
atomic::{AtomicBool, AtomicU32, Ordering},
Arc,
use extensions::{udp::UdpProtocolExtension, AnyProtocolExtension, ProtocolExtensionBuilder};
use flume as mpsc;
use futures::{channel::oneshot, select, Future, FutureExt};
use futures_timer::Delay;
use std::{
sync::{
atomic::{AtomicBool, AtomicU32, Ordering},
Arc,
},
time::Duration,
};
use ws::AppendingWebSocketRead;
/// Wisp version supported by this crate.
pub const WISP_VERSION: WispVersion = WispVersion { major: 2, minor: 0 };
/// The role of the multiplexor.
#[derive(Debug, PartialEq, Copy, Clone)]
@ -37,29 +46,29 @@ pub enum Role {
/// Errors the Wisp implementation can return.
#[derive(Debug)]
pub enum WispError {
/// The packet recieved did not have enough data.
/// The packet received did not have enough data.
PacketTooSmall,
/// The packet recieved had an invalid type.
/// The packet received had an invalid type.
InvalidPacketType,
/// The stream had an invalid type.
InvalidStreamType,
/// The stream had an invalid ID.
InvalidStreamId,
/// The close packet had an invalid reason.
InvalidCloseReason,
/// The URI recieved was invalid.
/// The URI received was invalid.
InvalidUri,
/// The URI recieved had no host.
/// The URI received had no host.
UriHasNoHost,
/// The URI recieved had no port.
/// The URI received had no port.
UriHasNoPort,
/// The max stream count was reached.
MaxStreamCountReached,
/// The Wisp protocol version was incompatible.
IncompatibleProtocolVersion,
/// The stream had already been closed.
StreamAlreadyClosed,
/// The websocket frame recieved had an invalid type.
/// The websocket frame received had an invalid type.
WsFrameInvalidType,
/// The websocket frame recieved was not finished.
/// The websocket frame received was not finished.
WsFrameNotFinished,
/// Error specific to the websocket implementation.
WsImplError(Box<dyn std::error::Error + Sync + Send>),
@ -67,17 +76,33 @@ pub enum WispError {
WsImplSocketClosed,
/// The websocket implementation did not support the action.
WsImplNotSupported,
/// Error specific to the protocol extension implementation.
ExtensionImplError(Box<dyn std::error::Error + Sync + Send>),
/// The protocol extension implementation did not support the action.
ExtensionImplNotSupported,
/// The UDP protocol extension is not supported by the server.
UdpExtensionNotSupported,
/// The string was invalid UTF-8.
Utf8Error(std::str::Utf8Error),
/// The integer failed to convert.
TryFromIntError(std::num::TryFromIntError),
/// Other error.
Other(Box<dyn std::error::Error + Sync + Send>),
/// Failed to send message to multiplexor task.
MuxMessageFailedToSend,
/// Failed to receive message from multiplexor task.
MuxMessageFailedToRecv,
}
impl From<std::str::Utf8Error> for WispError {
fn from(err: std::str::Utf8Error) -> WispError {
WispError::Utf8Error(err)
fn from(err: std::str::Utf8Error) -> Self {
Self::Utf8Error(err)
}
}
impl From<std::num::TryFromIntError> for WispError {
fn from(value: std::num::TryFromIntError) -> Self {
Self::TryFromIntError(value)
}
}
@ -86,13 +111,13 @@ impl std::fmt::Display for WispError {
match self {
Self::PacketTooSmall => write!(f, "Packet too small"),
Self::InvalidPacketType => write!(f, "Invalid packet type"),
Self::InvalidStreamType => write!(f, "Invalid stream type"),
Self::InvalidStreamId => write!(f, "Invalid stream id"),
Self::InvalidCloseReason => write!(f, "Invalid close reason"),
Self::InvalidUri => write!(f, "Invalid URI"),
Self::UriHasNoHost => write!(f, "URI has no host"),
Self::UriHasNoPort => write!(f, "URI has no port"),
Self::MaxStreamCountReached => write!(f, "Maximum stream count reached"),
Self::IncompatibleProtocolVersion => write!(f, "Incompatible Wisp protocol version"),
Self::StreamAlreadyClosed => write!(f, "Stream already closed"),
Self::WsFrameInvalidType => write!(f, "Invalid websocket frame type"),
Self::WsFrameNotFinished => write!(f, "Unfinished websocket frame"),
@ -103,9 +128,21 @@ impl std::fmt::Display for WispError {
Self::WsImplNotSupported => {
write!(f, "Websocket implementation error: unsupported feature")
}
Self::ExtensionImplError(err) => {
write!(f, "Protocol extension implementation error: {}", err)
}
Self::ExtensionImplNotSupported => {
write!(
f,
"Protocol extension implementation error: unsupported feature"
)
}
Self::UdpExtensionNotSupported => write!(f, "UDP protocol extension not supported"),
Self::Utf8Error(err) => write!(f, "UTF-8 error: {}", err),
Self::TryFromIntError(err) => write!(f, "Integer conversion error: {}", err),
Self::Other(err) => write!(f, "Other error: {}", err),
Self::MuxMessageFailedToSend => write!(f, "Failed to send multiplexor message"),
Self::MuxMessageFailedToRecv => write!(f, "Failed to receive multiplexor message"),
}
}
}
@ -113,36 +150,36 @@ impl std::fmt::Display for WispError {
impl std::error::Error for WispError {}
struct MuxMapValue {
stream: mpsc::UnboundedSender<Bytes>,
stream: mpsc::Sender<Bytes>,
stream_type: StreamType,
flow_control: Arc<AtomicU32>,
flow_control_event: Arc<Event>,
is_closed: Arc<AtomicBool>,
is_closed_event: Arc<Event>,
}
struct MuxInner<W>
where
W: ws::WebSocketWrite,
{
tx: ws::LockedWebSocketWrite<W>,
stream_map: Arc<DashMap<u32, MuxMapValue>>,
struct MuxInner {
tx: ws::LockedWebSocketWrite,
stream_map: DashMap<u32, MuxMapValue>,
buffer_size: u32,
}
impl<W: ws::WebSocketWrite> MuxInner<W> {
impl MuxInner {
pub async fn server_into_future<R>(
self,
rx: R,
extensions: Vec<AnyProtocolExtension>,
close_rx: mpsc::Receiver<WsEvent>,
muxstream_sender: mpsc::UnboundedSender<(ConnectPacket, MuxStream)>,
buffer_size: u32,
muxstream_sender: mpsc::Sender<(ConnectPacket, MuxStream)>,
close_tx: mpsc::Sender<WsEvent>,
) -> Result<(), WispError>
where
R: ws::WebSocketRead,
R: ws::WebSocketRead + Send,
{
self.into_future(
self.as_future(
close_rx,
self.server_loop(rx, muxstream_sender, buffer_size, close_tx),
close_tx.clone(),
self.server_loop(rx, extensions, muxstream_sender, close_tx),
)
.await
}
@ -150,34 +187,83 @@ impl<W: ws::WebSocketWrite> MuxInner<W> {
pub async fn client_into_future<R>(
self,
rx: R,
extensions: Vec<AnyProtocolExtension>,
close_rx: mpsc::Receiver<WsEvent>,
close_tx: mpsc::Sender<WsEvent>,
) -> Result<(), WispError>
where
R: ws::WebSocketRead,
R: ws::WebSocketRead + Send,
{
self.into_future(close_rx, self.client_loop(rx)).await
self.as_future(close_rx, close_tx, self.client_loop(rx, extensions))
.await
}
async fn into_future(
async fn as_future(
&self,
close_rx: mpsc::Receiver<WsEvent>,
close_tx: mpsc::Sender<WsEvent>,
wisp_fut: impl Future<Output = Result<(), WispError>>,
) -> Result<(), WispError> {
let ret = futures::select! {
_ = self.stream_loop(close_rx).fuse() => Ok(()),
_ = self.stream_loop(close_rx, close_tx).fuse() => Ok(()),
x = wisp_fut.fuse() => x,
};
self.stream_map.iter_mut().for_each(|mut x| {
for x in self.stream_map.iter_mut() {
x.is_closed.store(true, Ordering::Release);
x.stream.disconnect();
x.stream.close_channel();
});
x.is_closed_event.notify(usize::MAX);
}
self.stream_map.clear();
let _ = self.tx.close().await;
ret
}
async fn stream_loop(&self, mut stream_rx: mpsc::Receiver<WsEvent>) {
while let Some(msg) = stream_rx.next().await {
async fn create_new_stream(
&self,
stream_id: u32,
stream_type: StreamType,
role: Role,
stream_tx: mpsc::Sender<WsEvent>,
target_buffer_size: u32,
) -> Result<(MuxMapValue, MuxStream), WispError> {
let (ch_tx, ch_rx) = mpsc::bounded(self.buffer_size as usize);
let flow_control_event: Arc<Event> = Event::new().into();
let flow_control: Arc<AtomicU32> = AtomicU32::new(self.buffer_size).into();
let is_closed: Arc<AtomicBool> = AtomicBool::new(false).into();
let is_closed_event: Arc<Event> = Event::new().into();
Ok((
MuxMapValue {
stream: ch_tx,
stream_type,
flow_control: flow_control.clone(),
flow_control_event: flow_control_event.clone(),
is_closed: is_closed.clone(),
is_closed_event: is_closed_event.clone(),
},
MuxStream::new(
stream_id,
role,
stream_type,
ch_rx,
stream_tx.clone(),
is_closed,
is_closed_event,
flow_control,
flow_control_event,
target_buffer_size,
),
))
}
async fn stream_loop(
&self,
stream_rx: mpsc::Receiver<WsEvent>,
stream_tx: mpsc::Sender<WsEvent>,
) {
let mut next_free_stream_id: u32 = 1;
while let Ok(msg) = stream_rx.recv_async().await {
match msg {
WsEvent::SendPacket(packet, channel) => {
if self.stream_map.get(&packet.stream_id).is_some() {
@ -186,16 +272,55 @@ impl<W: ws::WebSocketWrite> MuxInner<W> {
let _ = channel.send(Err(WispError::InvalidStreamId));
}
}
WsEvent::CreateStream(stream_type, host, port, channel) => {
let ret: Result<MuxStream, WispError> = async {
let stream_id = next_free_stream_id;
let next_stream_id = next_free_stream_id
.checked_add(1)
.ok_or(WispError::MaxStreamCountReached)?;
let (map_value, stream) = self
.create_new_stream(
stream_id,
stream_type,
Role::Client,
stream_tx.clone(),
0,
)
.await?;
self.tx
.write_frame(
Packet::new_connect(stream_id, stream_type, port, host).into(),
)
.await?;
self.stream_map.insert(stream_id, map_value);
next_free_stream_id = next_stream_id;
Ok(stream)
}
.await;
let _ = channel.send(ret);
}
WsEvent::Close(packet, channel) => {
if let Some((_, mut stream)) = self.stream_map.remove(&packet.stream_id) {
stream.stream.disconnect();
stream.stream.close_channel();
if let Some((_, stream)) = self.stream_map.remove(&packet.stream_id) {
let _ = channel.send(self.tx.write_frame(packet.into()).await);
drop(stream.stream)
} else {
let _ = channel.send(Err(WispError::InvalidStreamId));
}
}
WsEvent::EndFut => break,
WsEvent::EndFut(x) => {
if let Some(reason) = x {
let _ = self
.tx
.write_frame(Packet::new_close(0, reason).into())
.await;
}
break;
}
}
}
}
@ -203,122 +328,115 @@ impl<W: ws::WebSocketWrite> MuxInner<W> {
async fn server_loop<R>(
&self,
mut rx: R,
muxstream_sender: mpsc::UnboundedSender<(ConnectPacket, MuxStream)>,
buffer_size: u32,
close_tx: mpsc::Sender<WsEvent>,
mut extensions: Vec<AnyProtocolExtension>,
muxstream_sender: mpsc::Sender<(ConnectPacket, MuxStream)>,
stream_tx: mpsc::Sender<WsEvent>,
) -> Result<(), WispError>
where
R: ws::WebSocketRead,
R: ws::WebSocketRead + Send,
{
// will send continues once flow_control is at 10% of max
let target_buffer_size = ((buffer_size as u64 * 90) / 100) as u32;
self.tx
.write_frame(Packet::new_continue(0, buffer_size).into())
.await?;
let target_buffer_size = ((self.buffer_size as u64 * 90) / 100) as u32;
loop {
let frame = rx.wisp_read_frame(&self.tx).await?;
if frame.opcode == ws::OpCode::Close {
break Ok(());
}
let packet = Packet::try_from(frame)?;
use PacketType::*;
match packet.packet_type {
Connect(inner_packet) => {
let (ch_tx, ch_rx) = mpsc::unbounded();
let stream_type = inner_packet.stream_type;
let flow_control: Arc<AtomicU32> = AtomicU32::new(buffer_size).into();
let flow_control_event: Arc<Event> = Event::new().into();
let is_closed: Arc<AtomicBool> = AtomicBool::new(false).into();
self.stream_map.insert(
packet.stream_id,
MuxMapValue {
stream: ch_tx,
stream_type,
flow_control: flow_control.clone(),
flow_control_event: flow_control_event.clone(),
is_closed: is_closed.clone(),
},
);
muxstream_sender
.unbounded_send((
inner_packet,
MuxStream::new(
if let Some(packet) =
Packet::maybe_handle_extension(frame, &mut extensions, &mut rx, &self.tx).await?
{
use PacketType::*;
match packet.packet_type {
Connect(inner_packet) => {
let (map_value, stream) = self
.create_new_stream(
packet.stream_id,
inner_packet.stream_type,
Role::Server,
stream_type,
ch_rx,
close_tx.clone(),
is_closed,
flow_control,
flow_control_event,
stream_tx.clone(),
target_buffer_size,
),
))
.map_err(|x| WispError::Other(Box::new(x)))?;
}
Data(data) => {
if let Some(stream) = self.stream_map.get(&packet.stream_id) {
let _ = stream.stream.unbounded_send(data);
if stream.stream_type == StreamType::Tcp {
stream.flow_control.store(
stream
.flow_control
.load(Ordering::Acquire)
.saturating_sub(1),
Ordering::Release,
);
)
.await?;
muxstream_sender
.send_async((inner_packet, stream))
.await
.map_err(|_| WispError::MuxMessageFailedToSend)?;
self.stream_map.insert(packet.stream_id, map_value);
}
Data(data) => {
if let Some(stream) = self.stream_map.get(&packet.stream_id) {
let _ = stream.stream.send_async(data).await;
if stream.stream_type == StreamType::Tcp {
stream.flow_control.store(
stream
.flow_control
.load(Ordering::Acquire)
.saturating_sub(1),
Ordering::Release,
);
}
}
}
}
Continue(_) => break Err(WispError::InvalidPacketType),
Close(_) => {
if let Some((_, mut stream)) = self.stream_map.remove(&packet.stream_id) {
stream.is_closed.store(true, Ordering::Release);
stream.stream.disconnect();
stream.stream.close_channel();
Continue(_) | Info(_) => break Err(WispError::InvalidPacketType),
Close(_) => {
if packet.stream_id == 0 {
break Ok(());
}
if let Some((_, stream)) = self.stream_map.remove(&packet.stream_id) {
stream.is_closed.store(true, Ordering::Release);
stream.is_closed_event.notify(usize::MAX);
drop(stream.stream)
}
}
}
}
}
}
async fn client_loop<R>(&self, mut rx: R) -> Result<(), WispError>
async fn client_loop<R>(
&self,
mut rx: R,
mut extensions: Vec<AnyProtocolExtension>,
) -> Result<(), WispError>
where
R: ws::WebSocketRead,
R: ws::WebSocketRead + Send,
{
loop {
let frame = rx.wisp_read_frame(&self.tx).await?;
if frame.opcode == ws::OpCode::Close {
break Ok(());
}
let packet = Packet::try_from(frame)?;
use PacketType::*;
match packet.packet_type {
Connect(_) => break Err(WispError::InvalidPacketType),
Data(data) => {
if let Some(stream) = self.stream_map.get(&packet.stream_id) {
let _ = stream.stream.unbounded_send(data);
}
}
Continue(inner_packet) => {
if let Some(stream) = self.stream_map.get(&packet.stream_id) {
if stream.stream_type == StreamType::Tcp {
stream
.flow_control
.store(inner_packet.buffer_remaining, Ordering::Release);
let _ = stream.flow_control_event.notify(u32::MAX);
if let Some(packet) =
Packet::maybe_handle_extension(frame, &mut extensions, &mut rx, &self.tx).await?
{
use PacketType::*;
match packet.packet_type {
Connect(_) | Info(_) => break Err(WispError::InvalidPacketType),
Data(data) => {
if let Some(stream) = self.stream_map.get(&packet.stream_id) {
let _ = stream.stream.send_async(data).await;
}
}
}
Close(_) => {
if let Some((_, mut stream)) = self.stream_map.remove(&packet.stream_id) {
stream.is_closed.store(true, Ordering::Release);
stream.stream.disconnect();
stream.stream.close_channel();
Continue(inner_packet) => {
if let Some(stream) = self.stream_map.get(&packet.stream_id) {
if stream.stream_type == StreamType::Tcp {
stream
.flow_control
.store(inner_packet.buffer_remaining, Ordering::Release);
let _ = stream.flow_control_event.notify(u32::MAX);
}
}
}
Close(_) => {
if packet.stream_id == 0 {
break Ok(());
}
if let Some((_, stream)) = self.stream_map.remove(&packet.stream_id) {
stream.is_closed.store(true, Ordering::Release);
stream.is_closed_event.notify(usize::MAX);
drop(stream.stream)
}
}
}
}
@ -332,7 +450,7 @@ impl<W: ws::WebSocketWrite> MuxInner<W> {
/// ```
/// use wisp_mux::ServerMux;
///
/// let (mux, fut) = ServerMux::new(rx, tx, 128);
/// let (mux, fut) = ServerMux::new(rx, tx, 128, Some([]));
/// tokio::spawn(async move {
/// if let Err(e) = fut.await {
/// println!("error in multiplexor: {:?}", e);
@ -346,39 +464,103 @@ impl<W: ws::WebSocketWrite> MuxInner<W> {
/// }
/// ```
pub struct ServerMux {
/// Whether the connection was downgraded to Wisp v1.
///
/// If this variable is true you must assume no extensions are supported.
pub downgraded: bool,
/// Extensions that are supported by both sides.
pub supported_extension_ids: Vec<u8>,
close_tx: mpsc::Sender<WsEvent>,
muxstream_recv: mpsc::UnboundedReceiver<(ConnectPacket, MuxStream)>,
muxstream_recv: mpsc::Receiver<(ConnectPacket, MuxStream)>,
}
impl ServerMux {
/// Create a new server-side multiplexor.
pub fn new<R, W: ws::WebSocketWrite>(
read: R,
///
/// If `extension_builders` is None a Wisp v1 connection is created otherwise a Wisp v2 connection is created.
/// **It is not guaranteed that all extensions you specify are available.** You must manually check
/// if the extensions you need are available after the multiplexor has been created.
pub async fn new<R, W>(
mut read: R,
write: W,
buffer_size: u32,
) -> (Self, impl Future<Output = Result<(), WispError>>)
extension_builders: Option<&[Box<dyn ProtocolExtensionBuilder + Send + Sync>]>,
) -> Result<(Self, impl Future<Output = Result<(), WispError>> + Send), WispError>
where
R: ws::WebSocketRead,
R: ws::WebSocketRead + Send,
W: ws::WebSocketWrite + Send + 'static,
{
let (close_tx, close_rx) = mpsc::channel::<WsEvent>(256);
let (close_tx, close_rx) = mpsc::bounded::<WsEvent>(256);
let (tx, rx) = mpsc::unbounded::<(ConnectPacket, MuxStream)>();
let write = ws::LockedWebSocketWrite::new(write);
(
let write = ws::LockedWebSocketWrite::new(Box::new(write));
write
.write_frame(Packet::new_continue(0, buffer_size).into())
.await?;
let mut supported_extensions = Vec::new();
let mut extra_packet = Vec::with_capacity(1);
let mut downgraded = true;
if let Some(builders) = extension_builders {
let extensions: Vec<_> = builders
.iter()
.map(|x| x.build_to_extension(Role::Server))
.collect();
let extension_ids: Vec<_> = extensions.iter().map(|x| x.get_id()).collect();
write
.write_frame(Packet::new_info(extensions).into())
.await?;
if let Some(frame) = select! {
x = read.wisp_read_frame(&write).fuse() => Some(x?),
_ = Delay::new(Duration::from_secs(5)).fuse() => None
} {
let packet = Packet::maybe_parse_info(frame, Role::Server, builders)?;
if let PacketType::Info(info) = packet.packet_type {
supported_extensions = info
.extensions
.into_iter()
.filter(|x| extension_ids.contains(&x.get_id()))
.collect();
downgraded = false;
} else {
extra_packet.push(packet.into());
}
}
}
Ok((
Self {
muxstream_recv: rx,
close_tx: close_tx.clone(),
downgraded,
supported_extension_ids: supported_extensions.iter().map(|x| x.get_id()).collect(),
},
MuxInner {
tx: write,
stream_map: DashMap::new().into(),
stream_map: DashMap::new(),
buffer_size,
}
.server_into_future(read, close_rx, tx, buffer_size, close_tx),
)
.server_into_future(
AppendingWebSocketRead(extra_packet, read),
supported_extensions,
close_rx,
tx,
close_tx,
),
))
}
/// Wait for a stream to be created.
pub async fn server_new_stream(&mut self) -> Option<(ConnectPacket, MuxStream)> {
self.muxstream_recv.next().await
self.muxstream_recv.recv_async().await.ok()
}
async fn close_internal(&mut self, reason: Option<CloseReason>) -> Result<(), WispError> {
self.close_tx
.send_async(WsEvent::EndFut(reason))
.await
.map_err(|_| WispError::MuxMessageFailedToSend)
}
/// Close all streams.
@ -386,19 +568,32 @@ impl ServerMux {
/// Also terminates the multiplexor future. Waiting for a new stream will never succeed after
/// this function is called.
pub async fn close(&mut self) -> Result<(), WispError> {
self.close_tx
.send(WsEvent::EndFut)
self.close_internal(None).await
}
/// Close all streams and send an extension incompatibility error to the client.
///
/// Also terminates the multiplexor future. Waiting for a new stream will never succed after
/// this function is called.
pub async fn close_extension_incompat(&mut self) -> Result<(), WispError> {
self.close_internal(Some(CloseReason::IncompatibleExtensions))
.await
.map_err(|_| WispError::MuxMessageFailedToSend)
}
}
impl Drop for ServerMux {
fn drop(&mut self) {
let _ = self.close_tx.send(WsEvent::EndFut(None));
}
}
/// Client side multiplexor.
///
/// # Example
/// ```
/// use wisp_mux::{ClientMux, StreamType};
///
/// let (mux, fut) = ClientMux::new(rx, tx).await?;
/// let (mux, fut) = ClientMux::new(rx, tx, Some([])).await?;
/// tokio::spawn(async move {
/// if let Err(e) = fut.await {
/// println!("error in multiplexor: {:?}", e);
@ -406,50 +601,93 @@ impl ServerMux {
/// });
/// let stream = mux.client_new_stream(StreamType::Tcp, "google.com", 80);
/// ```
pub struct ClientMux<W>
where
W: ws::WebSocketWrite,
{
tx: ws::LockedWebSocketWrite<W>,
stream_map: Arc<DashMap<u32, MuxMapValue>>,
next_free_stream_id: AtomicU32,
close_tx: mpsc::Sender<WsEvent>,
buf_size: u32,
target_buf_size: u32,
pub struct ClientMux {
/// Whether the connection was downgraded to Wisp v1.
///
/// If this variable is true you must assume no extensions are supported.
pub downgraded: bool,
/// Extensions that are supported by both sides.
pub supported_extension_ids: Vec<u8>,
stream_tx: mpsc::Sender<WsEvent>,
}
impl<W: ws::WebSocketWrite> ClientMux<W> {
impl ClientMux {
/// Create a new client side multiplexor.
pub async fn new<R>(
///
/// If `extension_builders` is None a Wisp v1 connection is created otherwise a Wisp v2 connection is created.
/// **It is not guaranteed that all extensions you specify are available.** You must manually check
/// if the extensions you need are available after the multiplexor has been created.
pub async fn new<R, W>(
mut read: R,
write: W,
) -> Result<(Self, impl Future<Output = Result<(), WispError>>), WispError>
extension_builders: Option<&[Box<dyn ProtocolExtensionBuilder + Send + Sync>]>,
) -> Result<(Self, impl Future<Output = Result<(), WispError>> + Send), WispError>
where
R: ws::WebSocketRead,
R: ws::WebSocketRead + Send,
W: ws::WebSocketWrite + Send + 'static,
{
let write = ws::LockedWebSocketWrite::new(write);
let write = ws::LockedWebSocketWrite::new(Box::new(write));
let first_packet = Packet::try_from(read.wisp_read_frame(&write).await?)?;
if first_packet.stream_id != 0 {
return Err(WispError::InvalidStreamId);
}
if let PacketType::Continue(packet) = first_packet.packet_type {
let (tx, rx) = mpsc::channel::<WsEvent>(256);
let map = Arc::new(DashMap::new());
let mut supported_extensions = Vec::new();
let mut extra_packet = Vec::with_capacity(1);
let mut downgraded = true;
if let Some(builders) = extension_builders {
let extensions: Vec<_> = builders
.iter()
.map(|x| x.build_to_extension(Role::Client))
.collect();
let extension_ids: Vec<_> = extensions.iter().map(|x| x.get_id()).collect();
if let Some(frame) = select! {
x = read.wisp_read_frame(&write).fuse() => Some(x?),
_ = Delay::new(Duration::from_secs(5)).fuse() => None
} {
let packet = Packet::maybe_parse_info(frame, Role::Client, builders)?;
if let PacketType::Info(info) = packet.packet_type {
supported_extensions = info
.extensions
.into_iter()
.filter(|x| extension_ids.contains(&x.get_id()))
.collect();
write
.write_frame(Packet::new_info(extensions).into())
.await?;
downgraded = false;
} else {
extra_packet.push(packet.into());
}
}
}
for extension in supported_extensions.iter_mut() {
extension.handle_handshake(&mut read, &write).await?;
}
let (tx, rx) = mpsc::bounded::<WsEvent>(256);
Ok((
Self {
tx: write.clone(),
stream_map: map.clone(),
next_free_stream_id: AtomicU32::new(1),
close_tx: tx.clone(),
buf_size: packet.buffer_remaining,
// server-only
target_buf_size: 0,
stream_tx: tx.clone(),
downgraded,
supported_extension_ids: supported_extensions
.iter()
.map(|x| x.get_id())
.collect(),
},
MuxInner {
tx: write.clone(),
stream_map: map.clone(),
tx: write,
stream_map: DashMap::new(),
buffer_size: packet.buffer_remaining,
}
.client_into_future(read, rx),
.client_into_future(
AppendingWebSocketRead(extra_packet, read),
supported_extensions,
rx,
tx,
),
))
} else {
Err(WispError::InvalidPacketType)
@ -458,51 +696,32 @@ impl<W: ws::WebSocketWrite> ClientMux<W> {
/// Create a new stream, multiplexed through Wisp.
pub async fn client_new_stream(
&self,
&mut self,
stream_type: StreamType,
host: String,
port: u16,
) -> Result<MuxStream, WispError> {
let (ch_tx, ch_rx) = mpsc::unbounded();
let stream_id = self.next_free_stream_id.load(Ordering::Acquire);
let next_stream_id = stream_id
.checked_add(1)
.ok_or(WispError::MaxStreamCountReached)?;
if stream_type == StreamType::Udp
&& !self
.supported_extension_ids
.iter()
.any(|x| *x == UdpProtocolExtension::ID)
{
return Err(WispError::UdpExtensionNotSupported);
}
let (tx, rx) = oneshot::channel();
self.stream_tx
.send_async(WsEvent::CreateStream(stream_type, host, port, tx))
.await
.map_err(|_| WispError::MuxMessageFailedToSend)?;
rx.await.map_err(|_| WispError::MuxMessageFailedToRecv)?
}
let flow_control_event: Arc<Event> = Event::new().into();
let flow_control: Arc<AtomicU32> = AtomicU32::new(self.buf_size).into();
let is_closed: Arc<AtomicBool> = AtomicBool::new(false).into();
self.tx
.write_frame(Packet::new_connect(stream_id, stream_type, port, host).into())
.await?;
self.next_free_stream_id
.store(next_stream_id, Ordering::Release);
self.stream_map.insert(
stream_id,
MuxMapValue {
stream: ch_tx,
stream_type,
flow_control: flow_control.clone(),
flow_control_event: flow_control_event.clone(),
is_closed: is_closed.clone(),
},
);
Ok(MuxStream::new(
stream_id,
Role::Client,
stream_type,
ch_rx,
self.close_tx.clone(),
is_closed,
flow_control,
flow_control_event,
self.target_buf_size,
))
async fn close_internal(&mut self, reason: Option<CloseReason>) -> Result<(), WispError> {
self.stream_tx
.send_async(WsEvent::EndFut(reason))
.await
.map_err(|_| WispError::MuxMessageFailedToSend)
}
/// Close all streams.
@ -510,9 +729,21 @@ impl<W: ws::WebSocketWrite> ClientMux<W> {
/// Also terminates the multiplexor future. Creating a stream is UB after calling this
/// function.
pub async fn close(&mut self) -> Result<(), WispError> {
self.close_tx
.send(WsEvent::EndFut)
self.close_internal(None).await
}
/// Close all streams and send an extension incompatibility error to the client.
///
/// Also terminates the multiplexor future. Creating a stream is UB after calling this
/// function.
pub async fn close_extension_incompat(&mut self) -> Result<(), WispError> {
self.close_internal(Some(CloseReason::IncompatibleExtensions))
.await
.map_err(|_| WispError::MuxMessageFailedToSend)
}
}
impl Drop for ClientMux {
fn drop(&mut self) {
let _ = self.stream_tx.send(WsEvent::EndFut(None));
}
}

270
wisp/src/packet.rs
View file

@ -1,23 +1,39 @@
use crate::{ws, WispError};
use crate::{
extensions::{AnyProtocolExtension, ProtocolExtensionBuilder},
ws::{self, Frame, LockedWebSocketWrite, OpCode, WebSocketRead},
Role, WispError, WISP_VERSION,
};
use bytes::{Buf, BufMut, Bytes, BytesMut};
/// Wisp stream type.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum StreamType {
/// TCP Wisp stream.
Tcp = 0x01,
Tcp,
/// UDP Wisp stream.
Udp = 0x02,
Udp,
/// Unknown Wisp stream type used for custom streams by protocol extensions.
Unknown(u8),
}
impl TryFrom<u8> for StreamType {
type Error = WispError;
fn try_from(stream_type: u8) -> Result<Self, Self::Error> {
use StreamType::*;
match stream_type {
0x01 => Ok(Tcp),
0x02 => Ok(Udp),
_ => Err(Self::Error::InvalidStreamType),
impl From<u8> for StreamType {
fn from(value: u8) -> Self {
use StreamType as S;
match value {
0x01 => S::Tcp,
0x02 => S::Udp,
x => S::Unknown(x),
}
}
}
impl From<StreamType> for u8 {
fn from(value: StreamType) -> Self {
use StreamType as S;
match value {
S::Tcp => 0x01,
S::Udp => 0x02,
S::Unknown(x) => x,
}
}
}
@ -34,6 +50,8 @@ pub enum CloseReason {
Voluntary = 0x02,
/// Unexpected stream closure due to a network error.
Unexpected = 0x03,
/// Incompatible extensions. Only used during the handshake.
IncompatibleExtensions = 0x04,
/// Stream creation failed due to invalid information.
ServerStreamInvalidInfo = 0x41,
/// Stream creation failed due to an unreachable destination host.
@ -54,21 +72,22 @@ pub enum CloseReason {
impl TryFrom<u8> for CloseReason {
type Error = WispError;
fn try_from(stream_type: u8) -> Result<Self, Self::Error> {
use CloseReason::*;
match stream_type {
0x01 => Ok(Unknown),
0x02 => Ok(Voluntary),
0x03 => Ok(Unexpected),
0x41 => Ok(ServerStreamInvalidInfo),
0x42 => Ok(ServerStreamUnreachable),
0x43 => Ok(ServerStreamConnectionTimedOut),
0x44 => Ok(ServerStreamConnectionRefused),
0x47 => Ok(ServerStreamTimedOut),
0x48 => Ok(ServerStreamBlockedAddress),
0x49 => Ok(ServerStreamThrottled),
0x81 => Ok(ClientUnexpected),
_ => Err(Self::Error::InvalidStreamType),
fn try_from(close_reason: u8) -> Result<Self, Self::Error> {
use CloseReason as R;
match close_reason {
0x01 => Ok(R::Unknown),
0x02 => Ok(R::Voluntary),
0x03 => Ok(R::Unexpected),
0x04 => Ok(R::IncompatibleExtensions),
0x41 => Ok(R::ServerStreamInvalidInfo),
0x42 => Ok(R::ServerStreamUnreachable),
0x43 => Ok(R::ServerStreamConnectionTimedOut),
0x44 => Ok(R::ServerStreamConnectionRefused),
0x47 => Ok(R::ServerStreamTimedOut),
0x48 => Ok(R::ServerStreamBlockedAddress),
0x49 => Ok(R::ServerStreamThrottled),
0x81 => Ok(R::ClientUnexpected),
_ => Err(Self::Error::InvalidCloseReason),
}
}
}
@ -108,7 +127,7 @@ impl TryFrom<Bytes> for ConnectPacket {
return Err(Self::Error::PacketTooSmall);
}
Ok(Self {
stream_type: bytes.get_u8().try_into()?,
stream_type: bytes.get_u8().into(),
destination_port: bytes.get_u16_le(),
destination_hostname: std::str::from_utf8(&bytes)?.to_string(),
})
@ -118,7 +137,7 @@ impl TryFrom<Bytes> for ConnectPacket {
impl From<ConnectPacket> for Bytes {
fn from(packet: ConnectPacket) -> Self {
let mut encoded = BytesMut::with_capacity(1 + 2 + packet.destination_hostname.len());
encoded.put_u8(packet.stream_type as u8);
encoded.put_u8(packet.stream_type.into());
encoded.put_u16_le(packet.destination_port);
encoded.extend(packet.destination_hostname.bytes());
encoded.freeze()
@ -198,6 +217,38 @@ impl From<ClosePacket> for Bytes {
}
}
/// Wisp version sent in the handshake.
#[derive(Debug, Clone)]
pub struct WispVersion {
/// Major Wisp version according to semver.
pub major: u8,
/// Minor Wisp version according to semver.
pub minor: u8,
}
/// Packet used in the initial handshake.
///
/// See [the docs](https://github.com/MercuryWorkshop/wisp-protocol/blob/main/protocol.md#0x05---info)
#[derive(Debug, Clone)]
pub struct InfoPacket {
/// Wisp version sent in the packet.
pub version: WispVersion,
/// List of protocol extensions sent in the packet.
pub extensions: Vec<AnyProtocolExtension>,
}
impl From<InfoPacket> for Bytes {
fn from(value: InfoPacket) -> Self {
let mut bytes = BytesMut::with_capacity(2);
bytes.put_u8(value.version.major);
bytes.put_u8(value.version.minor);
for extension in value.extensions {
bytes.extend(Bytes::from(extension));
}
bytes.freeze()
}
}
#[derive(Debug, Clone)]
/// Type of packet recieved.
pub enum PacketType {
@ -209,29 +260,33 @@ pub enum PacketType {
Continue(ContinuePacket),
/// Close packet.
Close(ClosePacket),
/// Info packet.
Info(InfoPacket),
}
impl PacketType {
/// Get the packet type used in the protocol.
pub fn as_u8(&self) -> u8 {
use PacketType::*;
use PacketType as P;
match self {
Connect(_) => 0x01,
Data(_) => 0x02,
Continue(_) => 0x03,
Close(_) => 0x04,
P::Connect(_) => 0x01,
P::Data(_) => 0x02,
P::Continue(_) => 0x03,
P::Close(_) => 0x04,
P::Info(_) => 0x05,
}
}
}
impl From<PacketType> for Bytes {
fn from(packet: PacketType) -> Self {
use PacketType::*;
use PacketType as P;
match packet {
Connect(x) => x.into(),
Data(x) => x,
Continue(x) => x.into(),
Close(x) => x.into(),
P::Connect(x) => x.into(),
P::Data(x) => x,
P::Continue(x) => x.into(),
P::Close(x) => x.into(),
P::Info(x) => x.into(),
}
}
}
@ -296,26 +351,141 @@ impl Packet {
packet_type: PacketType::Close(ClosePacket::new(reason)),
}
}
pub(crate) fn new_info(extensions: Vec<AnyProtocolExtension>) -> Self {
Self {
stream_id: 0,
packet_type: PacketType::Info(InfoPacket {
version: WISP_VERSION,
extensions,
}),
}
}
fn parse_packet(packet_type: u8, mut bytes: Bytes) -> Result<Self, WispError> {
use PacketType as P;
Ok(Self {
stream_id: bytes.get_u32_le(),
packet_type: match packet_type {
0x01 => P::Connect(ConnectPacket::try_from(bytes)?),
0x02 => P::Data(bytes),
0x03 => P::Continue(ContinuePacket::try_from(bytes)?),
0x04 => P::Close(ClosePacket::try_from(bytes)?),
// 0x05 is handled seperately
_ => return Err(WispError::InvalidPacketType),
},
})
}
pub(crate) fn maybe_parse_info(
frame: Frame,
role: Role,
extension_builders: &[Box<(dyn ProtocolExtensionBuilder + Send + Sync)>],
) -> Result<Self, WispError> {
if !frame.finished {
return Err(WispError::WsFrameNotFinished);
}
if frame.opcode != OpCode::Binary {
return Err(WispError::WsFrameInvalidType);
}
let mut bytes = frame.payload;
if bytes.remaining() < 1 {
return Err(WispError::PacketTooSmall);
}
let packet_type = bytes.get_u8();
if packet_type == 0x05 {
Self::parse_info(bytes, role, extension_builders)
} else {
Self::parse_packet(packet_type, bytes)
}
}
pub(crate) async fn maybe_handle_extension(
frame: Frame,
extensions: &mut [AnyProtocolExtension],
read: &mut (dyn WebSocketRead + Send),
write: &LockedWebSocketWrite,
) -> Result<Option<Self>, WispError> {
if !frame.finished {
return Err(WispError::WsFrameNotFinished);
}
if frame.opcode != OpCode::Binary {
return Err(WispError::WsFrameInvalidType);
}
let mut bytes = frame.payload;
if bytes.remaining() < 1 {
return Err(WispError::PacketTooSmall);
}
let packet_type = bytes.get_u8();
if let Some(extension) = extensions
.iter_mut()
.find(|x| x.get_supported_packets().iter().any(|x| *x == packet_type))
{
extension.handle_packet(bytes, read, write).await?;
Ok(None)
} else {
Ok(Some(Self::parse_packet(packet_type, bytes)?))
}
}
fn parse_info(
mut bytes: Bytes,
role: Role,
extension_builders: &[Box<(dyn ProtocolExtensionBuilder + Send + Sync)>],
) -> Result<Self, WispError> {
// packet type is already read by code that calls this
if bytes.remaining() < 4 + 2 {
return Err(WispError::PacketTooSmall);
}
if bytes.get_u32_le() != 0 {
return Err(WispError::InvalidStreamId);
}
let version = WispVersion {
major: bytes.get_u8(),
minor: bytes.get_u8(),
};
if version.major != WISP_VERSION.major {
return Err(WispError::IncompatibleProtocolVersion);
}
let mut extensions = Vec::new();
while bytes.remaining() > 4 {
// We have some extensions
let id = bytes.get_u8();
let length = usize::try_from(bytes.get_u32_le())?;
if bytes.remaining() < length {
return Err(WispError::PacketTooSmall);
}
if let Some(builder) = extension_builders.iter().find(|x| x.get_id() == id) {
if let Ok(extension) = builder.build_from_bytes(bytes.copy_to_bytes(length), role) {
extensions.push(extension)
}
} else {
bytes.advance(length)
}
}
Ok(Self {
stream_id: 0,
packet_type: PacketType::Info(InfoPacket {
version,
extensions,
}),
})
}
}
impl TryFrom<Bytes> for Packet {
type Error = WispError;
fn try_from(mut bytes: Bytes) -> Result<Self, Self::Error> {
if bytes.remaining() < 5 {
if bytes.remaining() < 1 {
return Err(Self::Error::PacketTooSmall);
}
let packet_type = bytes.get_u8();
use PacketType::*;
Ok(Self {
stream_id: bytes.get_u32_le(),
packet_type: match packet_type {
0x01 => Connect(ConnectPacket::try_from(bytes)?),
0x02 => Data(bytes),
0x03 => Continue(ContinuePacket::try_from(bytes)?),
0x04 => Close(ClosePacket::try_from(bytes)?),
_ => return Err(Self::Error::InvalidPacketType),
},
})
Self::parse_packet(packet_type, bytes)
}
}

8
wisp/src/sink_unfold.rs
View file

@ -1,8 +1,10 @@
//! futures sink unfold with a close function
use core::{future::Future, pin::Pin};
use futures::ready;
use futures::task::{Context, Poll};
use futures::Sink;
use futures::{
ready,
task::{Context, Poll},
Sink,
};
use pin_project_lite::pin_project;
pin_project! {

67
wisp/src/stream.rs
View file

@ -1,13 +1,14 @@
use crate::{sink_unfold, CloseReason, Packet, Role, StreamType, WispError};
use async_io_stream::IoStream;
pub use async_io_stream::IoStream;
use bytes::Bytes;
use event_listener::Event;
use flume as mpsc;
use futures::{
channel::{mpsc, oneshot},
stream,
channel::oneshot,
select, stream,
task::{Context, Poll},
Sink, SinkExt, Stream, StreamExt,
FutureExt, Sink, Stream,
};
use pin_project_lite::pin_project;
use std::{
@ -21,7 +22,13 @@ use std::{
pub(crate) enum WsEvent {
SendPacket(Packet, oneshot::Sender<Result<(), WispError>>),
Close(Packet, oneshot::Sender<Result<(), WispError>>),
EndFut,
CreateStream(
StreamType,
String,
u16,
oneshot::Sender<Result<MuxStream, WispError>>,
),
EndFut(Option<CloseReason>),
}
/// Read side of a multiplexor stream.
@ -32,8 +39,9 @@ pub struct MuxStreamRead {
pub stream_type: StreamType,
role: Role,
tx: mpsc::Sender<WsEvent>,
rx: mpsc::UnboundedReceiver<Bytes>,
rx: mpsc::Receiver<Bytes>,
is_closed: Arc<AtomicBool>,
is_closed_event: Arc<Event>,
flow_control: Arc<AtomicU32>,
flow_control_read: AtomicU32,
target_flow_control: u32,
@ -45,13 +53,16 @@ impl MuxStreamRead {
if self.is_closed.load(Ordering::Acquire) {
return None;
}
let bytes = self.rx.next().await?;
let bytes = select! {
x = self.rx.recv_async() => x.ok()?,
_ = self.is_closed_event.listen().fuse() => return None
};
if self.role == Role::Server && self.stream_type == StreamType::Tcp {
let val = self.flow_control_read.fetch_add(1, Ordering::AcqRel) + 1;
if val > self.target_flow_control {
let (tx, rx) = oneshot::channel::<Result<(), WispError>>();
self.tx
.send(WsEvent::SendPacket(
.send_async(WsEvent::SendPacket(
Packet::new_continue(
self.stream_id,
self.flow_control.fetch_add(val, Ordering::AcqRel) + val,
@ -101,13 +112,13 @@ impl MuxStreamWrite {
}
let (tx, rx) = oneshot::channel::<Result<(), WispError>>();
self.tx
.send(WsEvent::SendPacket(
.send_async(WsEvent::SendPacket(
Packet::new_data(self.stream_id, data),
tx,
))
.await
.map_err(|x| WispError::Other(Box::new(x)))?;
rx.await.map_err(|x| WispError::Other(Box::new(x)))??;
.map_err(|_| WispError::MuxMessageFailedToSend)?;
rx.await.map_err(|_| WispError::MuxMessageFailedToRecv)??;
if self.role == Role::Client && self.stream_type == StreamType::Tcp {
self.flow_control.store(
self.flow_control.load(Ordering::Acquire).saturating_sub(1),
@ -145,13 +156,13 @@ impl MuxStreamWrite {
let (tx, rx) = oneshot::channel::<Result<(), WispError>>();
self.tx
.send(WsEvent::Close(
.send_async(WsEvent::Close(
Packet::new_close(self.stream_id, reason),
tx,
))
.await
.map_err(|x| WispError::Other(Box::new(x)))?;
rx.await.map_err(|x| WispError::Other(Box::new(x)))??;
.map_err(|_| WispError::MuxMessageFailedToSend)?;
rx.await.map_err(|_| WispError::MuxMessageFailedToRecv)??;
Ok(())
}
@ -173,6 +184,19 @@ impl MuxStreamWrite {
}
}
impl Drop for MuxStreamWrite {
fn drop(&mut self) {
if !self.is_closed.load(Ordering::Acquire) {
self.is_closed.store(true, Ordering::Release);
let (tx, _) = oneshot::channel();
let _ = self.tx.send(WsEvent::Close(
Packet::new_close(self.stream_id, CloseReason::Unknown),
tx,
));
}
}
}
/// Multiplexor stream.
pub struct MuxStream {
/// ID of the stream.
@ -187,9 +211,10 @@ impl MuxStream {
stream_id: u32,
role: Role,
stream_type: StreamType,
rx: mpsc::UnboundedReceiver<Bytes>,
rx: mpsc::Receiver<Bytes>,
tx: mpsc::Sender<WsEvent>,
is_closed: Arc<AtomicBool>,
is_closed_event: Arc<Event>,
flow_control: Arc<AtomicU32>,
continue_recieved: Arc<Event>,
target_flow_control: u32,
@ -203,6 +228,7 @@ impl MuxStream {
tx: tx.clone(),
rx,
is_closed: is_closed.clone(),
is_closed_event: is_closed_event.clone(),
flow_control: flow_control.clone(),
flow_control_read: AtomicU32::new(0),
target_flow_control,
@ -282,13 +308,13 @@ impl MuxStreamCloser {
let (tx, rx) = oneshot::channel::<Result<(), WispError>>();
self.close_channel
.send(WsEvent::Close(
.send_async(WsEvent::Close(
Packet::new_close(self.stream_id, reason),
tx,
))
.await
.map_err(|x| WispError::Other(Box::new(x)))?;
rx.await.map_err(|x| WispError::Other(Box::new(x)))??;
.map_err(|_| WispError::MuxMessageFailedToSend)?;
rx.await.map_err(|_| WispError::MuxMessageFailedToRecv)??;
Ok(())
}
@ -317,7 +343,10 @@ impl MuxStreamIo {
impl Stream for MuxStreamIo {
type Item = Result<Vec<u8>, std::io::Error>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.project().rx.poll_next(cx).map(|x| x.map(|x| Ok(x.to_vec())))
self.project()
.rx
.poll_next(cx)
.map(|x| x.map(|x| Ok(x.to_vec())))
}
}

52
wisp/src/ws.rs
View file

@ -4,9 +4,10 @@
//! for other WebSocket implementations.
//!
//! [`fastwebsockets`]: https://github.com/MercuryWorkshop/epoxy-tls/blob/multiplexed/wisp/src/fastwebsockets.rs
use crate::WispError;
use async_trait::async_trait;
use bytes::Bytes;
use futures::lock::Mutex;
use std::sync::Arc;
/// Opcode of the WebSocket frame.
#[derive(Debug, PartialEq, Clone, Copy)]
@ -64,40 +65,55 @@ impl Frame {
}
/// Generic WebSocket read trait.
#[async_trait]
pub trait WebSocketRead {
/// Read a frame from the socket.
fn wisp_read_frame(
&mut self,
tx: &crate::ws::LockedWebSocketWrite<impl crate::ws::WebSocketWrite>,
) -> impl std::future::Future<Output = Result<Frame, crate::WispError>>;
async fn wisp_read_frame(&mut self, tx: &LockedWebSocketWrite) -> Result<Frame, WispError>;
}
/// Generic WebSocket write trait.
#[async_trait]
pub trait WebSocketWrite {
/// Write a frame to the socket.
fn wisp_write_frame(
&mut self,
frame: Frame,
) -> impl std::future::Future<Output = Result<(), crate::WispError>>;
async fn wisp_write_frame(&mut self, frame: Frame) -> Result<(), WispError>;
/// Close the socket.
async fn wisp_close(&mut self) -> Result<(), WispError>;
}
/// Locked WebSocket that can be shared between threads.
pub struct LockedWebSocketWrite<S>(Arc<Mutex<S>>);
/// Locked WebSocket.
pub struct LockedWebSocketWrite(Mutex<Box<dyn WebSocketWrite + Send>>);
impl<S: WebSocketWrite> LockedWebSocketWrite<S> {
impl LockedWebSocketWrite {
/// Create a new locked websocket.
pub fn new(ws: S) -> Self {
Self(Arc::new(Mutex::new(ws)))
pub fn new(ws: Box<dyn WebSocketWrite + Send>) -> Self {
Self(Mutex::new(ws))
}
/// Write a frame to the websocket.
pub async fn write_frame(&self, frame: Frame) -> Result<(), crate::WispError> {
pub async fn write_frame(&self, frame: Frame) -> Result<(), WispError> {
self.0.lock().await.wisp_write_frame(frame).await
}
/// Close the websocket.
pub async fn close(&self) -> Result<(), WispError> {
self.0.lock().await.wisp_close().await
}
}
impl<S: WebSocketWrite> Clone for LockedWebSocketWrite<S> {
fn clone(&self) -> Self {
Self(self.0.clone())
pub(crate) struct AppendingWebSocketRead<R>(pub Vec<Frame>, pub R)
where
R: WebSocketRead + Send;
#[async_trait]
impl<R> WebSocketRead for AppendingWebSocketRead<R>
where
R: WebSocketRead + Send,
{
async fn wisp_read_frame(&mut self, tx: &LockedWebSocketWrite) -> Result<Frame, WispError> {
if let Some(x) = self.0.pop() {
return Ok(x);
}
return self.1.wisp_read_frame(tx).await;
}
}