479 lines
16 KiB
Rust
479 lines
16 KiB
Rust
use tokio::io::{AsyncRead, AsyncReadExt};
|
||
|
||
// Frame type constants
|
||
pub const FRAME_OPEN: u8 = 0x01;
|
||
pub const FRAME_DATA: u8 = 0x02;
|
||
pub const FRAME_CLOSE: u8 = 0x03;
|
||
pub const FRAME_DATA_BACK: u8 = 0x04;
|
||
pub const FRAME_CLOSE_BACK: u8 = 0x05;
|
||
pub const FRAME_CONFIG: u8 = 0x06; // Hub -> Edge: configuration update
|
||
pub const FRAME_PING: u8 = 0x07; // Hub -> Edge: heartbeat probe
|
||
pub const FRAME_PONG: u8 = 0x08; // Edge -> Hub: heartbeat response
|
||
pub const FRAME_WINDOW_UPDATE: u8 = 0x09; // Edge -> Hub: per-stream flow control
|
||
pub const FRAME_WINDOW_UPDATE_BACK: u8 = 0x0A; // Hub -> Edge: per-stream flow control
|
||
|
||
// Frame header size: 4 (stream_id) + 1 (type) + 4 (length) = 9 bytes
|
||
pub const FRAME_HEADER_SIZE: usize = 9;
|
||
|
||
// Maximum payload size (16 MB)
|
||
pub const MAX_PAYLOAD_SIZE: u32 = 16 * 1024 * 1024;
|
||
|
||
// Per-stream flow control constants
|
||
/// Initial per-stream window size (4 MB). Sized for full throughput at high RTT:
|
||
/// at 100ms RTT, this sustains ~40 MB/s per stream.
|
||
pub const INITIAL_STREAM_WINDOW: u32 = 4 * 1024 * 1024;
|
||
/// Send WINDOW_UPDATE after consuming this many bytes (half the initial window).
|
||
pub const WINDOW_UPDATE_THRESHOLD: u32 = INITIAL_STREAM_WINDOW / 2;
|
||
/// Maximum window size to prevent overflow.
|
||
pub const MAX_WINDOW_SIZE: u32 = 16 * 1024 * 1024;
|
||
|
||
/// Encode a WINDOW_UPDATE frame for a specific stream.
|
||
pub fn encode_window_update(stream_id: u32, frame_type: u8, increment: u32) -> Vec<u8> {
|
||
encode_frame(stream_id, frame_type, &increment.to_be_bytes())
|
||
}
|
||
|
||
/// Compute the target per-stream window size based on the number of active streams.
|
||
/// Total memory budget is ~32MB shared across all streams. As more streams are active,
|
||
/// each gets a smaller window. This adapts to current demand — few streams get high
|
||
/// throughput, many streams save memory and reduce control frame pressure.
|
||
pub fn compute_window_for_stream_count(active: u32) -> u32 {
|
||
let per_stream = (32 * 1024 * 1024u64) / (active.max(1) as u64);
|
||
per_stream.clamp(64 * 1024, INITIAL_STREAM_WINDOW as u64) as u32
|
||
}
|
||
|
||
/// Decode a WINDOW_UPDATE payload into a byte increment. Returns None if payload is malformed.
|
||
pub fn decode_window_update(payload: &[u8]) -> Option<u32> {
|
||
if payload.len() != 4 {
|
||
return None;
|
||
}
|
||
Some(u32::from_be_bytes([payload[0], payload[1], payload[2], payload[3]]))
|
||
}
|
||
|
||
/// A single multiplexed frame.
|
||
#[derive(Debug, Clone)]
|
||
pub struct Frame {
|
||
pub stream_id: u32,
|
||
pub frame_type: u8,
|
||
pub payload: Vec<u8>,
|
||
}
|
||
|
||
/// Encode a frame into bytes: [stream_id:4][type:1][length:4][payload]
|
||
pub fn encode_frame(stream_id: u32, frame_type: u8, payload: &[u8]) -> Vec<u8> {
|
||
let len = payload.len() as u32;
|
||
let mut buf = Vec::with_capacity(FRAME_HEADER_SIZE + payload.len());
|
||
buf.extend_from_slice(&stream_id.to_be_bytes());
|
||
buf.push(frame_type);
|
||
buf.extend_from_slice(&len.to_be_bytes());
|
||
buf.extend_from_slice(payload);
|
||
buf
|
||
}
|
||
|
||
/// Build a PROXY protocol v1 header line.
|
||
/// Format: `PROXY TCP4 <client_ip> <edge_ip> <client_port> <dest_port>\r\n`
|
||
pub fn build_proxy_v1_header(
|
||
client_ip: &str,
|
||
edge_ip: &str,
|
||
client_port: u16,
|
||
dest_port: u16,
|
||
) -> String {
|
||
format!(
|
||
"PROXY TCP4 {} {} {} {}\r\n",
|
||
client_ip, edge_ip, client_port, dest_port
|
||
)
|
||
}
|
||
|
||
/// Stateful async frame reader that yields `Frame` values from an `AsyncRead`.
|
||
pub struct FrameReader<R> {
|
||
reader: R,
|
||
header_buf: [u8; FRAME_HEADER_SIZE],
|
||
}
|
||
|
||
impl<R: AsyncRead + Unpin> FrameReader<R> {
|
||
pub fn new(reader: R) -> Self {
|
||
Self {
|
||
reader,
|
||
header_buf: [0u8; FRAME_HEADER_SIZE],
|
||
}
|
||
}
|
||
|
||
/// Read the next frame. Returns `None` on EOF, `Err` on protocol violation.
|
||
pub async fn next_frame(&mut self) -> Result<Option<Frame>, std::io::Error> {
|
||
// Read header
|
||
match self.reader.read_exact(&mut self.header_buf).await {
|
||
Ok(_) => {}
|
||
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
|
||
Err(e) => return Err(e),
|
||
}
|
||
|
||
let stream_id = u32::from_be_bytes([
|
||
self.header_buf[0],
|
||
self.header_buf[1],
|
||
self.header_buf[2],
|
||
self.header_buf[3],
|
||
]);
|
||
let frame_type = self.header_buf[4];
|
||
let length = u32::from_be_bytes([
|
||
self.header_buf[5],
|
||
self.header_buf[6],
|
||
self.header_buf[7],
|
||
self.header_buf[8],
|
||
]);
|
||
|
||
if length > MAX_PAYLOAD_SIZE {
|
||
return Err(std::io::Error::new(
|
||
std::io::ErrorKind::InvalidData,
|
||
format!("frame payload too large: {} bytes", length),
|
||
));
|
||
}
|
||
|
||
let mut payload = vec![0u8; length as usize];
|
||
if length > 0 {
|
||
self.reader.read_exact(&mut payload).await?;
|
||
}
|
||
|
||
Ok(Some(Frame {
|
||
stream_id,
|
||
frame_type,
|
||
payload,
|
||
}))
|
||
}
|
||
|
||
/// Consume the reader and return the inner stream.
|
||
pub fn into_inner(self) -> R {
|
||
self.reader
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_encode_frame() {
|
||
let data = b"hello";
|
||
let encoded = encode_frame(42, FRAME_DATA, data);
|
||
assert_eq!(encoded.len(), FRAME_HEADER_SIZE + data.len());
|
||
// stream_id = 42 in BE
|
||
assert_eq!(&encoded[0..4], &42u32.to_be_bytes());
|
||
// frame type
|
||
assert_eq!(encoded[4], FRAME_DATA);
|
||
// length
|
||
assert_eq!(&encoded[5..9], &5u32.to_be_bytes());
|
||
// payload
|
||
assert_eq!(&encoded[9..], b"hello");
|
||
}
|
||
|
||
#[test]
|
||
fn test_encode_empty_frame() {
|
||
let encoded = encode_frame(1, FRAME_CLOSE, &[]);
|
||
assert_eq!(encoded.len(), FRAME_HEADER_SIZE);
|
||
assert_eq!(&encoded[5..9], &0u32.to_be_bytes());
|
||
}
|
||
|
||
#[test]
|
||
fn test_proxy_v1_header() {
|
||
let header = build_proxy_v1_header("1.2.3.4", "5.6.7.8", 12345, 443);
|
||
assert_eq!(header, "PROXY TCP4 1.2.3.4 5.6.7.8 12345 443\r\n");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_frame_reader() {
|
||
let frame1 = encode_frame(1, FRAME_OPEN, b"PROXY TCP4 1.2.3.4 5.6.7.8 1234 443\r\n");
|
||
let frame2 = encode_frame(1, FRAME_DATA, b"GET / HTTP/1.1\r\n");
|
||
let frame3 = encode_frame(1, FRAME_CLOSE, &[]);
|
||
|
||
let mut data = Vec::new();
|
||
data.extend_from_slice(&frame1);
|
||
data.extend_from_slice(&frame2);
|
||
data.extend_from_slice(&frame3);
|
||
|
||
let cursor = std::io::Cursor::new(data);
|
||
let mut reader = FrameReader::new(cursor);
|
||
|
||
let f1 = reader.next_frame().await.unwrap().unwrap();
|
||
assert_eq!(f1.stream_id, 1);
|
||
assert_eq!(f1.frame_type, FRAME_OPEN);
|
||
assert!(f1.payload.starts_with(b"PROXY"));
|
||
|
||
let f2 = reader.next_frame().await.unwrap().unwrap();
|
||
assert_eq!(f2.frame_type, FRAME_DATA);
|
||
|
||
let f3 = reader.next_frame().await.unwrap().unwrap();
|
||
assert_eq!(f3.frame_type, FRAME_CLOSE);
|
||
assert!(f3.payload.is_empty());
|
||
|
||
// EOF
|
||
assert!(reader.next_frame().await.unwrap().is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn test_encode_frame_config_type() {
|
||
let payload = b"{\"listenPorts\":[443]}";
|
||
let encoded = encode_frame(0, FRAME_CONFIG, payload);
|
||
assert_eq!(encoded[4], FRAME_CONFIG);
|
||
assert_eq!(&encoded[0..4], &0u32.to_be_bytes());
|
||
assert_eq!(&encoded[9..], payload.as_slice());
|
||
}
|
||
|
||
#[test]
|
||
fn test_encode_frame_data_back_type() {
|
||
let payload = b"response data";
|
||
let encoded = encode_frame(7, FRAME_DATA_BACK, payload);
|
||
assert_eq!(encoded[4], FRAME_DATA_BACK);
|
||
assert_eq!(&encoded[0..4], &7u32.to_be_bytes());
|
||
assert_eq!(&encoded[5..9], &(payload.len() as u32).to_be_bytes());
|
||
assert_eq!(&encoded[9..], payload.as_slice());
|
||
}
|
||
|
||
#[test]
|
||
fn test_encode_frame_close_back_type() {
|
||
let encoded = encode_frame(99, FRAME_CLOSE_BACK, &[]);
|
||
assert_eq!(encoded[4], FRAME_CLOSE_BACK);
|
||
assert_eq!(&encoded[0..4], &99u32.to_be_bytes());
|
||
assert_eq!(&encoded[5..9], &0u32.to_be_bytes());
|
||
assert_eq!(encoded.len(), FRAME_HEADER_SIZE);
|
||
}
|
||
|
||
#[test]
|
||
fn test_encode_frame_large_stream_id() {
|
||
let encoded = encode_frame(u32::MAX, FRAME_DATA, b"x");
|
||
assert_eq!(&encoded[0..4], &u32::MAX.to_be_bytes());
|
||
assert_eq!(encoded[4], FRAME_DATA);
|
||
assert_eq!(&encoded[5..9], &1u32.to_be_bytes());
|
||
assert_eq!(encoded[9], b'x');
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_frame_reader_max_payload_rejection() {
|
||
let mut data = Vec::new();
|
||
data.extend_from_slice(&1u32.to_be_bytes());
|
||
data.push(FRAME_DATA);
|
||
data.extend_from_slice(&(MAX_PAYLOAD_SIZE + 1).to_be_bytes());
|
||
|
||
let cursor = std::io::Cursor::new(data);
|
||
let mut reader = FrameReader::new(cursor);
|
||
|
||
let result = reader.next_frame().await;
|
||
assert!(result.is_err());
|
||
let err = result.unwrap_err();
|
||
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_frame_reader_eof_mid_header() {
|
||
// Only 5 bytes — not enough for a 9-byte header
|
||
let data = vec![0u8; 5];
|
||
let cursor = std::io::Cursor::new(data);
|
||
let mut reader = FrameReader::new(cursor);
|
||
|
||
// Should return Ok(None) on partial header EOF
|
||
let result = reader.next_frame().await;
|
||
assert!(result.unwrap().is_none());
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_frame_reader_eof_mid_payload() {
|
||
// Full header claiming 100 bytes of payload, but only 10 bytes present
|
||
let mut data = Vec::new();
|
||
data.extend_from_slice(&1u32.to_be_bytes());
|
||
data.push(FRAME_DATA);
|
||
data.extend_from_slice(&100u32.to_be_bytes());
|
||
data.extend_from_slice(&[0xAB; 10]);
|
||
|
||
let cursor = std::io::Cursor::new(data);
|
||
let mut reader = FrameReader::new(cursor);
|
||
|
||
let result = reader.next_frame().await;
|
||
assert!(result.is_err());
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_frame_reader_all_frame_types() {
|
||
let types = [
|
||
FRAME_OPEN,
|
||
FRAME_DATA,
|
||
FRAME_CLOSE,
|
||
FRAME_DATA_BACK,
|
||
FRAME_CLOSE_BACK,
|
||
FRAME_CONFIG,
|
||
FRAME_PING,
|
||
FRAME_PONG,
|
||
];
|
||
|
||
let mut data = Vec::new();
|
||
for (i, &ft) in types.iter().enumerate() {
|
||
let payload = format!("payload_{}", i);
|
||
data.extend_from_slice(&encode_frame(i as u32, ft, payload.as_bytes()));
|
||
}
|
||
|
||
let cursor = std::io::Cursor::new(data);
|
||
let mut reader = FrameReader::new(cursor);
|
||
|
||
for (i, &ft) in types.iter().enumerate() {
|
||
let frame = reader.next_frame().await.unwrap().unwrap();
|
||
assert_eq!(frame.stream_id, i as u32);
|
||
assert_eq!(frame.frame_type, ft);
|
||
assert_eq!(frame.payload, format!("payload_{}", i).as_bytes());
|
||
}
|
||
|
||
assert!(reader.next_frame().await.unwrap().is_none());
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_frame_reader_zero_length_payload() {
|
||
let data = encode_frame(42, FRAME_CLOSE, &[]);
|
||
let cursor = std::io::Cursor::new(data);
|
||
let mut reader = FrameReader::new(cursor);
|
||
|
||
let frame = reader.next_frame().await.unwrap().unwrap();
|
||
assert_eq!(frame.stream_id, 42);
|
||
assert_eq!(frame.frame_type, FRAME_CLOSE);
|
||
assert!(frame.payload.is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn test_encode_frame_ping_pong() {
|
||
// PING: stream_id=0, empty payload (control frame)
|
||
let ping = encode_frame(0, FRAME_PING, &[]);
|
||
assert_eq!(ping[4], FRAME_PING);
|
||
assert_eq!(&ping[0..4], &0u32.to_be_bytes());
|
||
assert_eq!(ping.len(), FRAME_HEADER_SIZE);
|
||
|
||
// PONG: stream_id=0, empty payload (control frame)
|
||
let pong = encode_frame(0, FRAME_PONG, &[]);
|
||
assert_eq!(pong[4], FRAME_PONG);
|
||
assert_eq!(&pong[0..4], &0u32.to_be_bytes());
|
||
assert_eq!(pong.len(), FRAME_HEADER_SIZE);
|
||
}
|
||
|
||
// --- compute_window_for_stream_count tests ---
|
||
|
||
#[test]
|
||
fn test_adaptive_window_zero_streams() {
|
||
// 0 streams treated as 1: 32MB/1 = 32MB → clamped to 4MB max
|
||
assert_eq!(compute_window_for_stream_count(0), INITIAL_STREAM_WINDOW);
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_one_stream() {
|
||
// 32MB/1 = 32MB → clamped to 4MB max
|
||
assert_eq!(compute_window_for_stream_count(1), INITIAL_STREAM_WINDOW);
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_at_max_boundary() {
|
||
// 32MB/8 = 4MB = exactly INITIAL_STREAM_WINDOW
|
||
assert_eq!(compute_window_for_stream_count(8), INITIAL_STREAM_WINDOW);
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_just_below_max() {
|
||
// 32MB/9 = 3,728,270 — first value below INITIAL_STREAM_WINDOW
|
||
let w = compute_window_for_stream_count(9);
|
||
assert!(w < INITIAL_STREAM_WINDOW);
|
||
assert_eq!(w, (32 * 1024 * 1024u64 / 9) as u32);
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_16_streams() {
|
||
// 32MB/16 = 2MB
|
||
assert_eq!(compute_window_for_stream_count(16), 2 * 1024 * 1024);
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_100_streams() {
|
||
// 32MB/100 = 335,544 bytes (~327KB)
|
||
let w = compute_window_for_stream_count(100);
|
||
assert_eq!(w, (32 * 1024 * 1024u64 / 100) as u32);
|
||
assert!(w > 64 * 1024); // above floor
|
||
assert!(w < INITIAL_STREAM_WINDOW as u32); // below ceiling
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_200_streams() {
|
||
// 32MB/200 = 167,772 bytes (~163KB), above 64KB floor
|
||
let w = compute_window_for_stream_count(200);
|
||
assert_eq!(w, (32 * 1024 * 1024u64 / 200) as u32);
|
||
assert!(w > 64 * 1024);
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_500_streams() {
|
||
// 32MB/500 = 67,108 bytes (~65.5KB), just above 64KB floor
|
||
let w = compute_window_for_stream_count(500);
|
||
assert_eq!(w, (32 * 1024 * 1024u64 / 500) as u32);
|
||
assert!(w > 64 * 1024);
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_at_min_boundary() {
|
||
// 32MB/512 = 65,536 = exactly 64KB floor
|
||
assert_eq!(compute_window_for_stream_count(512), 64 * 1024);
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_below_min_clamped() {
|
||
// 32MB/513 = 65,408 → clamped up to 64KB
|
||
assert_eq!(compute_window_for_stream_count(513), 64 * 1024);
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_1000_streams() {
|
||
// 32MB/1000 = 33,554 → clamped to 64KB
|
||
assert_eq!(compute_window_for_stream_count(1000), 64 * 1024);
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_max_u32() {
|
||
// Extreme: u32::MAX streams → tiny value → clamped to 64KB
|
||
assert_eq!(compute_window_for_stream_count(u32::MAX), 64 * 1024);
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_monotonically_decreasing() {
|
||
// Window should decrease (or stay same) as stream count increases
|
||
let mut prev = compute_window_for_stream_count(1);
|
||
for n in [2, 5, 10, 50, 100, 200, 500, 512, 1000] {
|
||
let w = compute_window_for_stream_count(n);
|
||
assert!(w <= prev, "window increased from {} to {} at n={}", prev, w, n);
|
||
prev = w;
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_adaptive_window_total_budget_bounded() {
|
||
// active × per_stream_window should never exceed 32MB (+ clamp overhead for high N)
|
||
for n in [1, 10, 50, 100, 200, 500] {
|
||
let w = compute_window_for_stream_count(n);
|
||
let total = w as u64 * n as u64;
|
||
assert!(total <= 32 * 1024 * 1024, "total {}MB exceeds budget at n={}", total / (1024*1024), n);
|
||
}
|
||
}
|
||
|
||
// --- encode/decode window_update roundtrip ---
|
||
|
||
#[test]
|
||
fn test_window_update_roundtrip() {
|
||
for &increment in &[0u32, 1, 64 * 1024, INITIAL_STREAM_WINDOW, MAX_WINDOW_SIZE, u32::MAX] {
|
||
let frame = encode_window_update(42, FRAME_WINDOW_UPDATE, increment);
|
||
assert_eq!(frame[4], FRAME_WINDOW_UPDATE);
|
||
let decoded = decode_window_update(&frame[FRAME_HEADER_SIZE..]);
|
||
assert_eq!(decoded, Some(increment));
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_window_update_back_roundtrip() {
|
||
let frame = encode_window_update(7, FRAME_WINDOW_UPDATE_BACK, 1234567);
|
||
assert_eq!(frame[4], FRAME_WINDOW_UPDATE_BACK);
|
||
assert_eq!(decode_window_update(&frame[FRAME_HEADER_SIZE..]), Some(1234567));
|
||
}
|
||
|
||
#[test]
|
||
fn test_decode_window_update_malformed() {
|
||
assert_eq!(decode_window_update(&[]), None);
|
||
assert_eq!(decode_window_update(&[0, 0, 0]), None);
|
||
assert_eq!(decode_window_update(&[0, 0, 0, 0, 0]), None);
|
||
}
|
||
}
|