ZIm/crates/cli/src/main.rs

#![cfg_attr(
    any(target_os = "linux", target_os = "freebsd", target_os = "windows"),
    allow(dead_code)
)]

use anyhow::{Context, Result};
use clap::Parser;
use cli::{ipc::IpcOneShotServer, CliRequest, CliResponse, IpcHandshake};
use collections::HashMap;
use parking_lot::Mutex;
use std::{
    env, fs, io,
    path::{Path, PathBuf},
    process::ExitStatus,
    sync::Arc,
    thread::{self, JoinHandle},
};
use tempfile::NamedTempFile;
use util::paths::PathWithPosition;

#[cfg(any(target_os = "linux", target_os = "freebsd"))]
use std::io::IsTerminal;

struct Detect;

trait InstalledApp {
    fn zed_version_string(&self) -> String;
    fn launch(&self, ipc_url: String) -> anyhow::Result<()>;
    fn run_foreground(&self, ipc_url: String) -> io::Result<ExitStatus>;
}

#[derive(Parser, Debug)]
#[command(
    name = "zed",
    disable_version_flag = true,
    after_help = "To read from stdin, append '-' (e.g. 'ps axf | zed -')"
)]
struct Args {
    /// Wait for all of the given paths to be opened/closed before exiting.
    #[arg(short, long)]
    wait: bool,
    /// Add files to the currently open workspace
    #[arg(short, long, overrides_with = "new")]
    add: bool,
    /// Create a new workspace
    #[arg(short, long, overrides_with = "add")]
    new: bool,
    /// A sequence of space-separated paths that you want to open.
    ///
    /// Use `path:line:row` syntax to open a file at a specific location.
    /// Non-existing paths and directories will ignore `:line:row` suffix.
    paths_with_position: Vec<String>,
    /// Print Zed's version and the app path.
    #[arg(short, long)]
    version: bool,
    /// Run zed in the foreground (useful for debugging)
    #[arg(long)]
    foreground: bool,
    /// Custom path to Zed.app or the zed binary
    #[arg(long)]
    zed: Option<PathBuf>,
    /// Run zed in dev-server mode
    #[arg(long)]
    dev_server_token: Option<String>,
    /// Uninstall Zed from user system
    #[cfg(all(
        any(target_os = "linux", target_os = "macos"),
        not(feature = "no-bundled-uninstall")
    ))]
    #[arg(long)]
    uninstall: bool,
}

fn parse_path_with_position(argument_str: &str) -> anyhow::Result<String> {
    let canonicalized = match Path::new(argument_str).canonicalize() {
        Ok(existing_path) => PathWithPosition::from_path(existing_path),
        Err(_) => {
            let path = PathWithPosition::parse_str(argument_str);
            let curdir = env::current_dir().context("retrieving current directory")?;
            path.map_path(|path| match fs::canonicalize(&path) {
                Ok(path) => Ok(path),
                Err(e) => {
                    if let Some(mut parent) = path.parent() {
                        if parent == Path::new("") {
                            parent = &curdir
                        }
                        match fs::canonicalize(parent) {
                            Ok(parent) => Ok(parent.join(path.file_name().unwrap())),
                            Err(_) => Err(e),
                        }
                    } else {
                        Err(e)
                    }
                }
            })
        }
        .with_context(|| format!("parsing as path with position {argument_str}"))?,
    };
    Ok(canonicalized.to_string(|path| path.to_string_lossy().to_string()))
}

fn main() -> Result<()> {
    // Exit flatpak sandbox if needed
    #[cfg(any(target_os = "linux", target_os = "freebsd"))]
    {
        flatpak::try_restart_to_host();
        flatpak::ld_extra_libs();
    }

    // Intercept version designators
    #[cfg(target_os = "macos")]
    if let Some(channel) = std::env::args().nth(1).filter(|arg| arg.starts_with("--")) {
        // When the first argument is a name of a release channel, we're gonna spawn off a cli of that version, with trailing args passed along.
        use std::str::FromStr as _;

        if let Ok(channel) = release_channel::ReleaseChannel::from_str(&channel[2..]) {
            return mac_os::spawn_channel_cli(channel, std::env::args().skip(2).collect());
        }
    }
    let args = Args::parse();

    #[cfg(any(target_os = "linux", target_os = "freebsd"))]
    let args = flatpak::set_bin_if_no_escape(args);

    let app = Detect::detect(args.zed.as_deref()).context("Bundle detection")?;

    if args.version {
        println!("{}", app.zed_version_string());
        return Ok(());
    }

    #[cfg(all(
        any(target_os = "linux", target_os = "macos"),
        not(feature = "no-bundled-uninstall")
    ))]
    if args.uninstall {
        static UNINSTALL_SCRIPT: &[u8] = include_bytes!("../../../script/uninstall.sh");

        let tmp_dir = tempfile::tempdir()?;
        let script_path = tmp_dir.path().join("uninstall.sh");
        fs::write(&script_path, UNINSTALL_SCRIPT)?;

        use std::os::unix::fs::PermissionsExt as _;
        fs::set_permissions(&script_path, fs::Permissions::from_mode(0o755))?;

        let status = std::process::Command::new("sh")
            .arg(&script_path)
            .env("ZED_CHANNEL", &*release_channel::RELEASE_CHANNEL_NAME)
            .status()
            .context("Failed to execute uninstall script")?;

        std::process::exit(status.code().unwrap_or(1));
    }

    let (server, server_name) =
        IpcOneShotServer::<IpcHandshake>::new().context("Handshake before Zed spawn")?;
    let url = format!("zed-cli://{server_name}");

    let open_new_workspace = if args.new {
        Some(true)
    } else if args.add {
        Some(false)
    } else {
        None
    };

    let env = {
        #[cfg(any(target_os = "linux", target_os = "freebsd"))]
        {
            // On Linux, the desktop entry uses `cli` to spawn `zed`.
            // We need to handle env vars correctly since std::env::vars() may not contain
            // project-specific vars (e.g. those set by direnv).
            // By setting env to None here, the LSP will use worktree env vars instead,
            // which is what we want.
            if !std::io::stdout().is_terminal() {
                None
            } else {
                Some(std::env::vars().collect::<HashMap<_, _>>())
            }
        }

        #[cfg(not(any(target_os = "linux", target_os = "freebsd")))]
        Some(std::env::vars().collect::<HashMap<_, _>>())
    };

    let exit_status = Arc::new(Mutex::new(None));
    let mut paths = vec![];
    let mut urls = vec![];
    let mut stdin_tmp_file: Option<fs::File> = None;
    for path in args.paths_with_position.iter() {
        if path.starts_with("zed://")
            || path.starts_with("http://")
            || path.starts_with("https://")
            || path.starts_with("file://")
            || path.starts_with("ssh://")
        {
            urls.push(path.to_string());
        } else if path == "-" && args.paths_with_position.len() == 1 {
            let file = NamedTempFile::new()?;
            paths.push(file.path().to_string_lossy().to_string());
            let (file, _) = file.keep()?;
            stdin_tmp_file = Some(file);
        } else {
            paths.push(parse_path_with_position(path)?)
        }
    }

    if let Some(_) = args.dev_server_token {
        return Err(anyhow::anyhow!(
            "Dev servers were removed in v0.157.x please upgrade to SSH remoting: https://zed.dev/docs/remote-development"
        ))?;
    }

    let sender: JoinHandle<anyhow::Result<()>> = thread::spawn({
        let exit_status = exit_status.clone();
        move || {
            let (_, handshake) = server.accept().context("Handshake after Zed spawn")?;
            let (tx, rx) = (handshake.requests, handshake.responses);

            tx.send(CliRequest::Open {
                paths,
                urls,
                wait: args.wait,
                open_new_workspace,
                env,
            })?;

            while let Ok(response) = rx.recv() {
                match response {
                    CliResponse::Ping => {}
                    CliResponse::Stdout { message } => println!("{message}"),
                    CliResponse::Stderr { message } => eprintln!("{message}"),
                    CliResponse::Exit { status } => {
                        exit_status.lock().replace(status);
                        return Ok(());
                    }
                }
            }

            Ok(())
        }
    });

    let pipe_handle: JoinHandle<anyhow::Result<()>> = thread::spawn(move || {
        if let Some(mut tmp_file) = stdin_tmp_file {
            let mut stdin = std::io::stdin().lock();
            if io::IsTerminal::is_terminal(&stdin) {
                return Ok(());
            }
            let mut buffer = [0; 8 * 1024];
            loop {
                let bytes_read = io::Read::read(&mut stdin, &mut buffer)?;
                if bytes_read == 0 {
                    break;
                }
                io::Write::write(&mut tmp_file, &buffer[..bytes_read])?;
            }
            io::Write::flush(&mut tmp_file)?;
        }
        Ok(())
    });

    if args.foreground {
        app.run_foreground(url)?;
    } else {
        app.launch(url)?;
        sender.join().unwrap()?;
        pipe_handle.join().unwrap()?;
    }

    if let Some(exit_status) = exit_status.lock().take() {
        std::process::exit(exit_status);
    }
    Ok(())
}

#[cfg(any(target_os = "linux", target_os = "freebsd"))]
mod linux {
    use std::{
        env,
        ffi::OsString,
        io,
        os::unix::net::{SocketAddr, UnixDatagram},
        path::{Path, PathBuf},
        process::{self, ExitStatus},
        sync::LazyLock,
        thread,
        time::Duration,
    };

    use anyhow::anyhow;
    use cli::FORCE_CLI_MODE_ENV_VAR_NAME;
    use fork::Fork;

    use crate::{Detect, InstalledApp};

    static RELEASE_CHANNEL: LazyLock<String> =
        LazyLock::new(|| include_str!("../../zed/RELEASE_CHANNEL").trim().to_string());

    struct App(PathBuf);

    impl Detect {
        pub fn detect(path: Option<&Path>) -> anyhow::Result<impl InstalledApp> {
            let path = if let Some(path) = path {
                path.to_path_buf().canonicalize()?
            } else {
                let cli = env::current_exe()?;
                let dir = cli
                    .parent()
                    .ok_or_else(|| anyhow!("no parent path for cli"))?;

                // libexec is the standard, lib/zed is for Arch (and other non-libexec distros),
                // ./zed is for the target directory in development builds.
                let possible_locations =
                    ["../libexec/zed-editor", "../lib/zed/zed-editor", "./zed"];
                possible_locations
                    .iter()
                    .find_map(|p| dir.join(p).canonicalize().ok().filter(|path| path != &cli))
                    .ok_or_else(|| {
                        anyhow!("could not find any of: {}", possible_locations.join(", "))
                    })?
            };

            Ok(App(path))
        }
    }

    impl InstalledApp for App {
        fn zed_version_string(&self) -> String {
            format!(
                "Zed {}{} – {}",
                if *RELEASE_CHANNEL == "stable" {
                    "".to_string()
                } else {
                    format!(" {} ", *RELEASE_CHANNEL)
                },
                option_env!("RELEASE_VERSION").unwrap_or_default(),
                self.0.display(),
            )
        }

        fn launch(&self, ipc_url: String) -> anyhow::Result<()> {
            let sock_path = paths::support_dir().join(format!("zed-{}.sock", *RELEASE_CHANNEL));
            let sock = UnixDatagram::unbound()?;
            if sock.connect(&sock_path).is_err() {
                self.boot_background(ipc_url)?;
            } else {
                sock.send(ipc_url.as_bytes())?;
            }
            Ok(())
        }

        fn run_foreground(&self, ipc_url: String) -> io::Result<ExitStatus> {
            std::process::Command::new(self.0.clone())
                .arg(ipc_url)
                .status()
        }
    }

    impl App {
        fn boot_background(&self, ipc_url: String) -> anyhow::Result<()> {
            let path = &self.0;

            match fork::fork() {
                Ok(Fork::Parent(_)) => Ok(()),
                Ok(Fork::Child) => {
                    std::env::set_var(FORCE_CLI_MODE_ENV_VAR_NAME, "");
                    if let Err(_) = fork::setsid() {
                        eprintln!("failed to setsid: {}", std::io::Error::last_os_error());
                        process::exit(1);
                    }
                    if let Err(_) = fork::close_fd() {
                        eprintln!("failed to close_fd: {}", std::io::Error::last_os_error());
                    }
                    let error =
                        exec::execvp(path.clone(), &[path.as_os_str(), &OsString::from(ipc_url)]);
                    // if exec succeeded, we never get here.
                    eprintln!("failed to exec {:?}: {}", path, error);
                    process::exit(1)
                }
                Err(_) => Err(anyhow!(io::Error::last_os_error())),
            }
        }

        fn wait_for_socket(
            &self,
            sock_addr: &SocketAddr,
            sock: &mut UnixDatagram,
        ) -> Result<(), std::io::Error> {
            for _ in 0..100 {
                thread::sleep(Duration::from_millis(10));
                if sock.connect_addr(&sock_addr).is_ok() {
                    return Ok(());
                }
            }
            sock.connect_addr(&sock_addr)
        }
    }
}

#[cfg(any(target_os = "linux", target_os = "freebsd"))]
mod flatpak {
    use std::ffi::OsString;
    use std::path::PathBuf;
    use std::process::Command;
    use std::{env, process};

    const EXTRA_LIB_ENV_NAME: &'static str = "ZED_FLATPAK_LIB_PATH";
    const NO_ESCAPE_ENV_NAME: &'static str = "ZED_FLATPAK_NO_ESCAPE";

    /// Adds bundled libraries to LD_LIBRARY_PATH if running under flatpak
    pub fn ld_extra_libs() {
        let mut paths = if let Ok(paths) = env::var("LD_LIBRARY_PATH") {
            env::split_paths(&paths).collect()
        } else {
            Vec::new()
        };

        if let Ok(extra_path) = env::var(EXTRA_LIB_ENV_NAME) {
            paths.push(extra_path.into());
        }

        env::set_var("LD_LIBRARY_PATH", env::join_paths(paths).unwrap());
    }

    /// Restarts outside of the sandbox if currently running within it
    pub fn try_restart_to_host() {
        if let Some(flatpak_dir) = get_flatpak_dir() {
            let mut args = vec!["/usr/bin/flatpak-spawn".into(), "--host".into()];
            args.append(&mut get_xdg_env_args());
            args.push("--env=ZED_UPDATE_EXPLANATION=Please use flatpak to update zed".into());
            args.push(
                format!(
                    "--env={EXTRA_LIB_ENV_NAME}={}",
                    flatpak_dir.join("lib").to_str().unwrap()
                )
                .into(),
            );
            args.push(flatpak_dir.join("bin").join("zed").into());

            let mut is_app_location_set = false;
            for arg in &env::args_os().collect::<Vec<_>>()[1..] {
                args.push(arg.clone());
                is_app_location_set |= arg == "--zed";
            }

            if !is_app_location_set {
                args.push("--zed".into());
                args.push(flatpak_dir.join("libexec").join("zed-editor").into());
            }

            let error = exec::execvp("/usr/bin/flatpak-spawn", args);
            eprintln!("failed restart cli on host: {:?}", error);
            process::exit(1);
        }
    }

    pub fn set_bin_if_no_escape(mut args: super::Args) -> super::Args {
        if env::var(NO_ESCAPE_ENV_NAME).is_ok()
            && env::var("FLATPAK_ID").map_or(false, |id| id.starts_with("dev.zed.Zed"))
        {
            if args.zed.is_none() {
                args.zed = Some("/app/libexec/zed-editor".into());
                env::set_var("ZED_UPDATE_EXPLANATION", "Please use flatpak to update zed");
            }
        }
        args
    }

    fn get_flatpak_dir() -> Option<PathBuf> {
        if env::var(NO_ESCAPE_ENV_NAME).is_ok() {
            return None;
        }

        if let Ok(flatpak_id) = env::var("FLATPAK_ID") {
            if !flatpak_id.starts_with("dev.zed.Zed") {
                return None;
            }

            let install_dir = Command::new("/usr/bin/flatpak-spawn")
                .arg("--host")
                .arg("flatpak")
                .arg("info")
                .arg("--show-location")
                .arg(flatpak_id)
                .output()
                .unwrap();
            let install_dir = PathBuf::from(String::from_utf8(install_dir.stdout).unwrap().trim());
            Some(install_dir.join("files"))
        } else {
            None
        }
    }

    fn get_xdg_env_args() -> Vec<OsString> {
        let xdg_keys = [
            "XDG_DATA_HOME",
            "XDG_CONFIG_HOME",
            "XDG_CACHE_HOME",
            "XDG_STATE_HOME",
        ];
        env::vars()
            .filter(|(key, _)| xdg_keys.contains(&key.as_str()))
            .map(|(key, val)| format!("--env=FLATPAK_{}={}", key, val).into())
            .collect()
    }
}

// todo("windows")
#[cfg(target_os = "windows")]
mod windows {
    use crate::{Detect, InstalledApp};
    use std::io;
    use std::path::Path;
    use std::process::ExitStatus;

    struct App;
    impl InstalledApp for App {
        fn zed_version_string(&self) -> String {
            unimplemented!()
        }
        fn launch(&self, _ipc_url: String) -> anyhow::Result<()> {
            unimplemented!()
        }
        fn run_foreground(&self, _ipc_url: String) -> io::Result<ExitStatus> {
            unimplemented!()
        }
    }

    impl Detect {
        pub fn detect(_path: Option<&Path>) -> anyhow::Result<impl InstalledApp> {
            Ok(App)
        }
    }
}

#[cfg(target_os = "macos")]
mod mac_os {
    use anyhow::{anyhow, Context, Result};
    use core_foundation::{
        array::{CFArray, CFIndex},
        string::kCFStringEncodingUTF8,
        url::{CFURLCreateWithBytes, CFURL},
    };
    use core_services::{kLSLaunchDefaults, LSLaunchURLSpec, LSOpenFromURLSpec, TCFType};
    use serde::Deserialize;
    use std::{
        ffi::OsStr,
        fs, io,
        path::{Path, PathBuf},
        process::{Command, ExitStatus},
        ptr,
    };

    use cli::FORCE_CLI_MODE_ENV_VAR_NAME;

    use crate::{Detect, InstalledApp};

    #[derive(Debug, Deserialize)]
    struct InfoPlist {
        #[serde(rename = "CFBundleShortVersionString")]
        bundle_short_version_string: String,
    }

    enum Bundle {
        App {
            app_bundle: PathBuf,
            plist: InfoPlist,
        },
        LocalPath {
            executable: PathBuf,
            plist: InfoPlist,
        },
    }

    fn locate_bundle() -> Result<PathBuf> {
        let cli_path = std::env::current_exe()?.canonicalize()?;
        let mut app_path = cli_path.clone();
        while app_path.extension() != Some(OsStr::new("app")) {
            if !app_path.pop() {
                return Err(anyhow!("cannot find app bundle containing {:?}", cli_path));
            }
        }
        Ok(app_path)
    }

    impl Detect {
        pub fn detect(path: Option<&Path>) -> anyhow::Result<impl InstalledApp> {
            let bundle_path = if let Some(bundle_path) = path {
                bundle_path
                    .canonicalize()
                    .with_context(|| format!("Args bundle path {bundle_path:?} canonicalization"))?
            } else {
                locate_bundle().context("bundle autodiscovery")?
            };

            match bundle_path.extension().and_then(|ext| ext.to_str()) {
                Some("app") => {
                    let plist_path = bundle_path.join("Contents/Info.plist");
                    let plist =
                        plist::from_file::<_, InfoPlist>(&plist_path).with_context(|| {
                            format!("Reading *.app bundle plist file at {plist_path:?}")
                        })?;
                    Ok(Bundle::App {
                        app_bundle: bundle_path,
                        plist,
                    })
                }
                _ => {
                    println!("Bundle path {bundle_path:?} has no *.app extension, attempting to locate a dev build");
                    let plist_path = bundle_path
                        .parent()
                        .with_context(|| format!("Bundle path {bundle_path:?} has no parent"))?
                        .join("WebRTC.framework/Resources/Info.plist");
                    let plist =
                        plist::from_file::<_, InfoPlist>(&plist_path).with_context(|| {
                            format!("Reading dev bundle plist file at {plist_path:?}")
                        })?;
                    Ok(Bundle::LocalPath {
                        executable: bundle_path,
                        plist,
                    })
                }
            }
        }
    }

    impl InstalledApp for Bundle {
        fn zed_version_string(&self) -> String {
            let is_dev = matches!(self, Self::LocalPath { .. });
            format!(
                "Zed {}{} – {}",
                self.plist().bundle_short_version_string,
                if is_dev { " (dev)" } else { "" },
                self.path().display(),
            )
        }

        fn launch(&self, url: String) -> anyhow::Result<()> {
            match self {
                Self::App { app_bundle, .. } => {
                    let app_path = app_bundle;

                    let status = unsafe {
                        let app_url = CFURL::from_path(app_path, true)
                            .with_context(|| format!("invalid app path {app_path:?}"))?;
                        let url_to_open = CFURL::wrap_under_create_rule(CFURLCreateWithBytes(
                            ptr::null(),
                            url.as_ptr(),
                            url.len() as CFIndex,
                            kCFStringEncodingUTF8,
                            ptr::null(),
                        ));
                        // equivalent to: open zed-cli:... -a /Applications/Zed\ Preview.app
                        let urls_to_open =
                            CFArray::from_copyable(&[url_to_open.as_concrete_TypeRef()]);
                        LSOpenFromURLSpec(
                            &LSLaunchURLSpec {
                                appURL: app_url.as_concrete_TypeRef(),
                                itemURLs: urls_to_open.as_concrete_TypeRef(),
                                passThruParams: ptr::null(),
                                launchFlags: kLSLaunchDefaults,
                                asyncRefCon: ptr::null_mut(),
                            },
                            ptr::null_mut(),
                        )
                    };

                    anyhow::ensure!(
                        status == 0,
                        "cannot start app bundle {}",
                        self.zed_version_string()
                    );
                }

                Self::LocalPath { executable, .. } => {
                    let executable_parent = executable
                        .parent()
                        .with_context(|| format!("Executable {executable:?} path has no parent"))?;
                    let subprocess_stdout_file = fs::File::create(
                        executable_parent.join("zed_dev.log"),
                    )
                    .with_context(|| format!("Log file creation in {executable_parent:?}"))?;
                    let subprocess_stdin_file =
                        subprocess_stdout_file.try_clone().with_context(|| {
                            format!("Cloning descriptor for file {subprocess_stdout_file:?}")
                        })?;
                    let mut command = std::process::Command::new(executable);
                    let command = command
                        .env(FORCE_CLI_MODE_ENV_VAR_NAME, "")
                        .stderr(subprocess_stdout_file)
                        .stdout(subprocess_stdin_file)
                        .arg(url);

                    command
                        .spawn()
                        .with_context(|| format!("Spawning {command:?}"))?;
                }
            }

            Ok(())
        }

        fn run_foreground(&self, ipc_url: String) -> io::Result<ExitStatus> {
            let path = match self {
                Bundle::App { app_bundle, .. } => app_bundle.join("Contents/MacOS/zed"),
                Bundle::LocalPath { executable, .. } => executable.clone(),
            };

            std::process::Command::new(path).arg(ipc_url).status()
        }
    }

    impl Bundle {
        fn plist(&self) -> &InfoPlist {
            match self {
                Self::App { plist, .. } => plist,
                Self::LocalPath { plist, .. } => plist,
            }
        }

        fn path(&self) -> &Path {
            match self {
                Self::App { app_bundle, .. } => app_bundle,
                Self::LocalPath { executable, .. } => executable,
            }
        }
    }

    pub(super) fn spawn_channel_cli(
        channel: release_channel::ReleaseChannel,
        leftover_args: Vec<String>,
    ) -> Result<()> {
        use anyhow::bail;

        let app_id_prompt = format!("id of app \"{}\"", channel.display_name());
        let app_id_output = Command::new("osascript")
            .arg("-e")
            .arg(&app_id_prompt)
            .output()?;
        if !app_id_output.status.success() {
            bail!("Could not determine app id for {}", channel.display_name());
        }
        let app_name = String::from_utf8(app_id_output.stdout)?.trim().to_owned();
        let app_path_prompt = format!("kMDItemCFBundleIdentifier == '{app_name}'");
        let app_path_output = Command::new("mdfind").arg(app_path_prompt).output()?;
        if !app_path_output.status.success() {
            bail!(
                "Could not determine app path for {}",
                channel.display_name()
            );
        }
        let app_path = String::from_utf8(app_path_output.stdout)?.trim().to_owned();
        let cli_path = format!("{app_path}/Contents/MacOS/cli");
        Command::new(cli_path).args(leftover_args).spawn()?;
        Ok(())
    }
}