From 66bf56fc4f571c94486c5e2a1552aebf66f5f797 Mon Sep 17 00:00:00 2001
From: Julia <floc@unpromptedtirade.com>
Date: Thu, 6 Jul 2023 17:18:45 -0400
Subject: [PATCH] Prevent duplicate instances by coordinating via a socket

---
 crates/cli/src/main.rs          |  1 +
 crates/zed/src/main.rs          |  9 +++-
 crates/zed/src/only_instance.rs | 82 +++++++++++++++++++++++++++++++++
 crates/zed/src/zed.rs           |  1 +
 4 files changed, 91 insertions(+), 2 deletions(-)
 create mode 100644 crates/zed/src/only_instance.rs

diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs
index bdf677512c..2f742814a8 100644
--- a/crates/cli/src/main.rs
+++ b/crates/cli/src/main.rs
@@ -201,6 +201,7 @@ impl Bundle {
                     self.zed_version_string()
                 );
             }
+
             Self::LocalPath { executable, .. } => {
                 let executable_parent = executable
                     .parent()
diff --git a/crates/zed/src/main.rs b/crates/zed/src/main.rs
index 3da8c24617..5eed301367 100644
--- a/crates/zed/src/main.rs
+++ b/crates/zed/src/main.rs
@@ -57,8 +57,9 @@ use staff_mode::StaffMode;
 use util::{channel::RELEASE_CHANNEL, paths, ResultExt, TryFutureExt};
 use workspace::{item::ItemHandle, notifications::NotifyResultExt, AppState, Workspace};
 use zed::{
-    assets::Assets, build_window_options, handle_keymap_file_changes, initialize_workspace,
-    languages, menus,
+    assets::Assets,
+    build_window_options, handle_keymap_file_changes, initialize_workspace, languages, menus,
+    only_instance::{ensure_only_instance, IsOnlyInstance},
 };
 
 fn main() {
@@ -66,6 +67,10 @@ fn main() {
     init_paths();
     init_logger();
 
+    if ensure_only_instance() != IsOnlyInstance::Yes {
+        return;
+    }
+
     log::info!("========== starting zed ==========");
     let mut app = gpui::App::new(Assets).unwrap();
 
diff --git a/crates/zed/src/only_instance.rs b/crates/zed/src/only_instance.rs
new file mode 100644
index 0000000000..c1358f7a33
--- /dev/null
+++ b/crates/zed/src/only_instance.rs
@@ -0,0 +1,82 @@
+use std::{
+    io::{Read, Write},
+    net::{Ipv4Addr, SocketAddr, SocketAddrV4, TcpListener, TcpStream},
+    thread,
+    time::Duration,
+};
+
+const PORT: u16 = 43739;
+const LOCALHOST: Ipv4Addr = Ipv4Addr::new(127, 0, 0, 1);
+const ADDRESS: SocketAddr = SocketAddr::V4(SocketAddrV4::new(LOCALHOST, PORT));
+const INSTANCE_HANDSHAKE: &str = "Zed Editor Instance Running";
+const CONNECT_TIMEOUT: Duration = Duration::from_millis(10);
+const RECEIVE_TIMEOUT: Duration = Duration::from_millis(35);
+const SEND_TIMEOUT: Duration = Duration::from_millis(20);
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum IsOnlyInstance {
+    Yes,
+    No,
+}
+
+pub fn ensure_only_instance() -> IsOnlyInstance {
+    if check_got_handshake() {
+        return IsOnlyInstance::No;
+    }
+
+    let listener = match TcpListener::bind(ADDRESS) {
+        Ok(listener) => listener,
+
+        Err(err) => {
+            log::warn!("Error binding to single instance port: {err}");
+            if check_got_handshake() {
+                return IsOnlyInstance::No;
+            }
+
+            // Avoid failing to start when some other application by chance already has
+            // a claim on the port. This is sub-par as any other instance that gets launched
+            // will be unable to communicate with this instance and will duplicate
+            log::warn!("Backup handshake request failed, continuing without handshake");
+            return IsOnlyInstance::Yes;
+        }
+    };
+
+    thread::spawn(move || {
+        for stream in listener.incoming() {
+            let mut stream = match stream {
+                Ok(stream) => stream,
+                Err(_) => return,
+            };
+
+            _ = stream.set_nodelay(true);
+            _ = stream.set_read_timeout(Some(SEND_TIMEOUT));
+            _ = stream.write_all(INSTANCE_HANDSHAKE.as_bytes());
+        }
+    });
+
+    IsOnlyInstance::Yes
+}
+
+fn check_got_handshake() -> bool {
+    match TcpStream::connect_timeout(&ADDRESS, CONNECT_TIMEOUT) {
+        Ok(mut stream) => {
+            let mut buf = vec![0u8; INSTANCE_HANDSHAKE.len()];
+
+            stream.set_read_timeout(Some(RECEIVE_TIMEOUT)).unwrap();
+            if let Err(err) = stream.read_exact(&mut buf) {
+                log::warn!("Connected to single instance port but failed to read: {err}");
+                return false;
+            }
+
+            if buf == INSTANCE_HANDSHAKE.as_bytes() {
+                log::info!("Got instance handshake");
+                return true;
+            }
+
+            log::warn!("Got wrong instance handshake value");
+            false
+        }
+
+        Err(_) => false,
+    }
+}
diff --git a/crates/zed/src/zed.rs b/crates/zed/src/zed.rs
index 0df16f4bab..09bdbf65be 100644
--- a/crates/zed/src/zed.rs
+++ b/crates/zed/src/zed.rs
@@ -1,6 +1,7 @@
 pub mod assets;
 pub mod languages;
 pub mod menus;
+pub mod only_instance;
 #[cfg(any(test, feature = "test-support"))]
 pub mod test;