uitest/src/ws.rs
//! Copied from gunnar-ui/uitest/src/ws.rs (one implementation: the same zero-dependency client, not a rewrite). Keep the two in step by copying, never by editing one alone.
//! A client-side websocket, RFC 6455, text frames only — enough for one CDP
//! connection.
//!
//! What is deliberately NOT here: `Sec-WebSocket-Accept` verification. It is a
//! SHA-1 of the key we just sent, and checking it would mean carrying a SHA-1
//! implementation to defend against a server that answered `101 Switching
//! Protocols` without being a websocket. Chrome on loopback is not that threat.
use std::io::{Read, Write};
use std::net::TcpStream;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
pub struct Ws {
s: TcpStream,
}
/// How long one read waits before it is an error. Generous: a CDP call that
/// navigates and awaits a promise can sit quiet for a while, and a hang here
/// must still end as a row rather than as a wedged suite.
const READ_TIMEOUT: Duration = Duration::from_secs(60);
/// A tiny counter-plus-clock source for the 16-byte handshake nonce and the
/// per-frame mask. RFC 6455 wants a mask that a client cannot be made to
/// predict; nothing on loopback is trying, and this is not a cryptographic
/// claim — it is only kept from being a constant.
fn nonce(n: u64) -> [u8; 16] {
let t = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_nanos() as u64;
let mut out = [0u8; 16];
let mut x = t ^ (n.wrapping_mul(0x9e37_79b9_7f4a_7c15));
for chunk in out.chunks_mut(8) {
// splitmix64
x = x.wrapping_add(0x9e37_79b9_7f4a_7c15);
let mut z = x;
z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
z ^= z >> 31;
chunk.copy_from_slice(&z.to_le_bytes()[..chunk.len()]);
}
out
}
fn b64(data: &[u8]) -> String {
const A: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
let mut o = String::new();
for c in data.chunks(3) {
let b = [c[0], *c.get(1).unwrap_or(&0), *c.get(2).unwrap_or(&0)];
let v = ((b[0] as u32) << 16) | ((b[1] as u32) << 8) | b[2] as u32;
o.push(A[(v >> 18) as usize & 63] as char);
o.push(A[(v >> 12) as usize & 63] as char);
o.push(if c.len() > 1 { A[(v >> 6) as usize & 63] as char } else { '=' });
o.push(if c.len() > 2 { A[v as usize & 63] as char } else { '=' });
}
o
}
impl Ws {
/// `ws://host:port/path`
pub fn connect(url: &str) -> Result<Ws, String> {
let rest = url.strip_prefix("ws://").ok_or_else(|| format!("not a ws:// url: {url}"))?;
let (hostport, path) = match rest.find('/') {
Some(i) => (&rest[..i], &rest[i..]),
None => (rest, "/"),
};
let s = TcpStream::connect(hostport).map_err(|e| format!("connect {hostport}: {e}"))?;
s.set_nodelay(true).ok();
s.set_read_timeout(Some(READ_TIMEOUT)).ok();
let mut ws = Ws { s };
let key = b64(&nonce(1));
let req = format!(
"GET {path} HTTP/1.1\r\nHost: {hostport}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\
Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n\r\n"
);
ws.s.write_all(req.as_bytes()).map_err(|e| format!("handshake write: {e}"))?;
let mut head = Vec::new();
let mut byte = [0u8; 1];
while !head.ends_with(b"\r\n\r\n") {
let n = ws.s.read(&mut byte).map_err(|e| format!("handshake read: {e}"))?;
if n == 0 {
return Err("the server closed during the handshake".into());
}
head.push(byte[0]);
if head.len() > 8192 {
return Err("handshake response is not a header".into());
}
}
let head = String::from_utf8_lossy(&head);
if !head.starts_with("HTTP/1.1 101") {
return Err(format!("upgrade refused: {}", head.lines().next().unwrap_or("")));
}
Ok(ws)
}
pub fn send_text(&mut self, text: &str) -> Result<(), String> {
let payload = text.as_bytes();
let mut f: Vec<u8> = Vec::with_capacity(payload.len() + 14);
f.push(0x81); // FIN | text
let mask = nonce(payload.len() as u64);
let mask = &mask[..4];
match payload.len() {
n if n < 126 => f.push(0x80 | n as u8),
n if n <= u16::MAX as usize => {
f.push(0x80 | 126);
f.extend_from_slice(&(n as u16).to_be_bytes());
}
n => {
f.push(0x80 | 127);
f.extend_from_slice(&(n as u64).to_be_bytes());
}
}
f.extend_from_slice(mask);
f.extend(payload.iter().enumerate().map(|(i, b)| b ^ mask[i % 4]));
self.s.write_all(&f).map_err(|e| format!("send: {e}"))?;
self.s.flush().map_err(|e| format!("flush: {e}"))
}
/// One complete text message if one BEGINS to arrive within `wait`, else
/// `None`. Nothing is sent, and nothing is lost: the wait is a `peek`, and
/// once a byte is there the whole message is read as [`Ws::recv_text`]
/// reads it, with the ordinary timeout.
pub fn recv_text_within(&mut self, wait: Duration) -> Result<Option<String>, String> {
self.s.set_read_timeout(Some(wait.max(Duration::from_millis(1)))).map_err(|e| format!("set timeout: {e}"))?;
let mut one = [0u8; 1];
let peeked = self.s.peek(&mut one);
self.s.set_read_timeout(Some(READ_TIMEOUT)).map_err(|e| format!("set timeout: {e}"))?;
match peeked {
Ok(0) => Err("the peer closed the websocket".into()),
Ok(_) => self.recv_text().map(Some),
Err(e) if matches!(e.kind(), std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut) => Ok(None),
Err(e) => Err(format!("peek: {e}")),
}
}
fn exactly(&mut self, n: usize) -> Result<Vec<u8>, String> {
let mut v = vec![0u8; n];
self.s.read_exact(&mut v).map_err(|e| format!("read {n}: {e}"))?;
Ok(v)
}
/// One complete text message. Continuation frames are joined; a ping is
/// answered and the read continues; a close frame is an error, because for
/// this client the peer going away always means the situation failed.
pub fn recv_text(&mut self) -> Result<String, String> {
let mut msg: Vec<u8> = Vec::new();
loop {
let h = self.exactly(2)?;
let fin = h[0] & 0x80 != 0;
let opcode = h[0] & 0x0f;
let masked = h[1] & 0x80 != 0;
let len = match h[1] & 0x7f {
126 => {
let b = self.exactly(2)?;
u16::from_be_bytes([b[0], b[1]]) as usize
}
127 => {
let b = self.exactly(8)?;
u64::from_be_bytes(b.try_into().unwrap()) as usize
}
n => n as usize,
};
let mask = if masked { Some(self.exactly(4)?) } else { None };
let mut body = self.exactly(len)?;
if let Some(m) = mask {
for (i, b) in body.iter_mut().enumerate() {
*b ^= m[i % 4];
}
}
match opcode {
0x0 | 0x1 => {
msg.extend_from_slice(&body);
if fin {
return String::from_utf8(msg).map_err(|e| format!("not utf-8: {e}"));
}
}
0x8 => return Err("the peer closed the websocket".into()),
0x9 => {
// pong, masked, same payload
let mask = nonce(9);
let mut f = vec![0x8a, 0x80 | body.len() as u8];
f.extend_from_slice(&mask[..4]);
f.extend(body.iter().enumerate().map(|(i, b)| b ^ mask[i % 4]));
self.s.write_all(&f).map_err(|e| format!("pong: {e}"))?;
}
0xa => {}
other => return Err(format!("unexpected opcode {other}")),
}
}
}
}