use super::*;
mod control;
mod rfb;
#[derive(Debug, Args)]
#[command(
args_conflicts_with_subcommands = true,
subcommand_precedence_over_arg = true
)]
pub(crate) struct DesktopCommand {
#[command(subcommand)]
action: Option<control::Action>,
#[arg(default_value = "current")]
id: String,
#[arg(long)]
endpoint: Option<String>,
#[arg(long)]
session: Option<String>,
#[arg(long)]
no_viewer: bool,
#[arg(long, default_value = "vncviewer")]
viewer: String,
}
#[derive(serde::Deserialize)]
struct Session {
port: u16,
session: String,
}
#[derive(Debug)]
pub(crate) struct MissingDesktop {
pub(crate) box_id: String,
}
impl std::fmt::Display for MissingDesktop {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(formatter, "No desktop installed on {}", self.box_id)
}
}
impl std::error::Error for MissingDesktop {}
fn check_launcher(result: &ExecResult, box_id: &str) -> Result<()> {
if result.exited && result.code == 1 {
return Err(MissingDesktop {
box_id: box_id.to_owned(),
}
.into());
}
anyhow::ensure!(
result.exited && result.code == 0,
"Could not check the desktop installation: {}",
safe_terminal_text(&String::from_utf8_lossy(&result.stderr))
);
Ok(())
}
fn parse_session(bytes: &[u8]) -> Result<Session> {
let session: Session =
serde_json::from_slice(bytes).context("read desktop launcher response")?;
anyhow::ensure!(
session.port >= 1024,
"desktop launcher returned an invalid VNC port"
);
Ok(session)
}
async fn start_session(
agent: &mut AgentClient,
id: &str,
selected_session: Option<String>,
) -> Result<Session> {
agent.info().await?;
let installed = agent
.exec(
vec![
"/bin/sh".to_owned(),
"-c".to_owned(),
"test -e /usr/local/bin/pbox-desktop".to_owned(),
],
"/",
Vec::<(String, String)>::new(),
"root",
)
.await
.context("check desktop installation")?;
check_launcher(&installed, id)?;
let mut argv = vec!["/usr/local/bin/pbox-desktop".to_owned()];
if let Some(session) = selected_session {
argv.extend(["--session".to_owned(), session]);
}
let result = agent
.exec(argv, "/", Vec::<(String, String)>::new(), "root")
.await
.context("start the installed desktop")?;
anyhow::ensure!(
result.exited && result.code == 0,
"desktop startup failed: {}",
safe_terminal_text(&String::from_utf8_lossy(&result.stderr))
);
parse_session(&result.stdout)
}
pub(crate) fn run(store: &ConfigStore, command: DesktopCommand, json: bool) -> Result<()> {
if let Some(action) = command.action {
return control::run(store, action, json);
}
let config = load_config(store)?;
let (id, endpoint) = resolve_agent_endpoint(&config, &command.id, command.endpoint.as_deref())?;
let materials = agent_materials(&config, &id)?;
tokio::runtime::Builder::new_current_thread().enable_all().build()?.block_on(async {
let mut agent = relay::connect_agent(&config, &endpoint, &id, &materials.ca.certificate_pem, &materials.client).await?;
let session = start_session(&mut agent, &id, command.session).await?;
let listener = TcpListener::bind(("127.0.0.1", 0)).await?;
let local = listener.local_addr()?;
let mut viewer = if command.no_viewer || json { None } else {
Some(tokio::process::Command::new(&command.viewer)
.arg(format!("127.0.0.1::{}", local.port()))
.stdin(std::process::Stdio::null()).stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null()).kill_on_drop(true).spawn()
.context("open VNC viewer; install TigerVNC viewer or use --no-viewer")?)
};
if json {
ui::json_text(&serde_json::to_string(&serde_json::json!({"id": id, "session": session.session, "local": local.to_string(), "remote": format!("127.0.0.1:{}", session.port)}))?);
} else {
ui::desktop_ready(&session.session, &local.to_string());
}
let mut connections = tokio::task::JoinSet::new();
loop {
tokio::select! {
result = tokio::signal::ctrl_c() => { result?; break; }
status = async { match &mut viewer { Some(child) => child.wait().await, None => std::future::pending().await } } => {
anyhow::ensure!(status?.success(), "VNC viewer exited unsuccessfully");
break;
}
Some(result) = connections.join_next(), if !connections.is_empty() => {
match result {
Ok(Ok(())) => {},
Ok(Err(error)) => ui::stderr().error(&safe_terminal_text(&format!("VNC connection: {error:#}"))),
Err(error) => return Err(error.into()),
}
}
accepted = listener.accept(), if connections.len() < MAX_FORWARD_CONNECTIONS => {
let (socket, _) = accepted?;
let mut client = agent.clone();
connections.spawn(async move { client.forward_tcp(socket, "127.0.0.1", session.port).await });
}
}
}
Ok(())
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn missing_launcher_is_distinct_from_a_failed_probe() {
let mut result = ExecResult {
exited: true,
code: 1,
..Default::default()
};
assert!(
check_launcher(&result, "pbx_12345678")
.unwrap_err()
.is::<MissingDesktop>()
);
result.code = 0;
assert!(check_launcher(&result, "pbx_12345678").is_ok());
result.code = 2;
assert!(
!check_launcher(&result, "pbx_12345678")
.unwrap_err()
.is::<MissingDesktop>()
);
result.code = 1;
result.exited = false;
assert!(
!check_launcher(&result, "pbx_12345678")
.unwrap_err()
.is::<MissingDesktop>()
);
}
#[test]
fn launcher_protocol_rejects_invalid_ports_and_non_json_output() {
assert!(parse_session(br#"{"port":5901,"session":"xfce"}"#).is_ok());
assert!(parse_session(br#"{"port":0,"session":"xfce"}"#).is_err());
assert!(parse_session(b"started desktop").is_err());
}
}