1use crate::HeadlessProject;
  2use anyhow::{anyhow, Context, Result};
  3use fs::RealFs;
  4use futures::channel::mpsc;
  5use futures::{select, select_biased, AsyncRead, AsyncWrite, AsyncWriteExt, FutureExt, SinkExt};
  6use gpui::{AppContext, Context as _};
  7use remote::proxy::ProxyLaunchError;
  8use remote::ssh_session::ChannelClient;
  9use remote::{
 10    json_log::LogRecord,
 11    protocol::{read_message, write_message},
 12};
 13use rpc::proto::Envelope;
 14use smol::channel::{Receiver, Sender};
 15use smol::io::AsyncReadExt;
 16use smol::Async;
 17use smol::{net::unix::UnixListener, stream::StreamExt as _};
 18use std::{
 19    env,
 20    io::Write,
 21    mem,
 22    path::{Path, PathBuf},
 23    sync::Arc,
 24};
 25
 26fn init_logging_proxy() {
 27    env_logger::builder()
 28        .format(|buf, record| {
 29            let mut log_record = LogRecord::new(record);
 30            log_record.message = format!("(remote proxy) {}", log_record.message);
 31            serde_json::to_writer(&mut *buf, &log_record)?;
 32            buf.write_all(b"\n")?;
 33            Ok(())
 34        })
 35        .init();
 36}
 37
 38fn init_logging_server(log_file_path: PathBuf) -> Result<Receiver<Vec<u8>>> {
 39    struct MultiWrite {
 40        file: Box<dyn std::io::Write + Send + 'static>,
 41        channel: Sender<Vec<u8>>,
 42        buffer: Vec<u8>,
 43    }
 44
 45    impl std::io::Write for MultiWrite {
 46        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
 47            let written = self.file.write(buf)?;
 48            self.buffer.extend_from_slice(&buf[..written]);
 49            Ok(written)
 50        }
 51
 52        fn flush(&mut self) -> std::io::Result<()> {
 53            self.channel
 54                .send_blocking(self.buffer.clone())
 55                .map_err(|error| std::io::Error::new(std::io::ErrorKind::Other, error))?;
 56            self.buffer.clear();
 57            self.file.flush()
 58        }
 59    }
 60
 61    let log_file = Box::new(if log_file_path.exists() {
 62        std::fs::OpenOptions::new()
 63            .append(true)
 64            .open(&log_file_path)
 65            .context("Failed to open log file in append mode")?
 66    } else {
 67        std::fs::File::create(&log_file_path).context("Failed to create log file")?
 68    });
 69
 70    let (tx, rx) = smol::channel::unbounded();
 71
 72    let target = Box::new(MultiWrite {
 73        file: log_file,
 74        channel: tx,
 75        buffer: Vec::new(),
 76    });
 77
 78    env_logger::Builder::from_default_env()
 79        .target(env_logger::Target::Pipe(target))
 80        .format(|buf, record| {
 81            let mut log_record = LogRecord::new(record);
 82            log_record.message = format!("(remote server) {}", log_record.message);
 83            serde_json::to_writer(&mut *buf, &log_record)?;
 84            buf.write_all(b"\n")?;
 85            Ok(())
 86        })
 87        .init();
 88
 89    Ok(rx)
 90}
 91
 92fn init_panic_hook() {
 93    std::panic::set_hook(Box::new(|info| {
 94        let payload = info
 95            .payload()
 96            .downcast_ref::<&str>()
 97            .map(|s| s.to_string())
 98            .or_else(|| info.payload().downcast_ref::<String>().cloned())
 99            .unwrap_or_else(|| "Box<Any>".to_string());
100
101        let backtrace = backtrace::Backtrace::new();
102        let mut backtrace = backtrace
103            .frames()
104            .iter()
105            .flat_map(|frame| {
106                frame
107                    .symbols()
108                    .iter()
109                    .filter_map(|frame| Some(format!("{:#}", frame.name()?)))
110            })
111            .collect::<Vec<_>>();
112
113        // Strip out leading stack frames for rust panic-handling.
114        if let Some(ix) = backtrace
115            .iter()
116            .position(|name| name == "rust_begin_unwind")
117        {
118            backtrace.drain(0..=ix);
119        }
120
121        log::error!(
122            "panic occurred: {}\nBacktrace:\n{}",
123            payload,
124            backtrace.join("\n")
125        );
126
127        std::process::abort();
128    }));
129}
130
131struct ServerListeners {
132    stdin: UnixListener,
133    stdout: UnixListener,
134    stderr: UnixListener,
135}
136
137impl ServerListeners {
138    pub fn new(stdin_path: PathBuf, stdout_path: PathBuf, stderr_path: PathBuf) -> Result<Self> {
139        Ok(Self {
140            stdin: UnixListener::bind(stdin_path).context("failed to bind stdin socket")?,
141            stdout: UnixListener::bind(stdout_path).context("failed to bind stdout socket")?,
142            stderr: UnixListener::bind(stderr_path).context("failed to bind stderr socket")?,
143        })
144    }
145}
146
147fn start_server(
148    listeners: ServerListeners,
149    mut log_rx: Receiver<Vec<u8>>,
150    cx: &mut AppContext,
151) -> Arc<ChannelClient> {
152    // This is the server idle timeout. If no connection comes in in this timeout, the server will shut down.
153    const IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10 * 60);
154
155    let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
156    let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded::<Envelope>();
157    let (app_quit_tx, mut app_quit_rx) = mpsc::unbounded::<()>();
158
159    cx.on_app_quit(move |_| {
160        let mut app_quit_tx = app_quit_tx.clone();
161        async move {
162            log::info!("app quitting. sending signal to server main loop");
163            app_quit_tx.send(()).await.ok();
164        }
165    })
166    .detach();
167
168    cx.spawn(|cx| async move {
169        let mut stdin_incoming = listeners.stdin.incoming();
170        let mut stdout_incoming = listeners.stdout.incoming();
171        let mut stderr_incoming = listeners.stderr.incoming();
172
173        loop {
174            let streams = futures::future::join3(stdin_incoming.next(), stdout_incoming.next(), stderr_incoming.next());
175
176            log::info!("accepting new connections");
177            let result = select! {
178                streams = streams.fuse() => {
179                    log::warn!("stdin {:?}, stdout: {:?}, stderr: {:?}", streams.0, streams.1, streams.2);
180                    let (Some(Ok(stdin_stream)), Some(Ok(stdout_stream)), Some(Ok(stderr_stream))) = streams else {
181                        break;
182                    };
183                    anyhow::Ok((stdin_stream, stdout_stream, stderr_stream))
184                }
185                _ = futures::FutureExt::fuse(smol::Timer::after(IDLE_TIMEOUT)) => {
186                    log::warn!("timed out waiting for new connections after {:?}. exiting.", IDLE_TIMEOUT);
187                    cx.update(|cx| {
188                        // TODO: This is a hack, because in a headless project, shutdown isn't executed
189                        // when calling quit, but it should be.
190                        cx.shutdown();
191                        cx.quit();
192                    })?;
193                    break;
194                }
195                _ = app_quit_rx.next().fuse() => {
196                    break;
197                }
198            };
199
200            let Ok((mut stdin_stream, mut stdout_stream, mut stderr_stream)) = result else {
201                break;
202            };
203
204            log::info!("yep! we got connections");
205
206            let mut input_buffer = Vec::new();
207            let mut output_buffer = Vec::new();
208            loop {
209                select_biased! {
210                    _ = app_quit_rx.next().fuse() => {
211                        return anyhow::Ok(());
212                    }
213
214                    stdin_message = read_message(&mut stdin_stream, &mut input_buffer).fuse() => {
215                        let message = match stdin_message {
216                            Ok(message) => message,
217                            Err(error) => {
218                                log::warn!("error reading message on stdin: {}. exiting.", error);
219                                break;
220                            }
221                        };
222                        if let Err(error) = incoming_tx.unbounded_send(message) {
223                            log::error!("failed to send message to application: {:?}. exiting.", error);
224                            return Err(anyhow!(error));
225                        }
226                    }
227
228                    outgoing_message  = outgoing_rx.next().fuse() => {
229                        let Some(message) = outgoing_message else {
230                            log::error!("stdout handler, no message");
231                            break;
232                        };
233
234                        if let Err(error) =
235                            write_message(&mut stdout_stream, &mut output_buffer, message).await
236                        {
237                            log::error!("failed to write stdout message: {:?}", error);
238                            break;
239                        }
240                        if let Err(error) = stdout_stream.flush().await {
241                            log::error!("failed to flush stdout message: {:?}", error);
242                            break;
243                        }
244                    }
245
246                    // // TODO: How do we handle backpressure?
247                    log_message = log_rx.next().fuse() => {
248                        if let Some(log_message) = log_message {
249                            if let Err(error) = stderr_stream.write_all(&log_message).await {
250                                log::error!("failed to write log message to stderr: {:?}", error);
251                                break;
252                            }
253                            if let Err(error) = stderr_stream.flush().await {
254                                log::error!("failed to flush stderr stream: {:?}", error);
255                                break;
256                            }
257                        }
258                    }
259                }
260            }
261        }
262        anyhow::Ok(())
263    })
264    .detach();
265
266    ChannelClient::new(incoming_rx, outgoing_tx, cx)
267}
268
269pub fn execute_run(
270    log_file: PathBuf,
271    pid_file: PathBuf,
272    stdin_socket: PathBuf,
273    stdout_socket: PathBuf,
274    stderr_socket: PathBuf,
275) -> Result<()> {
276    let log_rx = init_logging_server(log_file)?;
277    init_panic_hook();
278
279    log::info!(
280        "starting up. pid_file: {:?}, stdin_socket: {:?}, stdout_socket: {:?}, stderr_socket: {:?}",
281        pid_file,
282        stdin_socket,
283        stdout_socket,
284        stderr_socket
285    );
286
287    write_pid_file(&pid_file)
288        .with_context(|| format!("failed to write pid file: {:?}", &pid_file))?;
289
290    let listeners = ServerListeners::new(stdin_socket, stdout_socket, stderr_socket)?;
291
292    log::debug!("starting gpui app");
293    gpui::App::headless().run(move |cx| {
294        settings::init(cx);
295        HeadlessProject::init(cx);
296
297        log::info!("gpui app started, initializing server");
298        let session = start_server(listeners, log_rx, cx);
299
300        let project = cx.new_model(|cx| {
301            HeadlessProject::new(session, Arc::new(RealFs::new(Default::default(), None)), cx)
302        });
303
304        mem::forget(project);
305    });
306    log::info!("gpui app is shut down. quitting.");
307    Ok(())
308}
309
310#[derive(Clone)]
311struct ServerPaths {
312    log_file: PathBuf,
313    pid_file: PathBuf,
314    stdin_socket: PathBuf,
315    stdout_socket: PathBuf,
316    stderr_socket: PathBuf,
317}
318
319impl ServerPaths {
320    fn new(identifier: &str) -> Result<Self> {
321        let project_dir = create_state_directory(identifier)?;
322
323        let pid_file = project_dir.join("server.pid");
324        let stdin_socket = project_dir.join("stdin.sock");
325        let stdout_socket = project_dir.join("stdout.sock");
326        let stderr_socket = project_dir.join("stderr.sock");
327        let log_file = project_dir.join("server.log");
328
329        Ok(Self {
330            pid_file,
331            stdin_socket,
332            stdout_socket,
333            stderr_socket,
334            log_file,
335        })
336    }
337}
338
339pub fn execute_proxy(identifier: String, is_reconnecting: bool) -> Result<()> {
340    init_logging_proxy();
341    init_panic_hook();
342
343    log::debug!("starting up. PID: {}", std::process::id());
344
345    let server_paths = ServerPaths::new(&identifier)?;
346
347    let server_pid = check_pid_file(&server_paths.pid_file)?;
348    let server_running = server_pid.is_some();
349    if is_reconnecting {
350        if !server_running {
351            log::error!("attempted to reconnect, but no server running");
352            return Err(anyhow!(ProxyLaunchError::ServerNotRunning));
353        }
354    } else {
355        if let Some(pid) = server_pid {
356            log::debug!("found server already running with PID {}. Killing process and cleaning up files...", pid);
357            kill_running_server(pid, &server_paths)?;
358        }
359
360        spawn_server(&server_paths)?;
361    }
362
363    let stdin_task = smol::spawn(async move {
364        let stdin = Async::new(std::io::stdin())?;
365        let stream = smol::net::unix::UnixStream::connect(&server_paths.stdin_socket).await?;
366        handle_io(stdin, stream, "stdin").await
367    });
368
369    let stdout_task: smol::Task<Result<()>> = smol::spawn(async move {
370        let stdout = Async::new(std::io::stdout())?;
371        let stream = smol::net::unix::UnixStream::connect(&server_paths.stdout_socket).await?;
372        handle_io(stream, stdout, "stdout").await
373    });
374
375    let stderr_task: smol::Task<Result<()>> = smol::spawn(async move {
376        let mut stderr = Async::new(std::io::stderr())?;
377        let mut stream = smol::net::unix::UnixStream::connect(&server_paths.stderr_socket).await?;
378        let mut stderr_buffer = vec![0; 2048];
379        loop {
380            match stream.read(&mut stderr_buffer).await {
381                Ok(0) => {
382                    return anyhow::Ok(());
383                }
384                Ok(n) => {
385                    stderr.write_all(&mut stderr_buffer[..n]).await?;
386                    stderr.flush().await?;
387                }
388                Err(error) => {
389                    Err(anyhow!("error reading stderr: {error:?}"))?;
390                }
391            }
392        }
393    });
394
395    if let Err(forwarding_result) = smol::block_on(async move {
396        futures::select! {
397            result = stdin_task.fuse() => result,
398            result = stdout_task.fuse() => result,
399            result = stderr_task.fuse() => result,
400        }
401    }) {
402        log::error!(
403            "failed to forward messages: {:?}, terminating...",
404            forwarding_result
405        );
406        return Err(forwarding_result);
407    }
408
409    Ok(())
410}
411
412fn create_state_directory(identifier: &str) -> Result<PathBuf> {
413    let home_dir = env::var("HOME").unwrap_or_else(|_| ".".to_string());
414    let server_dir = PathBuf::from(home_dir)
415        .join(".local")
416        .join("state")
417        .join("zed-remote-server")
418        .join(identifier);
419
420    std::fs::create_dir_all(&server_dir)?;
421
422    Ok(server_dir)
423}
424
425fn kill_running_server(pid: u32, paths: &ServerPaths) -> Result<()> {
426    log::info!("killing existing server with PID {}", pid);
427    std::process::Command::new("kill")
428        .arg(pid.to_string())
429        .output()
430        .context("failed to kill existing server")?;
431
432    for file in [
433        &paths.pid_file,
434        &paths.stdin_socket,
435        &paths.stdout_socket,
436        &paths.stderr_socket,
437    ] {
438        log::debug!("cleaning up file {:?} before starting new server", file);
439        std::fs::remove_file(file).ok();
440    }
441    Ok(())
442}
443
444fn spawn_server(paths: &ServerPaths) -> Result<()> {
445    if paths.stdin_socket.exists() {
446        std::fs::remove_file(&paths.stdin_socket)?;
447    }
448    if paths.stdout_socket.exists() {
449        std::fs::remove_file(&paths.stdout_socket)?;
450    }
451    if paths.stderr_socket.exists() {
452        std::fs::remove_file(&paths.stderr_socket)?;
453    }
454
455    let binary_name = std::env::current_exe()?;
456    let server_process = std::process::Command::new(binary_name)
457        .arg("run")
458        .arg("--log-file")
459        .arg(&paths.log_file)
460        .arg("--pid-file")
461        .arg(&paths.pid_file)
462        .arg("--stdin-socket")
463        .arg(&paths.stdin_socket)
464        .arg("--stdout-socket")
465        .arg(&paths.stdout_socket)
466        .arg("--stderr-socket")
467        .arg(&paths.stderr_socket)
468        .spawn()?;
469
470    log::debug!("server started. PID: {:?}", server_process.id());
471
472    let mut total_time_waited = std::time::Duration::from_secs(0);
473    let wait_duration = std::time::Duration::from_millis(20);
474    while !paths.stdout_socket.exists()
475        || !paths.stdin_socket.exists()
476        || !paths.stderr_socket.exists()
477    {
478        log::debug!("waiting for server to be ready to accept connections...");
479        std::thread::sleep(wait_duration);
480        total_time_waited += wait_duration;
481    }
482
483    log::info!(
484        "server ready to accept connections. total time waited: {:?}",
485        total_time_waited
486    );
487    Ok(())
488}
489
490fn check_pid_file(path: &Path) -> Result<Option<u32>> {
491    let Some(pid) = std::fs::read_to_string(&path)
492        .ok()
493        .and_then(|contents| contents.parse::<u32>().ok())
494    else {
495        return Ok(None);
496    };
497
498    log::debug!("Checking if process with PID {} exists...", pid);
499    match std::process::Command::new("kill")
500        .arg("-0")
501        .arg(pid.to_string())
502        .output()
503    {
504        Ok(output) if output.status.success() => {
505            log::debug!("Process with PID {} exists. NOT spawning new server, but attaching to existing one.", pid);
506            Ok(Some(pid))
507        }
508        _ => {
509            log::debug!(
510                "Found PID file, but process with that PID does not exist. Removing PID file."
511            );
512            std::fs::remove_file(&path).context("Failed to remove PID file")?;
513            Ok(None)
514        }
515    }
516}
517
518fn write_pid_file(path: &Path) -> Result<()> {
519    if path.exists() {
520        std::fs::remove_file(path)?;
521    }
522    let pid = std::process::id().to_string();
523    log::debug!("writing PID {} to file {:?}", pid, path);
524    std::fs::write(path, pid).context("Failed to write PID file")
525}
526
527async fn handle_io<R, W>(mut reader: R, mut writer: W, socket_name: &str) -> Result<()>
528where
529    R: AsyncRead + Unpin,
530    W: AsyncWrite + Unpin,
531{
532    use remote::protocol::read_message_raw;
533
534    let mut buffer = Vec::new();
535    loop {
536        read_message_raw(&mut reader, &mut buffer)
537            .await
538            .with_context(|| format!("failed to read message from {}", socket_name))?;
539
540        write_size_prefixed_buffer(&mut writer, &mut buffer)
541            .await
542            .with_context(|| format!("failed to write message to {}", socket_name))?;
543
544        writer.flush().await?;
545
546        buffer.clear();
547    }
548}
549
550async fn write_size_prefixed_buffer<S: AsyncWrite + Unpin>(
551    stream: &mut S,
552    buffer: &mut Vec<u8>,
553) -> Result<()> {
554    let len = buffer.len() as u32;
555    stream.write_all(len.to_le_bytes().as_slice()).await?;
556    stream.write_all(buffer).await?;
557    Ok(())
558}