Menu
Public repository / master
AkurAI-Framework / crates/ws/src/client.rs
Latest change 5aea767e0114a6b39d51aa068dc13de51fc610a5 - akurai-ws: WsSender — cloneable send handle for other threads; tolerate a hung-up peer when echoing its close by Ólafur Búi Ólafsson
//! The RFC 6455 §4.1 client opening handshake: connect to a WebSocket server
//! over TCP, send the upgrade request, verify `Sec-WebSocket-Accept`, and hand
//! back a [`WsConn`] in the client role (masked writes, unmasked reads).
//!
//! Plain `ws://` only — TLS terminates elsewhere, as everywhere in the
//! framework. Typical use: a service subscribing to another service's socket
//! (a home hub to Home Assistant, a worker to a job feed).
//!
//! ```ignore
//! let mut conn = akurai_ws::connect_url("ws://127.0.0.1:8123/api/websocket", &[])?;
//! conn.send_text(r#"{"type":"ping"}"#)?;
//! let reply = conn.recv()?;
//! ```
use crate::base64;
use crate::conn::WsConn;
use crate::handshake::{accept_key, WS_VERSION};
use std::io::{self, Read, Write};
use std::net::{TcpStream, ToSocketAddrs};
use std::time::Duration;
/// Largest response head we accept from the server before giving up.
const MAX_HEAD: usize = 16 * 1024;
/// Connect to `ws://host:port/path`. `headers` are extra request headers (an
/// `Authorization`, say). Only `ws://` is supported.
pub fn connect_url(url: &str, headers: &[(&str, &str)]) -> io::Result<WsConn> {
let (host, port, path) =
parse_ws_url(url).ok_or_else(|| invalid("not a ws://host[:port]/path URL"))?;
connect((host.as_str(), port), &host, &path, headers)
}
/// Connect to `addr` and perform the client handshake for `path`, announcing
/// `host` in the `Host` header. Returns the live connection once the server has
/// answered `101` with the right `Sec-WebSocket-Accept`.
pub fn connect<A: ToSocketAddrs>(
addr: A,
host: &str,
path: &str,
headers: &[(&str, &str)],
) -> io::Result<WsConn> {
let mut stream = TcpStream::connect(addr)?;
stream.set_nodelay(true)?;
// The handshake must not hang forever on a silent peer; the data phase is
// long-lived, so the limit is lifted again before the connection is returned.
stream.set_read_timeout(Some(Duration::from_secs(15)))?;
let key = fresh_key()?;
stream.write_all(request_bytes(host, path, &key, headers).as_bytes())?;
stream.flush()?;
let head = read_head(&mut stream)?;
check_response(&head, &key)?;
stream.set_read_timeout(None)?;
WsConn::client(stream)
}
/// Split `ws://host[:port]/path` into its parts; `None` for anything else.
pub(crate) fn parse_ws_url(url: &str) -> Option<(String, u16, String)> {
let rest = url.strip_prefix("ws://")?;
let (authority, path) = match rest.find('/') {
Some(i) => (&rest[..i], &rest[i..]),
None => (rest, "/"),
};
if authority.is_empty() {
return None;
}
let (host, port) = match authority.rsplit_once(':') {
Some((h, p)) if !h.contains(']') || h.ends_with(']') => (h, p.parse().ok()?),
_ => (authority, 80),
};
Some((host.to_string(), port, path.to_string()))
}
/// The §4.1 request: a GET with the upgrade headers and a random 16-byte key.
fn request_bytes(host: &str, path: &str, key: &str, headers: &[(&str, &str)]) -> String {
let mut req = format!(
"GET {path} HTTP/1.1\r\nHost: {host}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\
Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: {WS_VERSION}\r\n"
);
for (name, value) in headers {
req.push_str(name);
req.push_str(": ");
req.push_str(value);
req.push_str("\r\n");
}
req.push_str("\r\n");
req
}
/// 16 random bytes, base64-encoded (§4.1). Randomness guards against caching
/// proxies replaying a handshake; it is not a secret.
fn fresh_key() -> io::Result<String> {
let mut bytes = [0u8; 16];
std::fs::File::open("/dev/urandom")?.read_exact(&mut bytes)?;
Ok(base64::encode(&bytes))
}
/// Read the response head byte by byte up to the blank line, so nothing past
/// it (the first frame, possibly) is swallowed by a buffer we then discard.
fn read_head(stream: &mut TcpStream) -> io::Result<Vec<u8>> {
let mut head = Vec::with_capacity(512);
let mut byte = [0u8; 1];
while !head.ends_with(b"\r\n\r\n") {
if head.len() >= MAX_HEAD {
return Err(invalid("response head too large"));
}
if stream.read(&mut byte)? == 0 {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"server closed during the WebSocket handshake",
));
}
head.push(byte[0]);
}
Ok(head)
}
/// §4.1 client validation: status 101, `Upgrade: websocket`, and an accept
/// token derived from our key.
fn check_response(head: &[u8], key: &str) -> io::Result<()> {
let text = String::from_utf8_lossy(head);
let mut lines = text.split("\r\n");
let status = lines.next().unwrap_or("");
let code = status.split_whitespace().nth(1).unwrap_or("");
if code != "101" {
return Err(invalid(&format!(
"server did not switch protocols: {status}"
)));
}
let mut upgrade_ok = false;
let mut accept = None;
for line in lines {
if let Some((name, value)) = line.split_once(':') {
let value = value.trim();
if name.eq_ignore_ascii_case("Upgrade") {
upgrade_ok = value.eq_ignore_ascii_case("websocket");
} else if name.eq_ignore_ascii_case("Sec-WebSocket-Accept") {
accept = Some(value.to_string());
}
}
}
if !upgrade_ok {
return Err(invalid("missing Upgrade: websocket"));
}
match accept {
Some(a) if a == accept_key(key) => Ok(()),
Some(_) => Err(invalid("Sec-WebSocket-Accept does not match our key")),
None => Err(invalid("missing Sec-WebSocket-Accept")),
}
}
fn invalid(msg: &str) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, msg.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Message;
use std::io::BufRead;
use std::net::TcpListener;
use std::thread;
/// A minimal RFC 6455 server on a loopback port: one connection, the
/// handshake by hand, then `WsConn` in the server role runs `serve`.
fn loopback_server(
serve: impl FnOnce(WsConn) -> io::Result<()> + Send + 'static,
) -> (String, thread::JoinHandle<Vec<String>>) {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap().to_string();
let handle = thread::spawn(move || {
let (stream, _) = listener.accept().unwrap();
let mut reader = io::BufReader::new(stream.try_clone().unwrap());
let mut lines = Vec::new();
loop {
let mut line = String::new();
reader.read_line(&mut line).unwrap();
if line == "\r\n" {
break;
}
lines.push(line.trim_end().to_string());
}
let key = lines
.iter()
.find_map(|l| l.strip_prefix("Sec-WebSocket-Key: "))
.unwrap()
.to_string();
let mut w = stream.try_clone().unwrap();
write!(
w,
"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {}\r\n\r\n",
accept_key(&key)
)
.unwrap();
if let Err(e) = serve(WsConn::new(stream).unwrap()) {
eprintln!("loopback server: {e}");
}
lines
});
(addr, handle)
}
#[test]
fn parses_ws_urls() {
assert_eq!(
parse_ws_url("ws://10.0.0.5:8123/api/websocket"),
Some(("10.0.0.5".into(), 8123, "/api/websocket".into()))
);
assert_eq!(
parse_ws_url("ws://host"),
Some(("host".into(), 80, "/".into()))
);
assert_eq!(parse_ws_url("wss://host/x"), None);
assert_eq!(parse_ws_url("ws:///x"), None);
}
#[test]
fn request_has_the_upgrade_headers_and_extras() {
let req = request_bytes("h:1", "/ws", "abc", &[("Authorization", "Bearer t")]);
assert!(req.starts_with("GET /ws HTTP/1.1\r\nHost: h:1\r\n"));
assert!(req.contains("Upgrade: websocket\r\n"));
assert!(req.contains("Connection: Upgrade\r\n"));
assert!(req.contains("Sec-WebSocket-Key: abc\r\n"));
assert!(req.contains("Sec-WebSocket-Version: 13\r\n"));
assert!(req.contains("Authorization: Bearer t\r\n"));
assert!(req.ends_with("\r\n\r\n"));
}
#[test]
fn rejects_a_wrong_accept_token_or_status() {
let bad = b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nSec-WebSocket-Accept: nope\r\n\r\n";
assert!(check_response(bad, "key").is_err());
let denied = b"HTTP/1.1 403 Forbidden\r\n\r\n";
assert!(check_response(denied, "key").is_err());
let good = format!(
"HTTP/1.1 101 Switching Protocols\r\nupgrade: WebSocket\r\nSec-WebSocket-Accept: {}\r\n\r\n",
accept_key("key")
);
assert!(check_response(good.as_bytes(), "key").is_ok());
}
#[test]
fn echoes_text_end_to_end_with_masked_client_frames() {
let (addr, server) = loopback_server(|mut conn| {
while let Some(msg) = conn.recv()? {
match msg {
Message::Text(t) => conn.send_text(&format!("echo:{t}"))?,
Message::Close(_) => break,
_ => {}
}
}
Ok(())
});
let mut conn = connect(&addr, "localhost", "/ws", &[("X-Test", "1")]).unwrap();
assert_eq!(conn.role(), crate::conn::Role::Client);
conn.send_text("halló").unwrap();
assert_eq!(
conn.recv().unwrap(),
Some(Message::Text("echo:halló".into()))
);
conn.close(crate::close::NORMAL, "bye").unwrap();
let request_lines = server.join().unwrap();
assert!(request_lines.contains(&"X-Test: 1".to_string()));
assert!(request_lines.contains(&"Host: localhost".to_string()));
}
#[test]
fn refuses_a_server_that_does_not_upgrade() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap().to_string();
thread::spawn(move || {
let (mut stream, _) = listener.accept().unwrap();
let mut buf = [0u8; 1024];
let _ = stream.read(&mut buf);
stream
.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n")
.unwrap();
});
let err = match connect(&addr, "localhost", "/nope", &[]) {
Ok(_) => panic!("a 404 must not become a connection"),
Err(e) => e,
};
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn a_sender_handle_writes_from_another_thread_without_interleaving() {
let (addr, server) = loopback_server(|mut conn| {
let mut n = 0;
while let Some(msg) = conn.recv()? {
match msg {
Message::Text(t) => {
n += 1;
conn.send_text(&t)?;
}
Message::Close(_) => break,
_ => {}
}
}
assert_eq!(n, 400);
Ok(())
});
let mut conn = connect(&addr, "localhost", "/ws", &[]).unwrap();
let sender = conn.sender();
let writers: Vec<_> = (0..4)
.map(|w| {
let s = sender.clone();
thread::spawn(move || {
for i in 0..100 {
s.send_text(&format!("w{w}:{i}:{}", "x".repeat(300)))
.unwrap();
}
})
})
.collect();
let mut echoed = 0;
while echoed < 400 {
match conn.recv().unwrap() {
Some(Message::Text(t)) => {
assert!(t.starts_with('w') && t.ends_with('x'), "frame intact: {t}");
echoed += 1;
}
other => panic!("unexpected {other:?}"),
}
}
for w in writers {
w.join().unwrap();
}
sender.close(crate::close::NORMAL, "").unwrap();
assert!(matches!(conn.recv().unwrap(), Some(Message::Close(_))));
server.join().unwrap();
}
}