133 lines
3.9 KiB
Rust
133 lines
3.9 KiB
Rust
//! A scripted WebSocket server for tests: it accepts one handshake per connection and then lets
|
|
//! the test send raw frames and read the client's. Built on `tls_server`. Do not edit.
|
|
|
|
#![allow(dead_code)] // each test file uses a different part of this module
|
|
|
|
use std::io::{Read, Write};
|
|
use std::net::SocketAddr;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use gatewayd::ws::handshake::accept_for;
|
|
use rustls::ServerConfig;
|
|
|
|
use crate::tls_server::{Conn, serve};
|
|
|
|
pub struct Peer {
|
|
pub conn: Box<dyn Conn>,
|
|
/// The request head the client sent, for tests that check it.
|
|
pub request: String,
|
|
}
|
|
|
|
/// A frame from the client: opcode, whether it was masked, and the unmasked payload.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub struct ClientFrame {
|
|
pub opcode: u8,
|
|
pub masked: bool,
|
|
pub mask: [u8; 4],
|
|
pub payload: Vec<u8>,
|
|
}
|
|
|
|
impl Peer {
|
|
pub fn send(&mut self, bytes: &[u8]) {
|
|
let _ = self.conn.write_all(bytes);
|
|
let _ = self.conn.flush();
|
|
}
|
|
|
|
/// An unmasked server frame, FIN set.
|
|
pub fn frame(&mut self, opcode: u8, payload: &[u8]) {
|
|
let mut out = vec![0x80 | opcode];
|
|
if payload.len() < 126 {
|
|
out.push(payload.len() as u8);
|
|
} else {
|
|
out.push(126);
|
|
out.extend_from_slice(&(payload.len() as u16).to_be_bytes());
|
|
}
|
|
out.extend_from_slice(payload);
|
|
self.send(&out);
|
|
}
|
|
|
|
pub fn text(&mut self, text: &str) {
|
|
self.frame(0x1, text.as_bytes());
|
|
}
|
|
|
|
fn read_exact(&mut self, n: usize) -> Option<Vec<u8>> {
|
|
let mut buf = vec![0u8; n];
|
|
self.conn.read_exact(&mut buf).ok()?;
|
|
Some(buf)
|
|
}
|
|
|
|
/// The next frame from the client, or `None` when it has gone.
|
|
pub fn read_frame(&mut self) -> Option<ClientFrame> {
|
|
let head = self.read_exact(2)?;
|
|
let opcode = head[0] & 0x0F;
|
|
let masked = head[1] & 0x80 != 0;
|
|
let len = match head[1] & 0x7F {
|
|
126 => u16::from_be_bytes(self.read_exact(2)?.try_into().ok()?) as usize,
|
|
127 => u64::from_be_bytes(self.read_exact(8)?.try_into().ok()?) as usize,
|
|
n => n as usize,
|
|
};
|
|
let mask: [u8; 4] = if masked {
|
|
self.read_exact(4)?.try_into().ok()?
|
|
} else {
|
|
[0; 4]
|
|
};
|
|
let raw = self.read_exact(len)?;
|
|
let payload = raw
|
|
.iter()
|
|
.zip(mask.iter().cycle())
|
|
.map(|(b, m)| b ^ m)
|
|
.collect();
|
|
Some(ClientFrame {
|
|
opcode,
|
|
masked,
|
|
mask,
|
|
payload,
|
|
})
|
|
}
|
|
|
|
pub fn pause(&self, d: Duration) {
|
|
std::thread::sleep(d);
|
|
}
|
|
}
|
|
|
|
/// Serve WebSocket connections: complete the handshake (or answer `refuse_with` instead), then run
|
|
/// `script` on the connection.
|
|
pub fn serve_ws<F>(
|
|
tls: Option<Arc<ServerConfig>>,
|
|
refuse_with: Option<&'static str>,
|
|
script: F,
|
|
) -> SocketAddr
|
|
where
|
|
F: Fn(Peer) + Send + Sync + 'static,
|
|
{
|
|
serve(tls, move |mut conn| {
|
|
let mut head = Vec::new();
|
|
let mut byte = [0u8; 1];
|
|
while !head.ends_with(b"\r\n\r\n") {
|
|
if conn.read(&mut byte).map(|n| n == 0).unwrap_or(true) {
|
|
return;
|
|
}
|
|
head.push(byte[0]);
|
|
}
|
|
let request = String::from_utf8_lossy(&head).into_owned();
|
|
if let Some(reply) = refuse_with {
|
|
let _ = conn.write_all(reply.as_bytes());
|
|
return;
|
|
}
|
|
let key = request
|
|
.lines()
|
|
.find_map(|l| l.strip_prefix("Sec-WebSocket-Key: "))
|
|
.unwrap_or_default()
|
|
.trim()
|
|
.to_string();
|
|
let reply = format!(
|
|
"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {}\r\n\r\n",
|
|
accept_for(&key)
|
|
);
|
|
let _ = conn.write_all(reply.as_bytes());
|
|
let _ = conn.flush();
|
|
script(Peer { conn, request });
|
|
})
|
|
}
|