tenant.openagents/omega

1use std::{
2    io::{Read, Write},
3    net::{Ipv4Addr, SocketAddr, SocketAddrV4, TcpListener, TcpStream},
4    thread,
5    time::Duration,
6};
7
8use sysinfo::System;
9
10use release_channel::ReleaseChannel;
11
12const LOCALHOST: Ipv4Addr = Ipv4Addr::new(127, 0, 0, 1);
13const CONNECT_TIMEOUT: Duration = Duration::from_millis(10);
14const RECEIVE_TIMEOUT: Duration = Duration::from_millis(35);
15const SEND_TIMEOUT: Duration = Duration::from_millis(20);
16const USER_BLOCK: u16 = 100;
17
18fn address() -> SocketAddr {
19    // These port numbers are offset by the user ID to avoid conflicts between
20    // different users on the same machine. In addition to that the ports for each
21    // release channel are spaced out by 100 to avoid conflicts between different
22    // users running different release channels on the same machine. This ends up
23    // interleaving the ports between different users and different release channels.
24    //
25    // On macOS user IDs start at 501 and on Linux they start at 1000. The first user
26    // on a Mac with ID 501 running a dev channel build will use port 44238, and the
27    // second user with ID 502 will use port 44239, and so on. User 501 will use ports
28    // 44338, 44438, and 44538 for the preview, stable, and nightly channels,
29    // respectively. User 502 will use ports 44339, 44439, and 44539 for the preview,
30    // stable, and nightly channels, respectively.
31    let port = match *release_channel::RELEASE_CHANNEL {
32        ReleaseChannel::Dev => 44737,
33        ReleaseChannel::Preview => 44737 + USER_BLOCK,
34        ReleaseChannel::Stable => 44737 + (2 * USER_BLOCK),
35        ReleaseChannel::Nightly => 44737 + (3 * USER_BLOCK),
36    };
37    let mut user_port = port;
38    let mut sys = System::new_all();
39    sys.refresh_all();
40    if let Ok(current_pid) = sysinfo::get_current_pid()
41        && let Some(uid) = sys
42            .process(current_pid)
43            .and_then(|process| process.user_id())
44    {
45        let uid_u32 = get_uid_as_u32(uid);
46        // Ensure that the user ID is not too large to avoid overflow when
47        // calculating the port number. This seems unlikely but it doesn't
48        // hurt to be safe.
49        let max_port = 65535;
50        let max_uid: u32 = max_port - port as u32;
51        let wrapped_uid: u16 = (uid_u32 % max_uid) as u16;
52        user_port += wrapped_uid;
53    }
54
55    // A custom data directory is a different instance, so it gets its own lock.
56    //
57    // The lock exists to stop two processes sharing one profile. Two processes
58    // pointed at different `--user-data-dir` roots share nothing, yet the port
59    // above keys only on release channel and uid, so the second was refused
60    // with "omega is already running". `ZED_RELEASE_CHANNEL` is
61    // `debug_assertions`-only, so a release build had no way to move the port
62    // either, and every clean-profile proof had to begin by quitting the
63    // owner's live session.
64    //
65    // The offset is derived from the directory path, so the same profile keeps
66    // the same port across launches — a custom instance still refuses a second
67    // copy of *itself*, which is the property worth keeping.
68    if let Some(custom_dir) = paths::custom_data_dir() {
69        let mut hash: u32 = 2_166_136_261;
70        for byte in custom_dir.as_os_str().as_encoded_bytes() {
71            hash ^= u32::from(*byte);
72            hash = hash.wrapping_mul(16_777_619);
73        }
74        // Land well clear of the four release-channel blocks above.
75        const CUSTOM_BASE: u16 = 45_737;
76        const CUSTOM_SPAN: u32 = 4_000;
77        user_port = CUSTOM_BASE + (hash % CUSTOM_SPAN) as u16;
78    }
79
80    SocketAddr::V4(SocketAddrV4::new(LOCALHOST, user_port))
81}
82
83#[cfg(unix)]
84fn get_uid_as_u32(uid: &sysinfo::Uid) -> u32 {
85    *uid.clone()
86}
87
88#[cfg(windows)]
89fn get_uid_as_u32(uid: &sysinfo::Uid) -> u32 {
90    // Extract the RID which is an integer
91    uid.to_string()
92        .rsplit('-')
93        .next()
94        .and_then(|rid| rid.parse::<u32>().ok())
95        .unwrap_or(0)
96}
97
98fn instance_handshake() -> &'static str {
99    match *release_channel::RELEASE_CHANNEL {
100        ReleaseChannel::Dev => "Omega Dev Instance Running",
101        ReleaseChannel::Nightly => "Omega Nightly Instance Running",
102        ReleaseChannel::Preview => "Omega RC Instance Running",
103        ReleaseChannel::Stable => "Omega Instance Running",
104    }
105}
106
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub enum IsOnlyInstance {
109    Yes,
110    No,
111}
112
113pub fn ensure_only_instance() -> IsOnlyInstance {
114    if check_got_handshake() {
115        return IsOnlyInstance::No;
116    }
117
118    let listener = match TcpListener::bind(address()) {
119        Ok(listener) => listener,
120
121        Err(err) => {
122            log::warn!("Error binding to single instance port: {err}");
123            if check_got_handshake() {
124                return IsOnlyInstance::No;
125            }
126
127            // Avoid failing to start when some other application by chance already has
128            // a claim on the port. This is sub-par as any other instance that gets launched
129            // will be unable to communicate with this instance and will duplicate
130            log::warn!("Backup handshake request failed, continuing without handshake");
131            return IsOnlyInstance::Yes;
132        }
133    };
134
135    thread::Builder::new()
136        .name("EnsureSingleton".to_string())
137        .spawn(move || {
138            for stream in listener.incoming() {
139                let mut stream = match stream {
140                    Ok(stream) => stream,
141                    Err(_) => return,
142                };
143
144                _ = stream.set_nodelay(true);
145                _ = stream.set_read_timeout(Some(SEND_TIMEOUT));
146                _ = stream.write_all(instance_handshake().as_bytes());
147            }
148        })
149        .unwrap();
150
151    IsOnlyInstance::Yes
152}
153
154fn check_got_handshake() -> bool {
155    match TcpStream::connect_timeout(&address(), CONNECT_TIMEOUT) {
156        Ok(mut stream) => {
157            let mut buf = vec![0u8; instance_handshake().len()];
158
159            stream.set_read_timeout(Some(RECEIVE_TIMEOUT)).unwrap();
160            if let Err(err) = stream.read_exact(&mut buf) {
161                log::warn!("Connected to single instance port but failed to read: {err}");
162                return false;
163            }
164
165            if buf == instance_handshake().as_bytes() {
166                log::info!("Got instance handshake");
167                return true;
168            }
169
170            log::warn!("Got wrong instance handshake value");
171            false
172        }
173
174        Err(_) => false,
175    }
176}
177