use std::ffi::{OsStr, OsString};
use std::os::unix::ffi::{OsStrExt, OsStringExt};
use std::path::{Path, PathBuf};
use std::time::Duration;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as B64;
use crate::frame::{KIND_CONTROL, KIND_HELLO, KIND_SCREEN};
use crate::task::Lifecycle;
pub const PROTOCOL_VERSION: u32 = 3;
#[derive(Debug, Clone, PartialEq)]
pub struct LaunchContext {
pub env: Vec<(OsString, OsString)>,
pub cwd: PathBuf,
}
impl LaunchContext {
pub fn here() -> LaunchContext {
LaunchContext {
env: std::env::vars_os().collect(),
cwd: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum Command {
Spawn { command: String, cwd: PathBuf },
Kill { id: u64 },
Remove { id: u64 },
Restart { id: u64 },
Tag { id: u64, on: bool },
Resize { rows: u16, cols: u16 },
Watch { id: Option<u64> },
Input { id: u64, bytes: Vec<u8> },
Paste { id: u64, bytes: Vec<u8> },
Mouse {
id: u64,
kind: MouseKind,
col: u16,
row: u16,
},
Scrollback { id: u64, action: ScrollAction },
SaveSession { name: String },
LoadSession { name: String },
Shutdown,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScrollAction {
Up(u16),
Down(u16),
Top,
Live,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MouseBtn {
Left = 0,
Middle = 1,
Right = 2,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MouseKind {
WheelUp,
WheelDown,
Press(MouseBtn),
Drag(MouseBtn),
Release(MouseBtn),
}
#[derive(Debug, Clone, PartialEq)]
pub enum Event {
HelloOk,
Tasks(Vec<TaskView>),
Screen(ScreenView),
Status(String),
}
#[derive(Debug, Clone, PartialEq)]
pub struct TaskView {
pub id: u64,
pub command: String,
pub cwd: PathBuf,
pub tagged: bool,
pub lifecycle: Lifecycle,
pub preview: String,
pub started_ago: Duration,
}
#[derive(Debug, Clone, PartialEq)]
pub struct ScreenView {
pub id: u64,
pub lines: Vec<String>,
pub formatted: Vec<u8>,
pub cursor: (u16, u16),
pub hide_cursor: bool,
pub wants_mouse: bool,
pub alt_screen: bool,
pub scrollback: usize,
}
fn ps(p: &Path) -> String {
p.to_string_lossy().into_owned()
}
fn os_b64(s: &OsStr) -> String {
B64.encode(s.as_bytes())
}
fn os_from_b64(v: &jzon::JsonValue) -> Option<OsString> {
Some(OsString::from_vec(B64.decode(v.as_str()?).ok()?))
}
fn lifecycle_str(l: Lifecycle) -> &'static str {
match l {
Lifecycle::Active => "active",
Lifecycle::Idle => "idle",
Lifecycle::Ok => "ok",
Lifecycle::Failed => "failed",
}
}
fn lifecycle_from(s: &str) -> Option<Lifecycle> {
match s {
"active" => Some(Lifecycle::Active),
"idle" => Some(Lifecycle::Idle),
"ok" => Some(Lifecycle::Ok),
"failed" => Some(Lifecycle::Failed),
_ => None,
}
}
pub fn encode_hello(ctx: &LaunchContext) -> (u8, Vec<u8>) {
let mut o = jzon::JsonValue::new_object();
let _ = o.insert("v", PROTOCOL_VERSION);
let _ = o.insert("cwd", ps(&ctx.cwd));
let mut pairs = jzon::JsonValue::new_array();
for (k, v) in &ctx.env {
let mut pair = jzon::JsonValue::new_array();
let _ = pair.push(os_b64(k));
let _ = pair.push(os_b64(v));
let _ = pairs.push(pair);
}
let _ = o.insert("env", pairs);
(KIND_HELLO, o.dump().into_bytes())
}
pub fn decode_hello(kind: u8, payload: &[u8]) -> Option<(u32, LaunchContext)> {
if kind != KIND_HELLO {
return None;
}
let v = jzon::parse(std::str::from_utf8(payload).ok()?).ok()?;
let mut env = Vec::new();
for pair in v["env"].members() {
env.push((os_from_b64(&pair[0])?, os_from_b64(&pair[1])?));
}
Some((
v["v"].as_u32()?,
LaunchContext {
env,
cwd: PathBuf::from(v["cwd"].as_str()?),
},
))
}
pub fn encode_command(cmd: &Command) -> (u8, Vec<u8>) {
let mut o = jzon::JsonValue::new_object();
match cmd {
Command::Spawn { command, cwd } => {
let _ = o.insert("t", "spawn");
let _ = o.insert("command", command.as_str());
let _ = o.insert("cwd", ps(cwd));
}
Command::Kill { id } => {
let _ = o.insert("t", "kill");
let _ = o.insert("id", *id);
}
Command::Remove { id } => {
let _ = o.insert("t", "remove");
let _ = o.insert("id", *id);
}
Command::Restart { id } => {
let _ = o.insert("t", "restart");
let _ = o.insert("id", *id);
}
Command::Tag { id, on } => {
let _ = o.insert("t", "tag");
let _ = o.insert("id", *id);
let _ = o.insert("on", *on);
}
Command::Resize { rows, cols } => {
let _ = o.insert("t", "resize");
let _ = o.insert("rows", *rows as u64);
let _ = o.insert("cols", *cols as u64);
}
Command::Watch { id } => {
let _ = o.insert("t", "watch");
match id {
Some(n) => {
let _ = o.insert("id", *n);
}
None => {
let _ = o.insert("id", jzon::JsonValue::Null);
}
}
}
Command::Input { id, bytes } => {
let _ = o.insert("t", "input");
let _ = o.insert("id", *id);
let mut arr = jzon::JsonValue::new_array();
for b in bytes {
let _ = arr.push(*b as u64);
}
let _ = o.insert("bytes", arr);
}
Command::Paste { id, bytes } => {
let _ = o.insert("t", "paste");
let _ = o.insert("id", *id);
let mut arr = jzon::JsonValue::new_array();
for b in bytes {
let _ = arr.push(*b as u64);
}
let _ = o.insert("bytes", arr);
}
Command::Mouse { id, kind, col, row } => {
let _ = o.insert("t", "mouse");
let _ = o.insert("id", *id);
let (k, btn) = match kind {
MouseKind::WheelUp => ("wu", None),
MouseKind::WheelDown => ("wd", None),
MouseKind::Press(b) => ("p", Some(*b)),
MouseKind::Drag(b) => ("d", Some(*b)),
MouseKind::Release(b) => ("r", Some(*b)),
};
let _ = o.insert("k", k);
if let Some(b) = btn {
let _ = o.insert("b", b as u64);
}
let _ = o.insert("col", *col as u64);
let _ = o.insert("row", *row as u64);
}
Command::Scrollback { id, action } => {
let _ = o.insert("t", "sb");
let _ = o.insert("id", *id);
let (a, n) = match action {
ScrollAction::Up(n) => ("u", Some(*n)),
ScrollAction::Down(n) => ("d", Some(*n)),
ScrollAction::Top => ("t", None),
ScrollAction::Live => ("l", None),
};
let _ = o.insert("a", a);
if let Some(n) = n {
let _ = o.insert("n", n as u64);
}
}
Command::SaveSession { name } => {
let _ = o.insert("t", "save");
let _ = o.insert("name", name.as_str());
}
Command::LoadSession { name } => {
let _ = o.insert("t", "load");
let _ = o.insert("name", name.as_str());
}
Command::Shutdown => {
let _ = o.insert("t", "shutdown");
}
}
(KIND_CONTROL, o.dump().into_bytes())
}
pub fn decode_command(kind: u8, payload: &[u8]) -> Option<Command> {
if kind != KIND_CONTROL {
return None;
}
let v = jzon::parse(std::str::from_utf8(payload).ok()?).ok()?;
let cmd = match v["t"].as_str()? {
"spawn" => Command::Spawn {
command: v["command"].as_str()?.to_string(),
cwd: PathBuf::from(v["cwd"].as_str()?),
},
"kill" => Command::Kill {
id: v["id"].as_u64()?,
},
"remove" => Command::Remove {
id: v["id"].as_u64()?,
},
"restart" => Command::Restart {
id: v["id"].as_u64()?,
},
"tag" => Command::Tag {
id: v["id"].as_u64()?,
on: v["on"].as_bool()?,
},
"resize" => Command::Resize {
rows: v["rows"].as_u64()? as u16,
cols: v["cols"].as_u64()? as u16,
},
"watch" => Command::Watch {
id: if v["id"].is_null() {
None
} else {
Some(v["id"].as_u64()?)
},
},
"input" => Command::Input {
id: v["id"].as_u64()?,
bytes: v["bytes"]
.members()
.filter_map(|m| m.as_u64().map(|n| n as u8))
.collect(),
},
"paste" => Command::Paste {
id: v["id"].as_u64()?,
bytes: v["bytes"]
.members()
.filter_map(|m| m.as_u64().map(|n| n as u8))
.collect(),
},
"mouse" => {
let btn = || -> Option<MouseBtn> {
match v["b"].as_u64()? {
0 => Some(MouseBtn::Left),
1 => Some(MouseBtn::Middle),
2 => Some(MouseBtn::Right),
_ => None,
}
};
Command::Mouse {
id: v["id"].as_u64()?,
kind: match v["k"].as_str()? {
"wu" => MouseKind::WheelUp,
"wd" => MouseKind::WheelDown,
"p" => MouseKind::Press(btn()?),
"d" => MouseKind::Drag(btn()?),
"r" => MouseKind::Release(btn()?),
_ => return None,
},
col: v["col"].as_u64()? as u16,
row: v["row"].as_u64()? as u16,
}
}
"sb" => Command::Scrollback {
id: v["id"].as_u64()?,
action: match v["a"].as_str()? {
"u" => ScrollAction::Up(v["n"].as_u64()? as u16),
"d" => ScrollAction::Down(v["n"].as_u64()? as u16),
"t" => ScrollAction::Top,
"l" => ScrollAction::Live,
_ => return None,
},
},
"save" => Command::SaveSession {
name: v["name"].as_str()?.to_string(),
},
"load" => Command::LoadSession {
name: v["name"].as_str()?.to_string(),
},
"shutdown" => Command::Shutdown,
_ => return None,
};
Some(cmd)
}
pub fn encode_event(ev: &Event) -> (u8, Vec<u8>) {
match ev {
Event::HelloOk => {
let mut o = jzon::JsonValue::new_object();
let _ = o.insert("t", "hello_ok");
(KIND_CONTROL, o.dump().into_bytes())
}
Event::Tasks(views) => {
let mut arr = jzon::JsonValue::new_array();
for tv in views {
let mut o = jzon::JsonValue::new_object();
let _ = o.insert("id", tv.id);
let _ = o.insert("command", tv.command.as_str());
let _ = o.insert("cwd", ps(&tv.cwd));
let _ = o.insert("tagged", tv.tagged);
let _ = o.insert("life", lifecycle_str(tv.lifecycle));
let _ = o.insert("preview", tv.preview.as_str());
let _ = o.insert("started_ms", tv.started_ago.as_millis() as u64);
let _ = arr.push(o);
}
let mut root = jzon::JsonValue::new_object();
let _ = root.insert("t", "tasks");
let _ = root.insert("tasks", arr);
(KIND_CONTROL, root.dump().into_bytes())
}
Event::Status(msg) => {
let mut o = jzon::JsonValue::new_object();
let _ = o.insert("t", "status");
let _ = o.insert("msg", msg.as_str());
(KIND_CONTROL, o.dump().into_bytes())
}
Event::Screen(sv) => {
let mut header = jzon::JsonValue::new_object();
let _ = header.insert("id", sv.id);
let mut cur = jzon::JsonValue::new_array();
let _ = cur.push(sv.cursor.0 as u64);
let _ = cur.push(sv.cursor.1 as u64);
let _ = header.insert("cursor", cur);
let _ = header.insert("hide", sv.hide_cursor);
let _ = header.insert("mouse", sv.wants_mouse);
let _ = header.insert("alt", sv.alt_screen);
let _ = header.insert("sb", sv.scrollback as u64);
let mut lines = jzon::JsonValue::new_array();
for l in &sv.lines {
let _ = lines.push(l.as_str());
}
let _ = header.insert("lines", lines);
let hbytes = header.dump().into_bytes();
let mut payload = Vec::with_capacity(4 + hbytes.len() + sv.formatted.len());
payload.extend_from_slice(&(hbytes.len() as u32).to_be_bytes());
payload.extend_from_slice(&hbytes);
payload.extend_from_slice(&sv.formatted);
(KIND_SCREEN, payload)
}
}
}
pub fn decode_event(kind: u8, payload: &[u8]) -> Option<Event> {
match kind {
KIND_CONTROL => {
let v = jzon::parse(std::str::from_utf8(payload).ok()?).ok()?;
match v["t"].as_str()? {
"hello_ok" => Some(Event::HelloOk),
"tasks" => {
let mut views = Vec::new();
for tv in v["tasks"].members() {
views.push(TaskView {
id: tv["id"].as_u64()?,
command: tv["command"].as_str()?.to_string(),
cwd: PathBuf::from(tv["cwd"].as_str()?),
tagged: tv["tagged"].as_bool()?,
lifecycle: lifecycle_from(tv["life"].as_str()?)?,
preview: tv["preview"].as_str()?.to_string(),
started_ago: Duration::from_millis(tv["started_ms"].as_u64()?),
});
}
Some(Event::Tasks(views))
}
"status" => Some(Event::Status(v["msg"].as_str()?.to_string())),
_ => None,
}
}
KIND_SCREEN => {
let hlen = u32::from_be_bytes(payload.get(0..4)?.try_into().ok()?) as usize;
let header_bytes = payload.get(4..4 + hlen)?;
let formatted = payload.get(4 + hlen..)?.to_vec();
let h = jzon::parse(std::str::from_utf8(header_bytes).ok()?).ok()?;
let cursor = (
h["cursor"][0].as_u64()? as u16,
h["cursor"][1].as_u64()? as u16,
);
let lines = h["lines"]
.members()
.filter_map(|m| m.as_str().map(str::to_string))
.collect();
Some(Event::Screen(ScreenView {
id: h["id"].as_u64()?,
lines,
formatted,
cursor,
hide_cursor: h["hide"].as_bool()?,
wants_mouse: h["mouse"].as_bool()?,
alt_screen: h["alt"].as_bool()?,
scrollback: h["sb"].as_u64()? as usize,
}))
}
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn command_round_trips() {
let cases = [
Command::Spawn {
command: "echo hi".into(),
cwd: PathBuf::from("/tmp"),
},
Command::Kill { id: 7 },
Command::Remove { id: 3 },
Command::Restart { id: 4 },
Command::Tag { id: 2, on: true },
Command::Resize {
rows: 30,
cols: 100,
},
Command::Watch { id: Some(5) },
Command::Watch { id: None },
Command::Input {
id: 1,
bytes: vec![0, 27, 91, 255],
},
Command::Paste {
id: 6,
bytes: b"line1\nline2\x1b[201~\xff".to_vec(),
},
Command::Mouse {
id: 8,
kind: MouseKind::WheelDown,
col: 79,
row: 23,
},
Command::Mouse {
id: 8,
kind: MouseKind::Press(MouseBtn::Left),
col: 0,
row: 0,
},
Command::Mouse {
id: 8,
kind: MouseKind::Drag(MouseBtn::Middle),
col: 10,
row: 5,
},
Command::Mouse {
id: 8,
kind: MouseKind::Release(MouseBtn::Right),
col: 10,
row: 5,
},
Command::Scrollback {
id: 3,
action: ScrollAction::Up(23),
},
Command::Scrollback {
id: 3,
action: ScrollAction::Down(1),
},
Command::Scrollback {
id: 3,
action: ScrollAction::Top,
},
Command::Scrollback {
id: 3,
action: ScrollAction::Live,
},
Command::SaveSession {
name: "work".into(),
},
Command::LoadSession {
name: "home".into(),
},
Command::Shutdown,
];
for c in cases {
let (k, p) = encode_command(&c);
assert_eq!(decode_command(k, &p).as_ref(), Some(&c), "round-trip {c:?}");
}
}
#[test]
fn hello_ok_round_trips() {
let ack = Event::HelloOk;
let (k, p) = encode_event(&ack);
assert_eq!(k, KIND_CONTROL);
assert_eq!(decode_event(k, &p), Some(ack));
}
#[test]
fn hello_round_trips() {
let ctx = LaunchContext {
env: vec![
("PATH".into(), "/usr/bin:/bin".into()),
(
OsString::from_vec(b"BAD\xff\xfe".to_vec()),
OsString::from_vec(b"v\xff".to_vec()),
),
],
cwd: PathBuf::from("/home/x"),
};
let (k, p) = encode_hello(&ctx);
assert_eq!(k, KIND_HELLO);
assert_eq!(decode_hello(k, &p), Some((PROTOCOL_VERSION, ctx)));
let empty = LaunchContext {
env: Vec::new(),
cwd: PathBuf::from("/"),
};
let (k, p) = encode_hello(&empty);
assert_eq!(decode_hello(k, &p), Some((PROTOCOL_VERSION, empty)));
}
#[test]
fn hello_requires_its_own_frame_kind() {
let (_, p) = encode_hello(&LaunchContext {
env: Vec::new(),
cwd: PathBuf::from("/"),
});
assert_eq!(decode_hello(KIND_CONTROL, &p), None);
assert_eq!(decode_command(KIND_CONTROL, &p), None);
}
#[test]
fn hello_with_malformed_env_is_rejected() {
for env in [
r#"[["P@TH","L2Jpbg=="]]"#, r#"[["QUFBQUE","L2Jpbg=="]]"#, r#"[["UEFUSA==","AAAA="]]"#, r#"[[[80],[65]]]"#, r#"["PATH=/bin"]"#, ] {
let json = format!(r#"{{"v":3,"cwd":"/","env":{env}}}"#);
assert_eq!(
decode_hello(KIND_HELLO, json.as_bytes()),
None,
"should reject env {env}"
);
}
}
#[test]
fn tasks_and_status_round_trip() {
let tasks = Event::Tasks(vec![TaskView {
id: 1,
command: "vim".into(),
cwd: PathBuf::from("/home/x"),
tagged: true,
lifecycle: Lifecycle::Idle,
preview: "~ line".into(),
started_ago: Duration::from_millis(4200),
}]);
let (k, p) = encode_event(&tasks);
assert_eq!(k, KIND_CONTROL);
assert_eq!(decode_event(k, &p), Some(tasks));
let status = Event::Status("saved 'x'".into());
let (k, p) = encode_event(&status);
assert_eq!(decode_event(k, &p), Some(status));
}
#[test]
fn screen_round_trips_raw_bytes() {
let screen = Event::Screen(ScreenView {
id: 9,
lines: vec!["row0".into(), "row1".into()],
formatted: vec![0x1b, b'[', b'm', 0, 255, b'x'],
cursor: (3, 12),
hide_cursor: false,
wants_mouse: true,
alt_screen: false,
scrollback: 42,
});
let (k, p) = encode_event(&screen);
assert_eq!(k, KIND_SCREEN);
assert_eq!(decode_event(k, &p), Some(screen));
}
}