Skip to main content

hx_remote/
server.rs

1use crate::{
2    SocketRequest, SocketResponse, absolute_path, read_lsp_message, show_document_params,
3    write_lsp_message,
4};
5use serde_json::{Value, json};
6use std::collections::{HashMap, VecDeque};
7use std::fs;
8use std::io::{self, BufReader, ErrorKind, Read, Write};
9use std::os::unix::fs::{FileTypeExt, MetadataExt, PermissionsExt};
10use std::os::unix::net::{UnixListener, UnixStream};
11use std::path::{Path, PathBuf};
12use std::sync::mpsc::{self, Receiver, Sender};
13use std::thread;
14use std::time::Duration;
15use tempfile::{Builder, NamedTempFile};
16
17const MAX_SOCKET_REQUEST_BYTES: u64 = 32 * 1024 * 1024;
18
19enum Event {
20    Lsp(Value),
21    LspClosed,
22    FatalInput(String),
23    Open {
24        request: SocketRequest,
25        reply: Sender<Result<String, String>>,
26    },
27    Stop,
28}
29
30#[derive(Clone, Copy, PartialEq, Eq)]
31enum ConnectionControl {
32    None,
33    Stop,
34    ForceStop,
35}
36
37#[derive(Clone)]
38struct SocketIdentity {
39    path: PathBuf,
40    device: u64,
41    inode: u64,
42}
43
44struct SocketGuard {
45    identity: SocketIdentity,
46}
47
48impl SocketIdentity {
49    fn remove_if_same(&self) {
50        let Ok(metadata) = fs::symlink_metadata(&self.path) else {
51            return;
52        };
53        if metadata.dev() == self.device && metadata.ino() == self.inode {
54            let _ = fs::remove_file(&self.path);
55        }
56    }
57}
58
59impl Drop for SocketGuard {
60    fn drop(&mut self) {
61        self.identity.remove_if_same();
62    }
63}
64
65pub fn run_server(socket_path: PathBuf) -> Result<(), Box<dyn std::error::Error>> {
66    let (listener, socket_guard) = bind_socket(&socket_path)?;
67    let (event_tx, event_rx) = mpsc::channel();
68
69    spawn_lsp_reader(event_tx.clone());
70    spawn_socket_listener(listener, event_tx, socket_guard.identity.clone());
71
72    event_loop(event_rx)
73}
74
75fn bind_socket(socket_path: &Path) -> io::Result<(UnixListener, SocketGuard)> {
76    if let Some(parent) = socket_path
77        .parent()
78        .filter(|parent| !parent.as_os_str().is_empty())
79    {
80        fs::create_dir_all(parent)?;
81    }
82
83    let listener = match UnixListener::bind(socket_path) {
84        Ok(listener) => listener,
85        Err(error) if error.kind() == ErrorKind::AddrInUse => {
86            if UnixStream::connect(socket_path).is_ok() {
87                return Err(io::Error::new(
88                    ErrorKind::AddrInUse,
89                    format!(
90                        "another bridge is already listening on {}",
91                        socket_path.display()
92                    ),
93                ));
94            }
95
96            let metadata = fs::symlink_metadata(socket_path)?;
97            if !metadata.file_type().is_socket() {
98                return Err(io::Error::new(
99                    ErrorKind::AlreadyExists,
100                    format!(
101                        "refusing to replace non-socket path {}",
102                        socket_path.display()
103                    ),
104                ));
105            }
106            fs::remove_file(socket_path)?;
107            UnixListener::bind(socket_path)?
108        }
109        Err(error) => return Err(error),
110    };
111
112    fs::set_permissions(socket_path, fs::Permissions::from_mode(0o600))?;
113    let metadata = fs::symlink_metadata(socket_path)?;
114    let guard = SocketGuard {
115        identity: SocketIdentity {
116            path: socket_path.to_path_buf(),
117            device: metadata.dev(),
118            inode: metadata.ino(),
119        },
120    };
121    Ok((listener, guard))
122}
123
124fn spawn_lsp_reader(event_tx: Sender<Event>) {
125    thread::spawn(move || {
126        let stdin = io::stdin();
127        let mut reader = BufReader::new(stdin.lock());
128        loop {
129            match read_lsp_message(&mut reader) {
130                Ok(Some(message)) => {
131                    if event_tx.send(Event::Lsp(message)).is_err() {
132                        return;
133                    }
134                }
135                Ok(None) => {
136                    let _ = event_tx.send(Event::LspClosed);
137                    return;
138                }
139                Err(error) => {
140                    let _ = event_tx.send(Event::FatalInput(error.to_string()));
141                    return;
142                }
143            }
144        }
145    });
146}
147
148fn spawn_socket_listener(
149    listener: UnixListener,
150    event_tx: Sender<Event>,
151    socket_identity: SocketIdentity,
152) {
153    thread::spawn(move || {
154        for connection in listener.incoming() {
155            match connection {
156                Ok(stream) => {
157                    let event_tx = event_tx.clone();
158                    let socket_identity = socket_identity.clone();
159                    thread::spawn(move || {
160                        handle_socket_connection(stream, event_tx, socket_identity)
161                    });
162                }
163                Err(error) => {
164                    eprintln!("hxr: socket accept failed: {error}");
165                    return;
166                }
167            }
168        }
169    });
170}
171
172fn handle_socket_connection(
173    mut stream: UnixStream,
174    event_tx: Sender<Event>,
175    socket_identity: SocketIdentity,
176) {
177    let mut control = ConnectionControl::None;
178    let result = (|| -> Result<String, String> {
179        stream
180            .set_read_timeout(Some(Duration::from_secs(10)))
181            .map_err(|error| error.to_string())?;
182        stream
183            .set_write_timeout(Some(Duration::from_secs(10)))
184            .map_err(|error| error.to_string())?;
185
186        let mut request_bytes = Vec::new();
187        BufReader::new(&stream)
188            .take(MAX_SOCKET_REQUEST_BYTES + 1)
189            .read_to_end(&mut request_bytes)
190            .map_err(|error| error.to_string())?;
191        if request_bytes.len() as u64 > MAX_SOCKET_REQUEST_BYTES {
192            return Err(format!(
193                "request exceeds the {} MiB limit",
194                MAX_SOCKET_REQUEST_BYTES / 1024 / 1024
195            ));
196        }
197
198        let request: SocketRequest = serde_json::from_slice(&request_bytes)
199            .map_err(|error| format!("invalid bridge request: {error}"))?;
200        match request {
201            SocketRequest::Stop => {
202                control = ConnectionControl::Stop;
203                Ok("server stopping".to_owned())
204            }
205            SocketRequest::ForceStop => {
206                control = ConnectionControl::ForceStop;
207                Ok("server force-stopping".to_owned())
208            }
209            request => {
210                let (reply_tx, reply_rx) = mpsc::channel();
211                event_tx
212                    .send(Event::Open {
213                        request,
214                        reply: reply_tx,
215                    })
216                    .map_err(|_| "the LSP bridge has stopped".to_owned())?;
217                reply_rx
218                    .recv_timeout(Duration::from_secs(10))
219                    .map_err(|_| "the LSP bridge did not accept the request in time".to_owned())?
220            }
221        }
222    })();
223
224    let response = match result {
225        Ok(message) => SocketResponse::success(message),
226        Err(message) => SocketResponse::error(message),
227    };
228    if serde_json::to_writer(&mut stream, &response).is_ok() {
229        let _ = stream.write_all(b"\n");
230        let _ = stream.flush();
231    }
232
233    if control == ConnectionControl::ForceStop {
234        socket_identity.remove_if_same();
235        // SAFETY: Sending SIGKILL to our own process is the requested forced-stop
236        // behavior. getpid always returns the id of this process.
237        let signal_result = unsafe { libc::kill(libc::getpid(), libc::SIGKILL) };
238        if signal_result == -1 {
239            std::process::abort();
240        }
241        loop {
242            thread::park();
243        }
244    } else if control == ConnectionControl::Stop {
245        let _ = event_tx.send(Event::Stop);
246    }
247}
248
249fn event_loop(event_rx: Receiver<Event>) -> Result<(), Box<dyn std::error::Error>> {
250    let mut stdout = io::stdout().lock();
251    let mut initialized = false;
252    let mut shutting_down = false;
253    let mut queued = VecDeque::new();
254    let mut next_request_id = 1_u64;
255    let mut pending_requests: HashMap<u64, String> = HashMap::new();
256    let mut scratch_files: Vec<NamedTempFile> = Vec::new();
257
258    while let Ok(event) = event_rx.recv() {
259        match event {
260            Event::Lsp(message) => {
261                if let Some(method) = message.get("method").and_then(Value::as_str) {
262                    match method {
263                        "initialize" => {
264                            if let Some(id) = message.get("id") {
265                                let response = json!({
266                                    "jsonrpc": "2.0",
267                                    "id": id,
268                                    "result": {
269                                        "capabilities": {},
270                                        "serverInfo": {
271                                            "name": "hx-remote",
272                                            "version": env!("CARGO_PKG_VERSION")
273                                        }
274                                    }
275                                });
276                                write_lsp_message(&mut stdout, &response)?;
277                            }
278                        }
279                        "initialized" => {
280                            initialized = true;
281                            while let Some(request) = queued.pop_front() {
282                                dispatch_open(
283                                    request,
284                                    &mut next_request_id,
285                                    &mut stdout,
286                                    &mut pending_requests,
287                                    &mut scratch_files,
288                                )?;
289                            }
290                        }
291                        "shutdown" => {
292                            shutting_down = true;
293                            if let Some(id) = message.get("id") {
294                                write_lsp_message(
295                                    &mut stdout,
296                                    &json!({"jsonrpc": "2.0", "id": id, "result": null}),
297                                )?;
298                            }
299                        }
300                        "exit" => return Ok(()),
301                        _ => {}
302                    }
303                } else if let Some(id) = message.get("id").and_then(Value::as_u64)
304                    && let Some(target) = pending_requests.remove(&id)
305                {
306                    let success = message
307                        .get("result")
308                        .and_then(|result| result.get("success"))
309                        .and_then(Value::as_bool);
310                    if success == Some(false) || message.get("error").is_some() {
311                        eprintln!("hxr: Helix did not open {target}");
312                    }
313                }
314            }
315            Event::Open { request, reply } => {
316                if shutting_down {
317                    let _ = reply.send(Err("Helix is shutting down".into()));
318                } else if initialized {
319                    let result = dispatch_open(
320                        request,
321                        &mut next_request_id,
322                        &mut stdout,
323                        &mut pending_requests,
324                        &mut scratch_files,
325                    )
326                    .map(|target| format!("sent {target} to Helix"))
327                    .map_err(|error| error.to_string());
328                    let _ = reply.send(result);
329                } else {
330                    queued.push_back(request);
331                    let _ = reply.send(Ok("queued until Helix finishes initializing".into()));
332                }
333            }
334            Event::Stop => return Ok(()),
335            Event::LspClosed => return Ok(()),
336            Event::FatalInput(error) => {
337                return Err(format!("invalid LSP input: {error}").into());
338            }
339        }
340    }
341    Ok(())
342}
343
344fn dispatch_open(
345    request: SocketRequest,
346    next_request_id: &mut u64,
347    stdout: &mut impl Write,
348    pending_requests: &mut HashMap<u64, String>,
349    scratch_files: &mut Vec<NamedTempFile>,
350) -> io::Result<String> {
351    let (path, line, column) = match request {
352        SocketRequest::Open { path, line, column } => (absolute_path(&path)?, line, column),
353        SocketRequest::OpenStdin { contents, name } => {
354            let safe_name = sanitize_scratch_name(&name);
355            let mut file = Builder::new()
356                .prefix("hx-remote-")
357                .suffix(&format!("-{safe_name}"))
358                .tempfile()?;
359            file.write_all(contents.as_bytes())?;
360            file.flush()?;
361            let path = file.path().to_path_buf();
362            scratch_files.push(file);
363            (path, None, None)
364        }
365        SocketRequest::Stop | SocketRequest::ForceStop => {
366            return Err(io::Error::new(
367                ErrorKind::InvalidInput,
368                "control request cannot be dispatched as an open request",
369            ));
370        }
371    };
372
373    let params = show_document_params(&path, line, column)
374        .map_err(|message| io::Error::new(ErrorKind::InvalidInput, message))?;
375    let id = *next_request_id;
376    *next_request_id = next_request_id
377        .checked_add(1)
378        .ok_or_else(|| io::Error::other("LSP request id overflow"))?;
379    let request = json!({
380        "jsonrpc": "2.0",
381        "id": id,
382        "method": "window/showDocument",
383        "params": params
384    });
385    write_lsp_message(stdout, &request)?;
386
387    let target = path.display().to_string();
388    pending_requests.insert(id, target.clone());
389    Ok(target)
390}
391
392fn sanitize_scratch_name(name: &str) -> String {
393    let name = Path::new(name)
394        .file_name()
395        .and_then(|name| name.to_str())
396        .unwrap_or("stdin.txt");
397    let sanitized: String = name
398        .chars()
399        .map(|character| {
400            if character.is_ascii_alphanumeric() || matches!(character, '.' | '-' | '_') {
401                character
402            } else {
403                '_'
404            }
405        })
406        .take(80)
407        .collect();
408    if sanitized.is_empty() {
409        "stdin.txt".into()
410    } else {
411        sanitized
412    }
413}
414
415#[cfg(test)]
416mod tests {
417    use super::sanitize_scratch_name;
418
419    #[test]
420    fn scratch_names_cannot_escape_the_temp_directory() {
421        assert_eq!(
422            sanitize_scratch_name("../../my patch.diff"),
423            "my_patch.diff"
424        );
425        assert_eq!(sanitize_scratch_name(""), "stdin.txt");
426    }
427}