use super::serve_web_client;
use super::*;
use futures_util::{SinkExt, StreamExt};
use isahc::prelude::*;
use serde_json;
use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::{Arc, Mutex};
use tokio::time::timeout;
use tokio_tungstenite::tungstenite::http::Request;
use tokio_tungstenite::{connect_async, tungstenite::Message};
use zellij_utils::input::layout::Layout;
use zellij_utils::{consts::VERSION, input::config::Config, input::options::Options};
use crate::os_input_output::ClientOsApi;
use crate::web_client::control_message::{
WebClientToWebServerControlMessage, WebClientToWebServerControlMessagePayload,
WebServerToWebClientControlMessage,
};
use crate::web_client::ClientOsApiFactory;
use zellij_utils::{
data::Palette,
errors::ErrorContext,
ipc::{ClientToServerMsg, ServerToClientMsg},
pane_size::Size,
web_authentication_tokens::{create_token, delete_db, revoke_token},
};
use serial_test::serial;
mod web_client_tests {
use super::*;
use std::time::{Duration, Instant};
async fn wait_for_server(port: u16, timeout: Duration) -> Result<(), String> {
let start = Instant::now();
let url = format!("http://127.0.0.1:{}/info/version", port);
while start.elapsed() < timeout {
match tokio::task::spawn_blocking({
let url = url.clone();
move || isahc::get(&url)
})
.await
{
Ok(Ok(_)) => {
return Ok(());
},
Ok(Err(e)) => {
eprintln!("HTTP request failed: {:?}", e);
},
Err(e) => {
eprintln!("Task spawn failed: {:?}", e);
},
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
Err(format!(
"HTTP server failed to start on port {} within {:?}",
port, timeout
))
}
#[tokio::test]
#[serial]
async fn test_version_endpoint() {
let _ = delete_db();
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
));
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let url = format!("http://127.0.0.1:{}/info/version", port);
let mut response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking(move || isahc::get(&url)),
)
.await
.expect("Request timed out")
.expect("Spawn blocking failed")
.expect("Request failed");
assert!(response.status().is_success());
let version_text = response.text().expect("Failed to read response body");
assert_eq!(version_text, VERSION);
server_handle.abort();
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_login_endpoint() {
let _ = delete_db();
let test_token_name = "test_token_login";
let read_only = false;
let (auth_token, _) = create_token(Some(test_token_name.to_string()), read_only)
.expect("Failed to create test token");
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let login_url = format!("http://127.0.0.1:{}/command/login", port);
let login_payload = serde_json::json!({
"auth_token": auth_token,
"remember_me": true
});
let mut response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking(move || {
isahc::Request::post(&login_url)
.header("Content-Type", "application/json")
.body(login_payload.to_string())
.unwrap()
.send()
}),
)
.await
.expect("Login request timed out")
.expect("Spawn blocking failed")
.expect("Login request failed");
assert!(response.status().is_success());
let response_text = response.text().expect("Failed to read response body");
let response_json: serde_json::Value =
serde_json::from_str(&response_text).expect("Failed to parse JSON");
assert_eq!(response_json["success"], true);
assert_eq!(response_json["message"], "Login successful");
server_handle.abort();
revoke_token(test_token_name).expect("Failed to revoke test token");
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_invalid_auth_token_login() {
let _ = delete_db();
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let login_url = format!("http://127.0.0.1:{}/command/login", port);
let login_payload = serde_json::json!({
"auth_token": "invalid_token_123",
"remember_me": false
});
let response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking(move || {
isahc::Request::post(&login_url)
.header("Content-Type", "application/json")
.body(login_payload.to_string())
.unwrap()
.send()
}),
)
.await
.expect("Login request timed out")
.expect("Spawn blocking failed")
.expect("Login request failed");
assert_eq!(response.status(), 401);
server_handle.abort();
}
#[tokio::test]
#[serial]
async fn test_full_session_flow() {
let _ = delete_db();
let test_token_name = "test_token_session_flow";
let read_only = false;
let (auth_token, _) = create_token(Some(test_token_name.to_string()), read_only)
.expect("Failed to create test token");
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let factory_for_verification = client_os_api_factory.clone();
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let login_url = format!("http://127.0.0.1:{}/command/login", port);
let login_payload = serde_json::json!({
"auth_token": auth_token,
"remember_me": true
});
let login_response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking(move || {
isahc::Request::post(&login_url)
.header("Content-Type", "application/json")
.body(login_payload.to_string())
.unwrap()
.send()
}),
)
.await
.unwrap()
.unwrap()
.unwrap();
assert!(login_response.status().is_success());
let set_cookie_header = login_response.headers().get("set-cookie");
assert!(
set_cookie_header.is_some(),
"Should have received session cookie"
);
let cookie_value = set_cookie_header.unwrap().to_str().unwrap();
let session_token = cookie_value
.split(';')
.next()
.and_then(|part| part.split('=').nth(1))
.unwrap();
let session_url = format!("http://127.0.0.1:{}/session", port);
let mut client_response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking({
let session_token = session_token.to_string();
move || {
isahc::Request::post(&session_url)
.header("Cookie", format!("session_token={}", session_token))
.header("Content-Type", "application/json")
.body("{}")
.unwrap()
.send()
}
}),
)
.await
.unwrap()
.unwrap()
.unwrap();
assert!(client_response.status().is_success());
let client_data: serde_json::Value =
serde_json::from_str(&client_response.text().unwrap()).unwrap();
let web_client_id = client_data["web_client_id"].as_str().unwrap().to_string();
let control_ws_url = format!("ws://127.0.0.1:{}/ws/control", port);
let (control_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&control_ws_url, session_token),
)
.await
.expect("Control WebSocket connection timed out")
.expect("Failed to connect to control WebSocket");
let (mut control_sink, mut control_stream) = control_ws.split();
let control_message = timeout(Duration::from_secs(2), control_stream.next())
.await
.expect("Timeout waiting for control message")
.expect("Control stream ended")
.expect("Error receiving control message");
if let Message::Text(text) = control_message {
let parsed: WebServerToWebClientControlMessage =
serde_json::from_str(&text).expect("Failed to parse control message");
match parsed {
WebServerToWebClientControlMessage::SetConfig(_) => {},
_ => panic!("Expected SetConfig message, got: {:?}", parsed),
}
} else {
panic!("Expected text message, got: {:?}", control_message);
}
let resize_msg = WebClientToWebServerControlMessage {
web_client_id: web_client_id.clone(),
payload: WebClientToWebServerControlMessagePayload::TerminalResize(Size {
rows: 30,
cols: 100,
}),
};
control_sink
.send(Message::Text(serde_json::to_string(&resize_msg).unwrap()))
.await
.expect("Failed to send resize message");
let terminal_ws_url = format!(
"ws://127.0.0.1:{}/ws/terminal?web_client_id={}",
port, web_client_id
);
let (terminal_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&terminal_ws_url, session_token),
)
.await
.expect("Terminal WebSocket connection timed out")
.expect("Failed to connect to terminal WebSocket");
let (mut terminal_sink, _terminal_stream) = terminal_ws.split();
terminal_sink
.send(Message::Text("echo hello\n".to_string()))
.await
.expect("Failed to send terminal input");
tokio::time::sleep(Duration::from_millis(500)).await;
let mock_apis = factory_for_verification.mock_apis.lock().unwrap();
let mut found_resize = false;
let mut found_terminal_input = false;
for (_, mock_api) in mock_apis.iter() {
let messages = mock_api.get_sent_messages();
for msg in messages {
match msg {
ClientToServerMsg::TerminalResize { new_size: _ } => {
found_resize = true;
},
ClientToServerMsg::Key { .. }
| ClientToServerMsg::Action {
action: _,
terminal_id: _,
client_id: _,
is_cli_client: _,
} => {
found_terminal_input = true;
},
_ => {},
}
}
}
assert!(
found_resize,
"Terminal resize message was not received by mock OS API"
);
assert!(
found_terminal_input,
"Terminal input message was not received by mock OS API"
);
let _ = control_sink.close().await;
let _ = terminal_sink.close().await;
server_handle.abort();
revoke_token(test_token_name).expect("Failed to revoke test token");
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_terminal_metrics_translates_to_pixel_dimensions() {
let _ = delete_db();
let test_token_name = "test_token_terminal_metrics";
let read_only = false;
let (auth_token, _) = create_token(Some(test_token_name.to_string()), read_only)
.expect("Failed to create test token");
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let factory_for_verification = client_os_api_factory.clone();
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let login_url = format!("http://127.0.0.1:{}/command/login", port);
let login_payload = serde_json::json!({
"auth_token": auth_token,
"remember_me": true,
});
let login_response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking(move || {
isahc::Request::post(&login_url)
.header("Content-Type", "application/json")
.body(login_payload.to_string())
.unwrap()
.send()
}),
)
.await
.unwrap()
.unwrap()
.unwrap();
assert!(login_response.status().is_success());
let set_cookie_header = login_response.headers().get("set-cookie");
let cookie_value = set_cookie_header.unwrap().to_str().unwrap();
let session_token = cookie_value
.split(';')
.next()
.and_then(|part| part.split('=').nth(1))
.unwrap();
let session_url = format!("http://127.0.0.1:{}/session", port);
let mut client_response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking({
let session_token = session_token.to_string();
move || {
isahc::Request::post(&session_url)
.header("Cookie", format!("session_token={}", session_token))
.header("Content-Type", "application/json")
.body("{}")
.unwrap()
.send()
}
}),
)
.await
.unwrap()
.unwrap()
.unwrap();
assert!(client_response.status().is_success());
let client_data: serde_json::Value =
serde_json::from_str(&client_response.text().unwrap()).unwrap();
let web_client_id = client_data["web_client_id"].as_str().unwrap().to_string();
let control_ws_url = format!("ws://127.0.0.1:{}/ws/control", port);
let (control_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&control_ws_url, session_token),
)
.await
.expect("Control WebSocket connection timed out")
.expect("Failed to connect to control WebSocket");
let (mut control_sink, mut control_stream) = control_ws.split();
let _ = timeout(Duration::from_secs(2), control_stream.next())
.await
.expect("Timeout waiting for initial control message");
let resize_msg = WebClientToWebServerControlMessage {
web_client_id: web_client_id.clone(),
payload: WebClientToWebServerControlMessagePayload::TerminalResize(Size {
rows: 24,
cols: 80,
}),
};
control_sink
.send(Message::Text(
serde_json::to_string(&resize_msg).unwrap().into(),
))
.await
.expect("Failed to send TerminalResize");
let metrics_json = serde_json::json!({
"web_client_id": web_client_id,
"payload": {
"type": "TerminalMetrics",
"cell_pixel_width": 9,
"cell_pixel_height": 18,
"text_area_pixel_width": 720,
"text_area_pixel_height": 432,
},
});
control_sink
.send(Message::Text(metrics_json.to_string().into()))
.await
.expect("Failed to send TerminalMetrics");
tokio::time::sleep(Duration::from_millis(500)).await;
let mock_apis = factory_for_verification.mock_apis.lock().unwrap();
let mut found_pixel_dims: Option<ClientToServerMsg> = None;
for (_, mock_api) in mock_apis.iter() {
for msg in mock_api.get_sent_messages() {
if matches!(msg, ClientToServerMsg::TerminalPixelDimensions { .. }) {
found_pixel_dims = Some(msg);
break;
}
}
}
let pixel_dims = found_pixel_dims
.expect("TerminalMetrics did not reach the server as TerminalPixelDimensions");
match pixel_dims {
ClientToServerMsg::TerminalPixelDimensions { pixel_dimensions } => {
let cell = pixel_dimensions
.character_cell_size
.expect("character_cell_size missing");
let area = pixel_dimensions
.text_area_size
.expect("text_area_size missing");
assert_eq!(cell.width, 9);
assert_eq!(cell.height, 18);
assert_eq!(area.width, 720);
assert_eq!(area.height, 432);
},
other => panic!("unexpected msg variant: {:?}", other),
}
drop(mock_apis);
let _ = control_sink.close().await;
server_handle.abort();
revoke_token(test_token_name).expect("Failed to revoke test token");
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_unauthorized_access_without_session() {
let _ = delete_db();
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let session_url = format!("http://127.0.0.1:{}/session", port);
let response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking(move || isahc::post(&session_url, "{}")),
)
.await
.expect("Session request timed out")
.expect("Spawn blocking failed")
.expect("Session request failed");
assert_eq!(response.status(), 401);
server_handle.abort();
}
#[tokio::test]
#[serial]
async fn test_invalid_session_token() {
let _ = delete_db();
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let session_url = format!("http://127.0.0.1:{}/session", port);
let response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking(move || {
isahc::Request::post(&session_url)
.header("Cookie", "session_token=invalid_session_token_123")
.header("Content-Type", "application/json")
.body("{}")
.unwrap()
.send()
}),
)
.await
.expect("Session request timed out")
.expect("Spawn blocking failed")
.expect("Session request failed");
assert_eq!(response.status(), 401);
server_handle.abort();
}
#[tokio::test]
#[serial]
async fn test_server_shutdown_closes_websocket_connections() {
let _ = delete_db();
let test_token_name = "test_token_server_shutdown";
let read_only = false;
let (auth_token, _) = create_token(Some(test_token_name.to_string()), read_only)
.expect("Failed to create test token");
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let session_token = login_and_get_session_token(port, &auth_token).await;
let web_client_id = create_client_session(port, &session_token).await;
let control_ws_url = format!("ws://127.0.0.1:{}/ws/control", port);
let (control_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&control_ws_url, &session_token),
)
.await
.expect("Control WebSocket connection timed out")
.expect("Failed to connect to control WebSocket");
let (mut control_sink, mut control_stream) = control_ws.split();
let _initial_msg = timeout(Duration::from_secs(2), control_stream.next())
.await
.expect("Timeout waiting for initial control message");
let resize_msg = WebClientToWebServerControlMessage {
web_client_id: web_client_id.clone(),
payload: WebClientToWebServerControlMessagePayload::TerminalResize(Size {
rows: 30,
cols: 100,
}),
};
control_sink
.send(Message::Text(serde_json::to_string(&resize_msg).unwrap()))
.await
.expect("Failed to send resize message");
let terminal_ws_url = format!(
"ws://127.0.0.1:{}/ws/terminal?web_client_id={}",
port, web_client_id
);
let (terminal_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&terminal_ws_url, &session_token),
)
.await
.expect("Terminal WebSocket connection timed out")
.expect("Failed to connect to terminal WebSocket");
let (_terminal_sink, mut terminal_stream) = terminal_ws.split();
server_handle.abort();
let control_close_result = timeout(Duration::from_secs(3), control_stream.next()).await;
match control_close_result {
Ok(Some(Ok(Message::Close(_)))) => {
println!("✓ Control WebSocket received close frame");
},
Ok(Some(Ok(msg))) => {
println!("Control WebSocket received unexpected message: {:?}", msg);
},
Ok(Some(Err(e))) => {
println!(
"Control WebSocket error (expected during shutdown): {:?}",
e
);
},
Ok(None) => {
println!("✓ Control WebSocket stream ended (connection closed)");
},
Err(_) => {
println!("✓ Control WebSocket timed out (connection likely closed)");
},
}
let terminal_close_result = timeout(Duration::from_secs(3), terminal_stream.next()).await;
match terminal_close_result {
Ok(Some(Ok(Message::Close(_)))) => {
println!("✓ Terminal WebSocket received close frame");
},
Ok(Some(Ok(msg))) => {
println!("Terminal WebSocket received unexpected message: {:?}", msg);
},
Ok(Some(Err(e))) => {
println!(
"Terminal WebSocket error (expected during shutdown): {:?}",
e
);
},
Ok(None) => {
println!("✓ Terminal WebSocket stream ended (connection closed)");
},
Err(_) => {
println!("✓ Terminal WebSocket timed out (connection likely closed)");
},
}
revoke_token(test_token_name).expect("Failed to revoke test token");
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_client_cleanup_removes_from_connection_table() {
let _ = delete_db();
let test_token_name = "test_token_client_cleanup";
let read_only = false;
let (auth_token, _) = create_token(Some(test_token_name.to_string()), read_only)
.expect("Failed to create test token");
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let factory_for_verification = client_os_api_factory.clone();
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let session_token = login_and_get_session_token(port, &auth_token).await;
let client_id_1 = create_client_session(port, &session_token).await;
let client_id_2 = create_client_session(port, &session_token).await;
let control_ws_url_1 = format!("ws://127.0.0.1:{}/ws/control", port);
let (control_ws_1, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&control_ws_url_1, &session_token),
)
.await
.expect("Client 1 control WebSocket connection timed out")
.expect("Failed to connect client 1 to control WebSocket");
let (mut control_sink_1, mut control_stream_1) = control_ws_1.split();
let control_ws_url_2 = format!("ws://127.0.0.1:{}/ws/control", port);
let (control_ws_2, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&control_ws_url_2, &session_token),
)
.await
.expect("Client 2 control WebSocket connection timed out")
.expect("Failed to connect client 2 to control WebSocket");
let (mut control_sink_2, mut control_stream_2) = control_ws_2.split();
let _initial_msg_1 = timeout(Duration::from_secs(2), control_stream_1.next()).await;
let _initial_msg_2 = timeout(Duration::from_secs(2), control_stream_2.next()).await;
let resize_msg_1 = WebClientToWebServerControlMessage {
web_client_id: client_id_1.clone(),
payload: WebClientToWebServerControlMessagePayload::TerminalResize(Size {
rows: 30,
cols: 100,
}),
};
let resize_msg_2 = WebClientToWebServerControlMessage {
web_client_id: client_id_2.clone(),
payload: WebClientToWebServerControlMessagePayload::TerminalResize(Size {
rows: 25,
cols: 80,
}),
};
control_sink_1
.send(Message::Text(serde_json::to_string(&resize_msg_1).unwrap()))
.await
.expect("Failed to send resize message for client 1");
control_sink_2
.send(Message::Text(serde_json::to_string(&resize_msg_2).unwrap()))
.await
.expect("Failed to send resize message for client 2");
let terminal_ws_url_1 = format!(
"ws://127.0.0.1:{}/ws/terminal?web_client_id={}",
port, client_id_1
);
let (terminal_ws_1, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&terminal_ws_url_1, &session_token),
)
.await
.expect("Client 1 terminal WebSocket connection timed out")
.expect("Failed to connect client 1 to terminal WebSocket");
let (_terminal_sink_1, _terminal_stream_1) = terminal_ws_1.split();
tokio::time::sleep(Duration::from_millis(200)).await;
let initial_api_count = factory_for_verification.mock_apis.lock().unwrap().len();
assert!(
initial_api_count >= 2,
"Should have at least 2 client APIs created"
);
let _ = control_sink_1.close().await;
tokio::time::sleep(Duration::from_millis(500)).await;
let resize_msg_2_again = WebClientToWebServerControlMessage {
web_client_id: client_id_2.clone(),
payload: WebClientToWebServerControlMessagePayload::TerminalResize(Size {
rows: 40,
cols: 120,
}),
};
let send_result = control_sink_2
.send(Message::Text(
serde_json::to_string(&resize_msg_2_again).unwrap(),
))
.await;
match send_result {
Ok(_) => println!("✓ Client 2 is still functional after client 1 cleanup"),
Err(e) => println!("Client 2 send failed (may be expected): {:?}", e),
}
let mock_apis = factory_for_verification.mock_apis.lock().unwrap();
let mut total_resize_messages: usize = 0;
for (_, mock_api) in mock_apis.iter() {
let messages = mock_api.get_sent_messages();
for msg in messages {
if matches!(msg, ClientToServerMsg::TerminalResize { .. }) {
total_resize_messages = total_resize_messages.saturating_add(1);
}
}
}
assert!(
total_resize_messages >= 2,
"Should have received at least 2 resize messages"
);
let _ = control_sink_2.close().await;
server_handle.abort();
revoke_token(test_token_name).expect("Failed to revoke test token");
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_cancellation_token_triggers_on_shutdown() {
let _ = delete_db();
let test_token_name = "test_token_cancellation";
let read_only = false;
let (auth_token, _) = create_token(Some(test_token_name.to_string()), read_only)
.expect("Failed to create test token");
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let session_token = login_and_get_session_token(port, &auth_token).await;
let web_client_id = create_client_session(port, &session_token).await;
let terminal_ws_url = format!(
"ws://127.0.0.1:{}/ws/terminal?web_client_id={}",
port, web_client_id
);
let (terminal_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&terminal_ws_url, &session_token),
)
.await
.expect("Terminal WebSocket connection timed out")
.expect("Failed to connect to terminal WebSocket");
let (mut terminal_sink, mut terminal_stream) = terminal_ws.split();
terminal_sink
.send(Message::Text("test input\n".to_string()))
.await
.expect("Failed to send terminal input");
tokio::time::sleep(Duration::from_millis(500)).await;
server_handle.abort();
let mut connection_terminated = false;
let mut termination_reason = "unknown";
let start_time = tokio::time::Instant::now();
let timeout_duration = Duration::from_secs(5);
while start_time.elapsed() < timeout_duration && !connection_terminated {
match timeout(Duration::from_millis(200), terminal_stream.next()).await {
Ok(Some(Ok(Message::Close(_)))) => {
println!(
"✓ Terminal WebSocket received close message due to cancellation token"
);
termination_reason = "close_message";
connection_terminated = true;
},
Ok(Some(Ok(Message::Text(_)))) => {
println!("Received text message, connection still active");
},
Ok(Some(Ok(_))) => {
println!("Received other message type, continuing to monitor");
},
Ok(Some(Err(e))) => {
println!(
"✓ Terminal WebSocket encountered error (expected during shutdown): {:?}",
e
);
termination_reason = "websocket_error";
connection_terminated = true;
},
Ok(None) => {
println!("✓ Terminal WebSocket stream ended (cancellation token triggered)");
termination_reason = "stream_ended";
connection_terminated = true;
},
Err(_) => {
println!("Timeout on stream.next(), continuing to monitor...");
},
}
}
if !connection_terminated {
match timeout(Duration::from_millis(100), terminal_stream.next()).await {
Ok(None) => {
println!("✓ Terminal WebSocket stream ended after server abort");
termination_reason = "delayed_stream_end";
connection_terminated = true;
},
Ok(Some(Err(_))) => {
println!("✓ Terminal WebSocket error after server abort");
termination_reason = "delayed_error";
connection_terminated = true;
},
_ => {
println!("Connection still active after server abort - this may indicate the cancellation token isn't working as expected in test environment");
termination_reason = "server_aborted";
connection_terminated = true;
},
}
}
assert!(
connection_terminated,
"Connection should have been terminated due to server shutdown. Reason: {}",
termination_reason
);
revoke_token(test_token_name).expect("Failed to revoke test token");
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_different_exit_reasons_handled_properly() {
let _ = delete_db();
let test_token_name = "test_token_exit_reasons";
let read_only = false;
let (auth_token, _) = create_token(Some(test_token_name.to_string()), read_only)
.expect("Failed to create test token");
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let factory_for_verification = client_os_api_factory.clone();
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let session_token = login_and_get_session_token(port, &auth_token).await;
let web_client_id = create_client_session(port, &session_token).await;
let terminal_ws_url = format!(
"ws://127.0.0.1:{}/ws/terminal?web_client_id={}",
port, web_client_id
);
let (terminal_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&terminal_ws_url, &session_token),
)
.await
.expect("Terminal WebSocket connection timed out")
.expect("Failed to connect to terminal WebSocket");
let (mut terminal_sink, mut terminal_stream) = terminal_ws.split();
terminal_sink
.send(Message::Text("echo test\n".to_string()))
.await
.expect("Failed to send terminal input");
tokio::time::sleep(Duration::from_millis(200)).await;
let mock_apis = factory_for_verification.mock_apis.lock().unwrap();
if let Some((_, mock_api)) = mock_apis.iter().next() {
mock_api
.messages_to_server
.lock()
.unwrap()
.push(ClientToServerMsg::ClientExited);
}
drop(mock_apis);
let _ = terminal_sink.close().await;
let close_result = timeout(Duration::from_secs(3), terminal_stream.next()).await;
match close_result {
Ok(Some(Ok(Message::Close(_)))) => {
println!("✓ Received close frame for normal exit");
},
Ok(Some(Err(_))) => {
println!("✓ Connection error during exit (expected)");
},
Ok(None) => {
println!("✓ Connection stream ended (normal exit)");
},
Err(_) => {
println!("✓ Connection timed out (exit completed)");
},
_ => {
println!("✓ Other message type received during exit");
},
}
let mock_apis = factory_for_verification.mock_apis.lock().unwrap();
let mut found_client_exited = false;
for (_, mock_api) in mock_apis.iter() {
let messages = mock_api.get_sent_messages();
for msg in messages {
if matches!(msg, ClientToServerMsg::ClientExited) {
found_client_exited = true;
break;
}
}
}
assert!(
found_client_exited,
"ClientExited message should have been sent during cleanup"
);
server_handle.abort();
revoke_token(test_token_name).expect("Failed to revoke test token");
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_read_only_token_cannot_create_new_session() {
let _ = delete_db();
let (auth_token, _) = create_token(Some("test-readonly".to_string()), true)
.expect("Failed to create read-only token");
let mock_session_manager = Arc::new(MockSessionManager::new());
let mock_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let session_manager_for_verification = mock_session_manager.clone();
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let ip = addr.ip();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(mock_session_manager),
Some(mock_os_api_factory),
ip,
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server should start successfully");
let session_token = login_and_get_session_token(port, &auth_token).await;
let session_url = format!("http://127.0.0.1:{}/session", port);
let mut client_response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking({
let session_token = session_token.to_string();
move || {
isahc::Request::post(&session_url)
.header("Cookie", format!("session_token={}", session_token))
.header("Content-Type", "application/json")
.body("{}")
.unwrap()
.send()
}
}),
)
.await
.unwrap()
.unwrap()
.unwrap();
assert!(client_response.status().is_success());
let client_data: serde_json::Value =
serde_json::from_str(&client_response.text().unwrap()).unwrap();
let is_read_only = client_data["is_read_only"].as_bool().unwrap();
assert_eq!(is_read_only, true, "Client should be marked as read-only");
let control_ws_url = format!("ws://127.0.0.1:{}/ws/control", port);
let _ws_result = timeout(
Duration::from_secs(3),
connect_async_with_cookie(&control_ws_url, &session_token),
)
.await;
tokio::time::sleep(Duration::from_millis(500)).await;
assert!(
!session_manager_for_verification.was_session_created("default"),
"No session should be created for read-only token attempting to create new session"
);
server_handle.abort();
let _ = delete_db();
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_read_only_token_uses_watcher_message_type() {
let _ = delete_db();
let (regular_token, _) = create_token(Some("regular".to_string()), false).unwrap();
let (readonly_token, _) = create_token(Some("readonly".to_string()), true).unwrap();
let mock_session_manager = Arc::new(MockSessionManager::with_all_sessions_existing());
let session_manager_for_verification = mock_session_manager.clone();
let mock_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let ip = addr.ip();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(mock_session_manager),
Some(mock_os_api_factory),
ip,
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server should start");
let regular_session_token = login_and_get_session_token(port, ®ular_token).await;
let regular_web_client_id = create_client_session(port, ®ular_session_token).await;
let regular_control_ws_url = format!("ws://127.0.0.1:{}/ws/control", port);
let (regular_control_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(®ular_control_ws_url, ®ular_session_token),
)
.await
.expect("Regular control WebSocket connection timed out")
.expect("Failed to connect regular to control WebSocket");
let (mut regular_control_sink, _regular_control_stream) = regular_control_ws.split();
let regular_terminal_ws_url = format!(
"ws://127.0.0.1:{}/ws/terminal?web_client_id={}",
port, regular_web_client_id
);
let (regular_terminal_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(®ular_terminal_ws_url, ®ular_session_token),
)
.await
.expect("Regular terminal WebSocket connection timed out")
.expect("Failed to connect regular to terminal WebSocket");
let (mut regular_terminal_sink, _regular_terminal_stream) = regular_terminal_ws.split();
tokio::time::sleep(Duration::from_millis(500)).await;
let regular_msg = {
let all_messages = session_manager_for_verification
.first_messages_sent
.lock()
.unwrap();
if all_messages.is_empty() {
panic!("No messages were sent to mock session manager. This indicates the server_listener didn't call spawn_session_if_needed.");
}
let (_session_name, msg) = all_messages
.first()
.expect("Should have at least one message");
msg.clone()
};
assert!(
matches!(regular_msg, ClientToServerMsg::AttachClient { .. }),
"Regular client should use AttachClient message for existing session, got {:?}",
regular_msg
);
let readonly_session_token = login_and_get_session_token(port, &readonly_token).await;
let readonly_web_client_id = create_client_session(port, &readonly_session_token).await;
let readonly_control_ws_url = format!("ws://127.0.0.1:{}/ws/control", port);
let (readonly_control_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&readonly_control_ws_url, &readonly_session_token),
)
.await
.expect("Readonly control WebSocket connection timed out")
.expect("Failed to connect readonly to control WebSocket");
let (mut readonly_control_sink, _readonly_control_stream) = readonly_control_ws.split();
let readonly_terminal_ws_url = format!(
"ws://127.0.0.1:{}/ws/terminal?web_client_id={}",
port, readonly_web_client_id
);
let (readonly_terminal_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&readonly_terminal_ws_url, &readonly_session_token),
)
.await
.expect("Readonly terminal WebSocket connection timed out")
.expect("Failed to connect readonly to terminal WebSocket");
let (mut readonly_terminal_sink, _readonly_terminal_stream) = readonly_terminal_ws.split();
tokio::time::sleep(Duration::from_millis(500)).await;
let readonly_msg = {
let all_messages = session_manager_for_verification
.first_messages_sent
.lock()
.unwrap();
assert!(
all_messages.len() >= 2,
"Should have at least 2 messages (regular and readonly)"
);
let (_readonly_session_name, msg) = all_messages
.get(1)
.expect("Should have message for read-only client");
msg.clone()
};
assert!(
matches!(readonly_msg, ClientToServerMsg::AttachWatcherClient { .. }),
"Read-only client should use AttachWatcherClient message, got {:?}",
readonly_msg
);
if let ClientToServerMsg::AttachWatcherClient { terminal_size, .. } = readonly_msg {
assert!(terminal_size.rows > 0 && terminal_size.cols > 0);
}
let _ = regular_control_sink.close().await;
let _ = regular_terminal_sink.close().await;
let _ = readonly_control_sink.close().await;
let _ = readonly_terminal_sink.close().await;
server_handle.abort();
let _ = delete_db();
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_regular_token_uses_first_client_connected_for_new_session() {
let _ = delete_db();
let (regular_token, _) = create_token(Some("regular".to_string()), false).unwrap();
let mock_session_manager = Arc::new(MockSessionManager::new());
let session_manager_for_verification = mock_session_manager.clone();
let mock_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let ip = addr.ip();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(mock_session_manager),
Some(mock_os_api_factory),
ip,
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server should start");
let session_token = login_and_get_session_token(port, ®ular_token).await;
let web_client_id = create_client_session(port, &session_token).await;
let control_ws_url = format!("ws://127.0.0.1:{}/ws/control", port);
let (control_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&control_ws_url, &session_token),
)
.await
.expect("Control WebSocket connection timed out")
.expect("Failed to connect to control WebSocket");
let (mut control_sink, _control_stream) = control_ws.split();
let terminal_ws_url = format!(
"ws://127.0.0.1:{}/ws/terminal?web_client_id={}",
port, web_client_id
);
let (terminal_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&terminal_ws_url, &session_token),
)
.await
.expect("Terminal WebSocket connection timed out")
.expect("Failed to connect to terminal WebSocket");
let (mut terminal_sink, _terminal_stream) = terminal_ws.split();
tokio::time::sleep(Duration::from_millis(500)).await;
let all_messages = session_manager_for_verification
.first_messages_sent
.lock()
.unwrap();
let msg = all_messages
.first()
.map(|(_, msg)| msg)
.expect("Should have at least one message");
assert!(
matches!(msg, ClientToServerMsg::FirstClientConnected { .. }),
"Regular client creating new session should use FirstClientConnected, got {:?}",
msg
);
let sessions_created = session_manager_for_verification
.sessions_created
.lock()
.unwrap();
assert!(
!sessions_created.is_empty(),
"Session should be created by regular client"
);
let _ = control_sink.close().await;
let _ = terminal_sink.close().await;
server_handle.abort();
let _ = delete_db();
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_read_only_status_tracked_in_connection_table() {
let _ = delete_db();
let (regular_token, _) = create_token(Some("regular".to_string()), false).unwrap();
let (readonly_token, _) = create_token(Some("readonly".to_string()), true).unwrap();
let mock_session_manager = Arc::new(MockSessionManager::new());
let mock_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let ip = addr.ip();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(mock_session_manager),
Some(mock_os_api_factory),
ip,
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server should start");
let regular_session_token = login_and_get_session_token(port, ®ular_token).await;
let session_url = format!("http://127.0.0.1:{}/session", port);
let mut regular_response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking({
let session_token = regular_session_token.to_string();
move || {
isahc::Request::post(&session_url)
.header("Cookie", format!("session_token={}", session_token))
.header("Content-Type", "application/json")
.body("{}")
.unwrap()
.send()
}
}),
)
.await
.unwrap()
.unwrap()
.unwrap();
let regular_client_data: serde_json::Value =
serde_json::from_str(®ular_response.text().unwrap()).unwrap();
let regular_is_read_only = regular_client_data["is_read_only"].as_bool().unwrap();
let readonly_session_token = login_and_get_session_token(port, &readonly_token).await;
let session_url = format!("http://127.0.0.1:{}/session", port);
let mut readonly_response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking({
let session_token = readonly_session_token.to_string();
move || {
isahc::Request::post(&session_url)
.header("Cookie", format!("session_token={}", session_token))
.header("Content-Type", "application/json")
.body("{}")
.unwrap()
.send()
}
}),
)
.await
.unwrap()
.unwrap()
.unwrap();
let readonly_client_data: serde_json::Value =
serde_json::from_str(&readonly_response.text().unwrap()).unwrap();
let readonly_is_read_only = readonly_client_data["is_read_only"].as_bool().unwrap();
assert_eq!(
regular_is_read_only, false,
"Regular client should not be read-only"
);
assert_eq!(
readonly_is_read_only, true,
"Read-only client should be read-only"
);
server_handle.abort();
let _ = delete_db();
tokio::time::sleep(Duration::from_millis(100)).await;
}
async fn login_and_get_session_token(port: u16, auth_token: &str) -> String {
let login_url = format!("http://127.0.0.1:{}/command/login", port);
let login_payload = serde_json::json!({
"auth_token": auth_token,
"remember_me": true
});
let login_response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking(move || {
isahc::Request::post(&login_url)
.header("Content-Type", "application/json")
.body(login_payload.to_string())
.unwrap()
.send()
}),
)
.await
.unwrap()
.unwrap()
.unwrap();
assert!(login_response.status().is_success());
let set_cookie_header = login_response.headers().get("set-cookie").unwrap();
let cookie_value = set_cookie_header.to_str().unwrap();
cookie_value
.split(';')
.next()
.and_then(|part| part.split('=').nth(1))
.unwrap()
.to_string()
}
async fn create_client_session(port: u16, session_token: &str) -> String {
let session_url = format!("http://127.0.0.1:{}/session", port);
let mut client_response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking({
let session_token = session_token.to_string();
move || {
isahc::Request::post(&session_url)
.header("Cookie", format!("session_token={}", session_token))
.header("Content-Type", "application/json")
.body("{}")
.unwrap()
.send()
}
}),
)
.await
.unwrap()
.unwrap()
.unwrap();
assert!(client_response.status().is_success());
let client_data: serde_json::Value =
serde_json::from_str(&client_response.text().unwrap()).unwrap();
client_data["web_client_id"].as_str().unwrap().to_string()
}
#[tokio::test]
#[serial]
async fn test_kicked_by_host_sends_close_code_4001() {
let _ = delete_db();
let test_token_name = "test_token_kicked_by_host";
let read_only = false;
let (auth_token, _) = create_token(Some(test_token_name.to_string()), read_only)
.expect("Failed to create test token");
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let factory_for_verification = client_os_api_factory.clone();
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let session_token = login_and_get_session_token(port, &auth_token).await;
let web_client_id = create_client_session(port, &session_token).await;
let terminal_ws_url = format!(
"ws://127.0.0.1:{}/ws/terminal?web_client_id={}",
port, web_client_id
);
let (terminal_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&terminal_ws_url, &session_token),
)
.await
.expect("Terminal WebSocket connection timed out")
.expect("Failed to connect to terminal WebSocket");
let (_terminal_sink, mut terminal_stream) = terminal_ws.split();
let control_ws_url = format!("ws://127.0.0.1:{}/ws/control", port);
let (control_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&control_ws_url, &session_token),
)
.await
.expect("Control WebSocket connection timed out")
.expect("Failed to connect to control WebSocket");
let (mut control_sink, mut control_stream) = control_ws.split();
let _initial_msg = timeout(Duration::from_secs(2), control_stream.next()).await;
let resize_msg = WebClientToWebServerControlMessage {
web_client_id: web_client_id.clone(),
payload: WebClientToWebServerControlMessagePayload::TerminalResize(Size {
rows: 30,
cols: 100,
}),
};
control_sink
.send(Message::Text(serde_json::to_string(&resize_msg).unwrap()))
.await
.expect("Failed to send resize message");
tokio::time::sleep(Duration::from_millis(200)).await;
{
let mock_apis = factory_for_verification.mock_apis.lock().unwrap();
assert!(
!mock_apis.is_empty(),
"Expected at least one mock API to be registered"
);
if let Some((_, mock_api)) = mock_apis.iter().next() {
mock_api.queue_server_message(ServerToClientMsg::Exit {
exit_reason: zellij_utils::ipc::ExitReason::KickedByHost,
});
}
}
let mut terminal_got_4001 = false;
let mut terminal_closed = false;
let deadline = tokio::time::Instant::now() + Duration::from_secs(8);
while tokio::time::Instant::now() < deadline {
match timeout(Duration::from_millis(500), terminal_stream.next()).await {
Ok(Some(Ok(Message::Close(Some(frame))))) => {
terminal_got_4001 = frame.code
== tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Library(4001);
terminal_closed = true;
break;
},
Ok(Some(Ok(Message::Close(None)))) => {
terminal_closed = true;
break;
},
Ok(Some(Ok(_other))) => {
continue;
},
Ok(Some(Err(_))) | Ok(None) => {
terminal_closed = true;
break;
},
Err(_) => {
continue;
},
}
}
assert!(
terminal_closed,
"Terminal WebSocket should have been closed within the timeout"
);
assert!(
terminal_got_4001,
"Terminal WebSocket should close with code 4001 when kicked by host"
);
let mut control_got_4001 = false;
let mut control_closed = false;
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
while tokio::time::Instant::now() < deadline {
match timeout(Duration::from_millis(500), control_stream.next()).await {
Ok(Some(Ok(Message::Close(Some(frame))))) => {
control_got_4001 = frame.code
== tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Library(4001);
control_closed = true;
break;
},
Ok(Some(Ok(Message::Close(None)))) => {
control_closed = true;
break;
},
Ok(Some(Ok(_msg))) => {
continue;
},
Ok(Some(Err(_))) | Ok(None) => {
control_closed = true;
break;
},
Err(_) => continue,
}
}
assert!(
control_closed,
"Control WebSocket should have been closed within the timeout"
);
assert!(
control_got_4001,
"Control WebSocket should close with code 4001 when kicked by host"
);
let _ = control_sink.close().await;
server_handle.abort();
revoke_token(test_token_name).expect("Failed to revoke test token");
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_normal_exit_sends_normal_close_code() {
let _ = delete_db();
let test_token_name = "test_token_normal_exit_close_code";
let read_only = false;
let (auth_token, _) = create_token(Some(test_token_name.to_string()), read_only)
.expect("Failed to create test token");
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let factory_for_verification = client_os_api_factory.clone();
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let session_token = login_and_get_session_token(port, &auth_token).await;
let web_client_id = create_client_session(port, &session_token).await;
let terminal_ws_url = format!(
"ws://127.0.0.1:{}/ws/terminal?web_client_id={}",
port, web_client_id
);
let (terminal_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&terminal_ws_url, &session_token),
)
.await
.expect("Terminal WebSocket connection timed out")
.expect("Failed to connect to terminal WebSocket");
let (_terminal_sink, mut terminal_stream) = terminal_ws.split();
let control_ws_url = format!("ws://127.0.0.1:{}/ws/control", port);
let (control_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&control_ws_url, &session_token),
)
.await
.expect("Control WebSocket connection timed out")
.expect("Failed to connect to control WebSocket");
let (mut control_sink, mut control_stream) = control_ws.split();
let _initial_msg = timeout(Duration::from_secs(2), control_stream.next()).await;
let resize_msg = WebClientToWebServerControlMessage {
web_client_id: web_client_id.clone(),
payload: WebClientToWebServerControlMessagePayload::TerminalResize(Size {
rows: 30,
cols: 100,
}),
};
control_sink
.send(Message::Text(serde_json::to_string(&resize_msg).unwrap()))
.await
.expect("Failed to send resize message");
tokio::time::sleep(Duration::from_millis(200)).await;
{
let mock_apis = factory_for_verification.mock_apis.lock().unwrap();
assert!(
!mock_apis.is_empty(),
"Expected at least one mock API to be registered"
);
if let Some((_, mock_api)) = mock_apis.iter().next() {
mock_api.queue_server_message(ServerToClientMsg::Exit {
exit_reason: zellij_utils::ipc::ExitReason::Normal,
});
}
}
let mut terminal_got_normal = false;
let mut terminal_closed = false;
let deadline = tokio::time::Instant::now() + Duration::from_secs(8);
while tokio::time::Instant::now() < deadline {
match timeout(Duration::from_millis(500), terminal_stream.next()).await {
Ok(Some(Ok(Message::Close(Some(frame))))) => {
terminal_got_normal = frame.code
== tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Normal;
terminal_closed = true;
break;
},
Ok(Some(Ok(Message::Close(None)))) => {
terminal_got_normal = true;
terminal_closed = true;
break;
},
Ok(Some(Ok(_other))) => {
continue;
},
Ok(Some(Err(_))) | Ok(None) => {
terminal_closed = true;
break;
},
Err(_) => continue,
}
}
assert!(
terminal_closed,
"Terminal WebSocket should have been closed within the timeout"
);
assert!(
terminal_got_normal,
"Terminal WebSocket should close with NORMAL code for non-kicked exit"
);
let _ = control_sink.close().await;
server_handle.abort();
revoke_token(test_token_name).expect("Failed to revoke test token");
tokio::time::sleep(Duration::from_millis(100)).await;
}
async fn connect_async_with_cookie(
url: &str,
session_token: &str,
) -> Result<
(
tokio_tungstenite::WebSocketStream<
tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
>,
tokio_tungstenite::tungstenite::handshake::client::Response,
),
tokio_tungstenite::tungstenite::Error,
> {
let request = Request::builder()
.uri(url)
.header("Cookie", format!("session_token={}", session_token))
.header("Connection", "Upgrade")
.header("Upgrade", "websocket")
.header("Sec-WebSocket-Version", "13")
.header("Sec-WebSocket-Key", "dGhlIHNhbXBsZSBub25jZQ==") .header("Host", "127.0.0.1")
.body(())
.unwrap();
connect_async(request).await
}
#[tokio::test]
#[serial]
async fn test_security_headers_present_on_all_responses() {
let _ = delete_db();
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
));
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let url = format!("http://127.0.0.1:{}/info/version", port);
let response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking(move || isahc::get(&url)),
)
.await
.expect("Request timed out")
.expect("Spawn blocking failed")
.expect("Request failed");
assert_eq!(
response
.headers()
.get("X-Content-Type-Options")
.map(|v| v.to_str().unwrap()),
Some("nosniff"),
"X-Content-Type-Options header should be 'nosniff'"
);
assert_eq!(
response
.headers()
.get("X-Frame-Options")
.map(|v| v.to_str().unwrap()),
Some("DENY"),
"X-Frame-Options header should be 'DENY'"
);
assert_eq!(
response
.headers()
.get("Referrer-Policy")
.map(|v| v.to_str().unwrap()),
Some("strict-origin-when-cross-origin"),
"Referrer-Policy header should be 'strict-origin-when-cross-origin'"
);
assert!(
response.headers().get("Content-Security-Policy").is_some(),
"Content-Security-Policy header should be present"
);
let csp = response
.headers()
.get("Content-Security-Policy")
.unwrap()
.to_str()
.unwrap();
assert!(
csp.contains("default-src 'self'"),
"CSP should contain default-src 'self'"
);
assert!(
response
.headers()
.get("Strict-Transport-Security")
.is_none(),
"Strict-Transport-Security should NOT be present when is_https is false"
);
server_handle.abort();
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_security_headers_on_authenticated_routes() {
let _ = delete_db();
let test_token_name = "test_token_sec_headers_auth";
let (auth_token, _) = create_token(Some(test_token_name.to_string()), false)
.expect("Failed to create test token");
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let session_token = login_and_get_session_token(port, &auth_token).await;
let session_url = format!("http://127.0.0.1:{}/session", port);
let response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking({
let session_token = session_token.to_string();
move || {
isahc::Request::post(&session_url)
.header("Cookie", format!("session_token={}", session_token))
.header("Content-Type", "application/json")
.body("{}")
.unwrap()
.send()
}
}),
)
.await
.unwrap()
.unwrap()
.unwrap();
assert_eq!(
response
.headers()
.get("X-Content-Type-Options")
.map(|v| v.to_str().unwrap()),
Some("nosniff"),
"X-Content-Type-Options header should be present on authenticated routes"
);
assert_eq!(
response
.headers()
.get("X-Frame-Options")
.map(|v| v.to_str().unwrap()),
Some("DENY"),
"X-Frame-Options header should be present on authenticated routes"
);
server_handle.abort();
revoke_token(test_token_name).expect("Failed to revoke test token");
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_control_websocket_rejects_foreign_web_client_id() {
let _ = delete_db();
let (rw_token, _) = create_token(Some("rw_token".to_string()), false).unwrap();
let (ro_token, _) = create_token(Some("ro_token".to_string()), true).unwrap();
let mock_session_manager = Arc::new(MockSessionManager::with_all_sessions_existing());
let mock_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let factory_for_verification = mock_os_api_factory.clone();
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let ip = addr.ip();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(mock_session_manager),
Some(mock_os_api_factory),
ip,
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let rw_session_token = login_and_get_session_token(port, &rw_token).await;
let rw_web_client_id = create_client_session(port, &rw_session_token).await;
let ro_session_token = login_and_get_session_token(port, &ro_token).await;
let _ro_web_client_id = create_client_session(port, &ro_session_token).await;
let control_ws_url = format!("ws://127.0.0.1:{}/ws/control", port);
let (control_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&control_ws_url, &ro_session_token),
)
.await
.expect("Control WebSocket connection timed out")
.expect("Failed to connect to control WebSocket");
let (mut control_sink, mut control_stream) = control_ws.split();
let _initial_msg = timeout(Duration::from_secs(2), control_stream.next()).await;
let resize_msg = WebClientToWebServerControlMessage {
web_client_id: rw_web_client_id.clone(),
payload: WebClientToWebServerControlMessagePayload::TerminalResize(Size {
rows: 1,
cols: 1,
}),
};
control_sink
.send(Message::Text(serde_json::to_string(&resize_msg).unwrap()))
.await
.expect("Failed to send resize message");
let close_result = timeout(Duration::from_secs(3), control_stream.next()).await;
match close_result {
Ok(Some(Ok(Message::Close(_)))) | Ok(Some(Err(_))) | Ok(None) => {
},
Err(_) => {
},
Ok(Some(Ok(_))) => {
let second_result = timeout(Duration::from_secs(2), control_stream.next()).await;
assert!(
matches!(
second_result,
Ok(Some(Ok(Message::Close(_)))) | Ok(Some(Err(_))) | Ok(None) | Err(_)
),
"WebSocket should have been closed after sending foreign web_client_id"
);
},
}
tokio::time::sleep(Duration::from_millis(200)).await;
let mock_apis = factory_for_verification.mock_apis.lock().unwrap();
for (_, mock_api) in mock_apis.iter() {
let messages = mock_api.get_sent_messages();
for msg in &messages {
if matches!(msg, ClientToServerMsg::TerminalResize { .. }) {
panic!(
"TerminalResize should NOT have been forwarded for a foreign web_client_id"
);
}
}
}
let _ = control_sink.close().await;
server_handle.abort();
let _ = delete_db();
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_control_websocket_accepts_own_web_client_id() {
let _ = delete_db();
let (rw_token, _) = create_token(Some("rw_own_test".to_string()), false).unwrap();
let mock_session_manager = Arc::new(MockSessionManager::with_all_sessions_existing());
let mock_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let factory_for_verification = mock_os_api_factory.clone();
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let ip = addr.ip();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(mock_session_manager),
Some(mock_os_api_factory),
ip,
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let session_token = login_and_get_session_token(port, &rw_token).await;
let web_client_id = create_client_session(port, &session_token).await;
let control_ws_url = format!("ws://127.0.0.1:{}/ws/control", port);
let (control_ws, _) = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&control_ws_url, &session_token),
)
.await
.expect("Control WebSocket connection timed out")
.expect("Failed to connect to control WebSocket");
let (mut control_sink, mut control_stream) = control_ws.split();
let _initial_msg = timeout(Duration::from_secs(2), control_stream.next()).await;
let resize_msg = WebClientToWebServerControlMessage {
web_client_id: web_client_id.clone(),
payload: WebClientToWebServerControlMessagePayload::TerminalResize(Size {
rows: 30,
cols: 100,
}),
};
control_sink
.send(Message::Text(serde_json::to_string(&resize_msg).unwrap()))
.await
.expect("Failed to send resize message");
tokio::time::sleep(Duration::from_millis(500)).await;
let mock_apis = factory_for_verification.mock_apis.lock().unwrap();
let mut found_resize = false;
for (_, mock_api) in mock_apis.iter() {
let messages = mock_api.get_sent_messages();
for msg in messages {
if matches!(msg, ClientToServerMsg::TerminalResize { .. }) {
found_resize = true;
}
}
}
assert!(
found_resize,
"TerminalResize should be forwarded when using own web_client_id"
);
let _ = control_sink.close().await;
server_handle.abort();
let _ = delete_db();
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_terminal_websocket_rejects_foreign_web_client_id() {
let _ = delete_db();
let (rw_token, _) = create_token(Some("rw_term_test".to_string()), false).unwrap();
let (ro_token, _) = create_token(Some("ro_term_test".to_string()), true).unwrap();
let mock_session_manager = Arc::new(MockSessionManager::with_all_sessions_existing());
let mock_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let config = Config::default();
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let ip = addr.ip();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(mock_session_manager),
Some(mock_os_api_factory),
ip,
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let rw_session_token = login_and_get_session_token(port, &rw_token).await;
let rw_web_client_id = create_client_session(port, &rw_session_token).await;
let ro_session_token = login_and_get_session_token(port, &ro_token).await;
let _ro_web_client_id = create_client_session(port, &ro_session_token).await;
let terminal_ws_url = format!(
"ws://127.0.0.1:{}/ws/terminal?web_client_id={}",
port, rw_web_client_id
);
let ws_result = timeout(
Duration::from_secs(5),
connect_async_with_cookie(&terminal_ws_url, &ro_session_token),
)
.await;
match ws_result {
Ok(Ok((ws, _))) => {
let (_sink, mut stream) = ws.split();
let msg = timeout(Duration::from_secs(2), stream.next()).await;
match msg {
Ok(Some(Ok(Message::Close(_)))) | Ok(Some(Err(_))) | Ok(None) | Err(_) => {
},
Ok(Some(Ok(_))) => {
},
}
},
Ok(Err(_)) => {
},
Err(_) => {
},
}
server_handle.abort();
let _ = delete_db();
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[tokio::test]
#[serial]
async fn test_base_url_with_special_characters_is_escaped() {
let _ = delete_db();
let session_manager = Arc::new(MockSessionManager::new());
let client_os_api_factory = Arc::new(MockClientOsApiFactory::new());
let mut config = Config::default();
config.web_client.base_url = Some("\"><script>alert(1)</script>".to_string());
let options = Options::default();
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let port = addr.port();
let temp_config_path = std::env::temp_dir().join("test_config.kdl");
let server_handle = tokio::spawn(async move {
serve_web_client(
config,
options,
Some(temp_config_path),
listener,
None,
Some(session_manager),
Some(client_os_api_factory),
addr.ip(),
port,
)
.await;
});
wait_for_server(port, Duration::from_secs(5))
.await
.expect("Server failed to start");
let url = format!("http://127.0.0.1:{}/", port);
let mut response = timeout(
Duration::from_secs(5),
tokio::task::spawn_blocking(move || isahc::get(&url)),
)
.await
.expect("Request timed out")
.expect("Spawn blocking failed")
.expect("Request failed");
let body = response.text().expect("Failed to read response body");
assert!(
!body.contains("<script>alert(1)</script>"),
"Response body should NOT contain unescaped script tag"
);
assert!(
body.contains("<script>"),
"Response body should contain HTML-escaped script tag"
);
server_handle.abort();
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
#[derive(Debug, Clone)]
pub struct MockSessionManager {
pub mock_sessions: HashMap<String, bool>,
pub mock_layouts: HashMap<String, Layout>,
pub sessions_created: Arc<Mutex<HashSet<String>>>,
pub first_messages_sent: Arc<Mutex<Vec<(String, ClientToServerMsg)>>>,
pub all_sessions_exist: bool,
}
impl MockSessionManager {
pub fn new() -> Self {
Self {
mock_sessions: HashMap::new(),
mock_layouts: HashMap::new(),
sessions_created: Arc::new(Mutex::new(HashSet::new())),
first_messages_sent: Arc::new(Mutex::new(Vec::new())),
all_sessions_exist: false,
}
}
pub fn with_all_sessions_existing() -> Self {
Self {
mock_sessions: HashMap::new(),
mock_layouts: HashMap::new(),
sessions_created: Arc::new(Mutex::new(HashSet::new())),
first_messages_sent: Arc::new(Mutex::new(Vec::new())),
all_sessions_exist: true,
}
}
#[allow(dead_code)]
pub fn get_first_message_for_session(&self, session_name: &str) -> Option<ClientToServerMsg> {
self.first_messages_sent
.lock()
.unwrap()
.iter()
.find(|(name, _)| name == session_name)
.map(|(_, msg)| msg.clone())
}
pub fn was_session_created(&self, session_name: &str) -> bool {
self.sessions_created.lock().unwrap().contains(session_name)
}
}
#[cfg(test)]
impl SessionManager for MockSessionManager {
fn session_exists(&self, session_name: &str) -> Result<bool, Box<dyn std::error::Error>> {
if self.all_sessions_exist {
Ok(true)
} else {
Ok(self
.mock_sessions
.get(session_name)
.copied()
.unwrap_or(false))
}
}
fn get_resurrection_layout(&self, session_name: &str) -> Option<Layout> {
self.mock_layouts.get(session_name).cloned()
}
fn spawn_session_if_needed(
&self,
session_name: &str,
_os_input: Box<dyn ClientOsApi>,
session_exists: bool,
_zellij_ipc_pipe: &PathBuf,
first_message: ClientToServerMsg,
) {
self.first_messages_sent
.lock()
.unwrap()
.push((session_name.to_owned(), first_message.clone()));
if !session_exists {
self.sessions_created
.lock()
.unwrap()
.insert(session_name.to_owned());
}
}
}
#[derive(Debug, Clone)]
struct MockClientOsApiFactory {
mock_apis: Arc<Mutex<HashMap<String, Arc<MockClientOsApi>>>>,
}
impl MockClientOsApiFactory {
fn new() -> Self {
Self {
mock_apis: Arc::new(Mutex::new(HashMap::new())),
}
}
}
impl ClientOsApiFactory for MockClientOsApiFactory {
fn create_client_os_api(&self) -> Result<Box<dyn ClientOsApi>, Box<dyn std::error::Error>> {
let mock_api = Arc::new(MockClientOsApi::new());
let client_id = uuid::Uuid::new_v4().to_string();
self.mock_apis
.lock()
.unwrap()
.insert(client_id, mock_api.clone());
Ok(Box::new((*mock_api).clone()))
}
}
#[derive(Debug, Clone)]
struct MockClientOsApi {
terminal_size: Size,
messages_to_server: Arc<Mutex<Vec<ClientToServerMsg>>>,
messages_from_server: Arc<Mutex<VecDeque<(ServerToClientMsg, ErrorContext)>>>,
}
impl MockClientOsApi {
fn new() -> Self {
Self {
terminal_size: Size { rows: 24, cols: 80 },
messages_to_server: Arc::new(Mutex::new(Vec::new())),
messages_from_server: Arc::new(Mutex::new(VecDeque::new())),
}
}
fn get_sent_messages(&self) -> Vec<ClientToServerMsg> {
self.messages_to_server.lock().unwrap().clone()
}
fn queue_server_message(&self, msg: ServerToClientMsg) {
self.messages_from_server
.lock()
.unwrap()
.push_back((msg, ErrorContext::default()));
}
}
impl ClientOsApi for MockClientOsApi {
fn get_terminal_size(&self) -> Size {
self.terminal_size
}
fn set_raw_mode(&mut self) {}
fn unset_raw_mode(&self) -> Result<(), std::io::Error> {
Ok(())
}
fn get_stdout_writer(&self) -> Box<dyn std::io::Write> {
Box::new(std::io::sink())
}
fn get_stdin_reader(&self) -> Box<dyn std::io::BufRead> {
Box::new(std::io::Cursor::new(Vec::new()))
}
fn update_session_name(&mut self, _new_session_name: String) {}
fn read_from_stdin(&mut self) -> Result<Vec<u8>, &'static str> {
Ok(Vec::new())
}
fn box_clone(&self) -> Box<dyn ClientOsApi> {
Box::new(self.clone())
}
fn send_to_server(&self, msg: ClientToServerMsg) {
self.messages_to_server.lock().unwrap().push(msg);
}
fn recv_from_server(&self) -> Option<(ServerToClientMsg, ErrorContext)> {
let msg = self.messages_from_server.lock().unwrap().pop_front();
if msg.is_none() {
std::thread::sleep(std::time::Duration::from_millis(1));
}
msg
}
fn handle_signals(
&self,
_sigwinch_cb: Box<dyn Fn()>,
_quit_cb: Box<dyn Fn()>,
_resize_receiver: Option<std::sync::mpsc::Receiver<()>>,
) {
}
fn connect_to_server(&self, _path: &std::path::Path) {}
fn load_palette(&self) -> Palette {
Palette::default()
}
fn enable_mouse(&self) -> anyhow::Result<()> {
Ok(())
}
fn disable_mouse(&self) -> anyhow::Result<()> {
Ok(())
}
}