use super::*; use futures_util::{SinkExt, StreamExt}; use tokio_tungstenite::tungstenite::{ Message as Frame, protocol::frame::{ Frame as RawFrame, coding::{Data, OpCode}, }, }; #[test] #[ignore = "explicit release WebSocket handshake performance probe"] fn websocket_upgrade_cost() { runtime() .run(async |scope| { let (listener, serving) = server(scope).await; let port = listener.port().unwrap(); println!("\ncase,sample,iterations,elapsed_ns"); for sample in 0..16 { let started = std::time::Instant::now(); for _ in 0..20 { let socket = ws::connect("127.0.0.1", port, "/echo", &limits(), scope) .await .unwrap(); socket.send(Message::Text("check".into())).await.unwrap(); assert_eq!( socket.receive().await.unwrap(), Some(Message::Text("check".into())) ); socket.close(1000, "").await.unwrap(); } let elapsed = started.elapsed().as_nanos(); if sample != 0 { println!("websocket_upgrade,{},20,{elapsed}", sample - 1); } } listener.close().unwrap(); serving.wait().await.unwrap().unwrap(); }) .unwrap(); } #[test] fn websocket_client_limits_handshake_without_limiting_later_frames() { runtime() .run(async |scope| { for header_bytes in [60000, 66000] { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); let peer = scope .spawn_leaf(move || async move { let (stream, _) = listener.accept().await.unwrap(); let handshake = tokio_tungstenite::accept_hdr_async( stream, #[expect(clippy::result_large_err, reason = "Tungstenite 的握手回调要求直接返回 HTTP 错误响应")] move |_: &tokio_tungstenite::tungstenite::handshake::server::Request, mut response: tokio_tungstenite::tungstenite::handshake::server::Response| { response.headers_mut().insert( "x-padding", "x".repeat(header_bytes).parse().unwrap(), ); Ok(response) }, ) .await; if header_bytes == 60000 { let mut wire = handshake.unwrap(); wire.send(Frame::Binary(vec![17; 70000].into())).await.unwrap(); assert!(matches!(wire.next().await, Some(Ok(Frame::Close(_))))); wire.flush().await.unwrap(); } }) .unwrap(); let connection = ws::connect("127.0.0.1", i64::from(port), "/", &limits(), scope).await; if header_bytes == 60000 { let socket = connection.unwrap(); assert_eq!(socket.receive().await.unwrap(), Some(Message::Binary(vec![17; 70000].into()))); socket.close(1000, "").await.unwrap(); } else { assert!(connection.is_err()); } peer.wait().await.unwrap(); } }) .unwrap(); } #[test] fn websocket_client_reassembles_fragments_and_rejects_cumulative_limit_and_utf8() { runtime() .run(async |scope| { for invalid in [false, true] { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); let peer = scope .spawn_leaf(move || async move { let (stream, _) = listener.accept().await.unwrap(); let mut wire = tokio_tungstenite::accept_async(stream).await.unwrap(); wire.send(Frame::Frame(RawFrame::message( vec![0xe4, 0xb8], OpCode::Data(Data::Text), false, ))) .await .unwrap(); wire.send(Frame::Frame(RawFrame::message( vec![0xad], OpCode::Data(Data::Continue), true, ))) .await .unwrap(); if invalid { wire.send(Frame::Frame(RawFrame::message( vec![0xff], OpCode::Data(Data::Text), true, ))) .await .unwrap(); } else { wire.send(Frame::Frame(RawFrame::message( b"12345".to_vec(), OpCode::Data(Data::Text), false, ))) .await .unwrap(); wire.send(Frame::Frame(RawFrame::message( b"6789".to_vec(), OpCode::Data(Data::Continue), true, ))) .await .unwrap(); } let _ = wire.next().await; }) .unwrap(); let socket = ws::connect( "127.0.0.1", i64::from(port), "/", &ws::Limits { message_bytes: 8, ..limits() }, scope, ) .await .unwrap(); assert_eq!( socket.receive().await.unwrap(), Some(Message::Text("中".into())) ); assert!(socket.receive().await.is_err()); assert!(socket.send(Message::Text("closed".into())).await.is_err()); peer.wait().await.unwrap(); } }) .unwrap(); } #[test] fn websocket_accept_failure_can_send_http_error_and_http2_cannot_upgrade() { runtime() .run(async |scope| { for http2 in [false, true] { let listener = net::listen("127.0.0.1", 0, scope).await.unwrap(); let port = listener.port().unwrap(); let incoming = listener.clone(); let limits = Limits { http2: http2.then(http::Http2Limits::default), ..Limits::default() }; let server_limits = limits.clone(); let server = scope .spawn(move |child| async move { http::serve_live( &incoming, &server_limits, &http::LiveLimits::default(), &child, |_, reply, scope| async move { assert!( ws::accept(&reply, &super::limits(), &scope).await.is_err() ); reply .respond(Response { status: 400, headers: vec![], body: b"not a websocket".to_vec().into(), }) .await .map_err(|error| error.to_string()) }, ) .await }) .unwrap(); let response = http::send("127.0.0.1", port, request("/", &[]), &limits) .await .unwrap(); assert_eq!(response.status, 400); assert_eq!(response.body.as_ref(), b"not a websocket"); listener.close().unwrap(); server.wait().await.unwrap().unwrap(); } }) .unwrap(); } #[test] fn websocket_upgrade_preserves_prefetched_frames_and_rejects_unmasked_clients() { runtime().run(async |scope| { let (listener, server) = server(scope).await; let port = listener.port().unwrap(); scope.blocking(move || { let mut peer = TcpStream::connect(("127.0.0.1", port as u16)).unwrap(); peer.set_read_timeout(Some(Duration::from_secs(2))).unwrap(); peer.set_write_timeout(Some(Duration::from_secs(2))).unwrap(); let mut request = b"GET /echo HTTP/1.1\r\nHost: localhost\r\nConnection: Upgrade\r\nUpgrade: websocket\r\nSec-WebSocket-Version: 13\r\nSec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\r\n".to_vec(); request.extend([0x81, 0x81, 1, 2, 3, 4, b'x' ^ 1]); peer.write_all(&request).unwrap(); let mut peer = BufReader::new(peer); let mut line = String::new(); peer.read_line(&mut line).unwrap(); assert!(line.starts_with("HTTP/1.1 101"), "{line}"); loop { line.clear(); assert!(peer.read_line(&mut line).unwrap() > 0); if line == "\r\n" { break; } } let mut echo = [0; 3]; peer.read_exact(&mut echo).unwrap(); assert_eq!(echo, [0x81, 1, b'x']); peer.get_mut().write_all(&[0x81, 1, b'y']).unwrap(); assert_eq!(peer.read(&mut echo).unwrap(), 0); }).await.unwrap(); listener.close().unwrap(); server.wait().await.unwrap().unwrap(); }).unwrap(); }