use std::ffi::OsString;
use anyhow::{Context, Result, bail};
pub const INTERNAL_ATTACH_FLAG: &str = "--__internal-attach";
const SOCKET_FLAG: &str = "--socket";
const SIZE_FLAG: &str = "--size";
const FALLBACK_COLS: u16 = 200;
const FALLBACK_ROWS: u16 = 50;
#[derive(Debug, Clone, PartialEq, Eq)]
struct AttachArgs {
session_name: String,
socket_path: std::path::PathBuf,
cols: u16,
rows: u16,
}
#[must_use]
pub fn intercept_from_env() -> Option<Result<()>> {
let mut args = std::env::args_os();
let _argv0 = args.next();
match args.next() {
Some(first) if first == INTERNAL_ATTACH_FLAG => Some(run_internal_attach(args)),
_ => None,
}
}
fn run_internal_attach<I>(args: I) -> Result<()>
where
I: IntoIterator<Item = OsString>,
{
let parsed = parse_attach_args(args)?;
drive_attach(&parsed)
}
fn parse_attach_args<I>(args: I) -> Result<AttachArgs>
where
I: IntoIterator<Item = OsString>,
{
let mut session_name: Option<String> = None;
let mut socket_path: Option<std::path::PathBuf> = None;
let mut cols = FALLBACK_COLS;
let mut rows = FALLBACK_ROWS;
let mut iter = args.into_iter();
while let Some(arg) = iter.next() {
if arg == SOCKET_FLAG {
let value = iter
.next()
.context("--socket requires an absolute daemon socket path")?;
socket_path = Some(std::path::PathBuf::from(value));
} else if arg == SIZE_FLAG {
if let Some(value) = iter.next()
&& let Some((c, r)) = parse_size(&value.to_string_lossy())
{
cols = c;
rows = r;
}
} else if !arg.as_encoded_bytes().starts_with(b"--") && session_name.is_none() {
session_name = Some(arg.to_string_lossy().into_owned());
}
}
let session_name =
session_name.context("the visual attach requires a session name argument")?;
let socket_path = socket_path.context("the visual attach requires a --socket path")?;
crate::shells::daemon::validate_socket_path(&socket_path)?;
rmux_sdk::SessionName::new(session_name.clone())
.map_err(|e| anyhow::anyhow!("invalid rmux session name {session_name:?}: {e}"))?;
Ok(AttachArgs {
session_name,
socket_path,
cols,
rows,
})
}
fn parse_size(raw: &str) -> Option<(u16, u16)> {
let (cols, rows) = raw.split_once('x')?;
let cols: u16 = cols.parse().ok()?;
let rows: u16 = rows.parse().ok()?;
Some((cols, rows))
}
#[cfg(any(unix, windows))]
fn drive_attach(parsed: &AttachArgs) -> Result<()> {
let _ = (parsed.cols, parsed.rows);
let name = rmux_sdk::SessionName::new(parsed.session_name.clone())
.map_err(|e| anyhow::anyhow!("invalid rmux session name: {e}"))?;
let connection = rmux_client::connect(&parsed.socket_path).with_context(|| {
format!(
"connect to embedded rmux daemon at {:?}",
parsed.socket_path
)
})?;
match connection
.begin_attach(name)
.context("begin attach to rmux session")?
{
rmux_client::AttachTransition::Upgraded(upgrade) => {
let (stream, initial) = upgrade.into_parts();
rmux_client::attach_terminal_with_initial_bytes(stream, initial)
.context("attach terminal to rmux session")?;
Ok(())
}
rmux_client::AttachTransition::Rejected(_) => {
bail!(
"attach rejected by the daemon: no such session {:?} (it may have exited)",
parsed.session_name
)
}
}
}
#[cfg(not(any(unix, windows)))]
fn drive_attach(_parsed: &AttachArgs) -> Result<()> {
bail!("the embedded visual attach is only supported on unix and Windows")
}
#[cfg(test)]
mod tests {
use super::*;
fn args(items: &[&str]) -> Vec<OsString> {
items.iter().map(OsString::from).collect()
}
fn valid_socket() -> &'static str {
if cfg!(windows) {
r"\\.\pipe\basemind-shells-test"
} else {
"/tmp/basemind/shells/rmux.sock"
}
}
#[test]
fn parses_valid_attach_args_with_size() {
let parsed = parse_attach_args(args(&[
"bmsh-1-2",
"--socket",
valid_socket(),
"--size",
"120x40",
]))
.expect("valid args parse");
assert_eq!(parsed.session_name, "bmsh-1-2");
assert_eq!(parsed.socket_path, std::path::PathBuf::from(valid_socket()));
assert_eq!(parsed.cols, 120);
assert_eq!(parsed.rows, 40);
}
#[test]
fn falls_back_to_default_size_when_size_is_malformed() {
let parsed = parse_attach_args(args(&[
"bmsh-1-2",
"--socket",
valid_socket(),
"--size",
"not-a-size",
]))
.expect("malformed size must not fail the parse");
assert_eq!(parsed.cols, FALLBACK_COLS);
assert_eq!(parsed.rows, FALLBACK_ROWS);
}
#[test]
fn falls_back_to_default_size_when_size_flag_is_absent() {
let parsed = parse_attach_args(args(&["bmsh-1-2", "--socket", valid_socket()]))
.expect("missing size is fine");
assert_eq!(parsed.cols, FALLBACK_COLS);
assert_eq!(parsed.rows, FALLBACK_ROWS);
}
#[test]
fn rejects_missing_socket() {
let err = parse_attach_args(args(&["bmsh-1-2", "--size", "80x24"]))
.expect_err("missing --socket must be rejected");
assert!(err.to_string().contains("--socket"), "{err}");
}
#[test]
fn rejects_missing_session_name() {
let err = parse_attach_args(args(&["--socket", "/tmp/rmux.sock"]))
.expect_err("missing session name must be rejected");
assert!(err.to_string().contains("session name"), "{err}");
}
#[cfg(unix)]
#[test]
fn rejects_relative_socket_path() {
let err = parse_attach_args(args(&["bmsh-1-2", "--socket", "relative/evil.sock"]))
.expect_err("relative socket must be rejected");
assert!(err.to_string().contains("must be absolute"), "{err}");
}
#[cfg(unix)]
#[test]
fn rejects_socket_path_with_parent_traversal() {
let err = parse_attach_args(args(&["bmsh-1-2", "--socket", "/var/run/../../evil.sock"]))
.expect_err("`..` socket must be rejected");
assert!(err.to_string().contains("must not contain `..`"), "{err}");
}
#[test]
fn rejects_empty_session_name() {
let err = parse_attach_args(args(&["", "--socket", valid_socket()]))
.expect_err("empty session name must be rejected by SessionName::new");
assert!(err.to_string().contains("session name"), "{err}");
}
#[test]
fn parse_size_handles_valid_and_invalid() {
assert_eq!(parse_size("200x50"), Some((200, 50)));
assert_eq!(parse_size("1x1"), Some((1, 1)));
assert_eq!(parse_size("200"), None);
assert_eq!(parse_size("x50"), None);
assert_eq!(parse_size("200x"), None);
assert_eq!(parse_size("axb"), None);
assert_eq!(parse_size("999999x10"), None); }
}