path: root/src/export_fun.rs

//! allow sub-programs to invoke arbitrary user-defined functions via a unix socket

use crate::Session;
use crate::run::Input;
use crate::run::Output;
use crate::run::get_command_kind;
use std::env::current_exe;
use std::ffi::OsStr;
use std::fs;
use std::fs::File;
use std::io;
use std::io::IoSliceMut;
use std::io::Read;
use std::io::Write;
use std::os::fd::FromRawFd;
use std::os::fd::OwnedFd;
use std::os::unix::ffi::OsStrExt;
use std::os::unix::fs::symlink;
use std::os::unix::net::AncillaryData;
use std::os::unix::net::SocketAncillary;
use std::os::unix::net::UnixListener;
use std::os::unix::net::UnixStream;
use std::path::Path;
use std::path::PathBuf;
use std::process::exit;
use std::sync::Arc;
use std::sync::Mutex;
use std::thread;
use std::time::Duration;

fn handle_server(session: Arc<Mutex<Session>>, mut stream: UnixStream) -> io::Result<()> {
    // TODO: figure out how to get a reasonable limit on CLI arg len
    let mut buf = [0u8; 8192];
    let mut iov = [IoSliceMut::new(&mut buf)];

    let mut ancillary_buf = [0u8; 128];
    let mut ancillary = SocketAncillary::new(&mut ancillary_buf);

    let mut fds: Vec<i32> = Vec::new();

    let bytelen = stream.recv_vectored_with_ancillary(&mut iov, &mut ancillary)?;

    for msg in ancillary.messages() {
        if let Ok(msg) = msg {
            match msg {
                AncillaryData::ScmRights(rights) => {
                    for fd in rights {
                        fds.push(fd);
                    }
                }
                _ => (),
            }
        }
    }

    if fds.len() != 3 {
        // malformed
        return Ok(());
    }

    let Ok(cli_args) = crate::serialization::deserialize_cli_args(&buf[..bytelen]) else {
        // cli args malformed
        return Ok(());
    };

    if cli_args.is_empty() {
        // malformed
        return Ok(());
    };

    let se = session.lock().unwrap();
    match get_command_kind(&se, cli_args[0].as_slice()) {
        crate::run::CommandKind::Fun(_) => (),
        crate::run::CommandKind::Path(_) | crate::run::CommandKind::Builtin(_) => {
            return Ok(());
        }
    }
    drop(se);

    let stdin = File::from(unsafe { OwnedFd::from_raw_fd(fds[0]) });
    let stdout = File::from(unsafe { OwnedFd::from_raw_fd(fds[1]) });

    let res = crate::run::Executor::execute_fun(
        session,
        cli_args,
        Input::File(stdin),
        Output::File(stdout),
    );

    let exit_code = match res {
        Ok(_) => 0,
        Err(e) => match e {
            crate::run::ExecError::UnknownVariable(_) => -3,
            crate::run::ExecError::ExecError(x) => x,
        },
    };

    let _ = stream.set_write_timeout(Some(Duration::from_secs(1)));
    stream.write_all(&exit_code.to_le_bytes())?;

    Ok(())
}

fn handle_client(mut stream: UnixStream) -> io::Result<()> {
    // give up all my file descriptors descriptors
    let mut ancillary_buffer = [0; 128];
    let mut ancillary = SocketAncillary::new(&mut ancillary_buffer[..]);
    ancillary.add_fds(&[0, 1, 2]);

    // cli params
    let buf = crate::serialization::serialize_cli_args();
    let bufs = &mut [io::IoSlice::new(&buf[..])][..];

    // send
    stream.send_vectored_with_ancillary(bufs, &mut ancillary)?;

    // recv exit code
    let mut exit_buf = [0; 4];
    let res = stream.read_exact(&mut exit_buf);

    let exit_code = match res {
        Ok(_) => i32::from_le_bytes(exit_buf),
        Err(_) => -2,
    };

    exit(exit_code)
}

pub fn maybe_run_defined_function() {
    if let Some(program_name) = std::env::args_os().next() {
        let program_name = program_name.as_bytes();
        if !program_name.contains(&b'/') && program_name != b"pish" {
            if let Some(socket) = std::env::var_os("PISH_SOCKET") {
                if let Ok(stream) = UnixStream::connect(socket) {
                    let _ = handle_client(stream);
                }
                exit(-1);
            }
        }
    }
}

fn unique_string() -> String {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};
    use std::process;
    use std::time::{SystemTime, UNIX_EPOCH};

    let pid = process::id();
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_nanos();

    let mut hasher = DefaultHasher::new();
    pid.hash(&mut hasher);
    now.hash(&mut hasher);

    let hash = hasher.finish();
    let hex = format!("{:016x}", hash);

    hex[..12.min(hex.len())].to_string()
}

pub struct SocketRunning {
    bin_dir: PathBuf,
    path: PathBuf,
}

impl SocketRunning {
    pub fn socket_path(&self) -> &Path {
        &self.path
    }
    pub fn path(&self) -> &Path {
        &self.bin_dir
    }
}

#[must_use]
struct SocketDropper {
    session: Arc<Mutex<Session>>,
}

impl Drop for SocketDropper {
    fn drop(&mut self) {
        // mark socket for closing by `take`ing the socket_running value
        let session = &self.session;
        let Ok(mut se) = session.lock() else { return };
        let Some(sr) = se.socket_running.take() else {
            return;
        };

        // connect to the socket for a brief moment to make it check that it should terminate
        if let Ok(mut con) = UnixStream::connect(&sr.path) {
            let _ = write!(con, "please shut down :))");
        };
    }
}

#[must_use]
pub fn listen(session: Arc<Mutex<Session>>) -> impl Drop {
    let session_id = unique_string();
    let session_dir = crate::basedir::data_dir().join("session").join(session_id);
    let bin_dir = session_dir.join("bin");
    std::fs::create_dir_all(&bin_dir).unwrap();
    let socket_path = session_dir.join("cmd.sock");

    {
        let mut se = session.lock().unwrap();
        assert!(se.socket_running.is_none());
        se.socket_running = Some(SocketRunning {
            bin_dir,
            path: socket_path.clone(),
        });
    }

    let se = session.clone();
    thread::spawn(move || {
        struct SessionRemover(PathBuf);
        impl Drop for SessionRemover {
            fn drop(&mut self) {
                let _ = fs::remove_dir_all(&self.0);
            }
        }
        let _session_remover = SessionRemover(session_dir);
        let listener = UnixListener::bind(socket_path).unwrap();
        let mut it = listener.incoming();
        while let Some(stream) = it.next() {
            match se.lock() {
                Err(_) => break,
                Ok(se) if se.socket_running.is_none() => break,
                _ => (),
            }
            if let Ok(stream) = stream {
                let se = se.clone();
                thread::spawn(move || handle_server(se, stream));
            }
        }
    });

    SocketDropper { session }
}

fn create_function_hook_res(
    session: Arc<Mutex<Session>>,
    fun_name: &[u8],
) -> Result<(), Box<dyn std::error::Error>> {
    let session = session.lock().map_err(|e| format!("{e:?}"))?;
    let sock_run = session.socket_running.as_ref().ok_or("no socket running")?;
    let exe_path = current_exe()?;
    let symlink_path = sock_run.bin_dir.join(OsStr::from_bytes(fun_name));
    symlink(exe_path, symlink_path)?;
    Ok(())
}

pub fn create_function_hook(session: Arc<Mutex<Session>>, fun_name: &[u8]) {
    if let Err(e) = create_function_hook_res(session, fun_name) {
        println!("failed to create function hook: {e:?}");
    }
}