Merge pull request #10 from WhyNotHugo/cleanups

Add support for sandboxed socket path

src/main.rs

@@ -3,6 +3,7 @@ extern crate nix;
 #[cfg(windows)]
 extern crate named_pipe;
 
+use std::convert::TryInto;
 use std::io::{self, Read, Write};
 use std::thread;
 use byteorder::{ByteOrder, NativeEndian, WriteBytesExt};
@@ -11,23 +12,23 @@ mod proxy_socket;
 
 use proxy_socket::ProxySocket;
 
-const BUFFER_SIZE: u32 = 1024 ^ 2;	 // 1024 ^ 2 is the maximum
+const BUFFER_SIZE: usize = 1024 ^ 2;	 // 1024 ^ 2 is the maximum
 
-fn valid_length(length: u32) -> bool {
+fn valid_length(length: usize) -> bool {
-    return length > 0 && length <= BUFFER_SIZE;
+    length > 0 && length <= BUFFER_SIZE
 }
 
-fn read_header() -> u32 {
+fn read_header() -> usize {
     let stdin = io::stdin();
     let mut buf = vec![0; 4];
     let mut handle = stdin.lock();
 
     handle.read_exact(&mut buf).unwrap();
-    NativeEndian::read_u32(&buf)
+    NativeEndian::read_u32(&buf).try_into().unwrap()
 }
 
-fn read_body<T: Read + Write>(length: u32, socket: &mut ProxySocket<T>) {
+fn read_body<T: Read + Write>(length: usize, socket: &mut ProxySocket<T>) {
-    let mut buffer = vec![0; length as usize];
+    let mut buffer = vec![0; length];
     let stdin = io::stdin();
     let mut handle = stdin.lock();
 
@@ -39,7 +40,7 @@ fn read_body<T: Read + Write>(length: u32, socket: &mut ProxySocket<T>) {
 }
 
 fn read_response<T: Read>(socket: &mut ProxySocket<T>) {
-    let mut buf = vec![0; BUFFER_SIZE as usize];
+    let mut buf = vec![0; BUFFER_SIZE];
     if let Ok(len) = socket.read(&mut buf) {
         write_response(&buf[0..len]);
     }

src/proxy_socket.rs

@@ -4,6 +4,7 @@ use std::io::{self, Read, Write};
 #[cfg(not(windows))]
 use std::os::unix::io::AsRawFd;
 use std::os::unix::net::UnixStream;
+use std::path::PathBuf;
 use nix::sys::socket;
 use nix::sys::socket::sockopt::SndBuf;
 use nix::sys::socket::sockopt::RcvBuf;
@@ -32,7 +33,7 @@ impl<W: Write> Write for ProxySocket<W> {
 }
 
 #[cfg(windows)]
-pub fn connect(buffer_size: u32) -> io::Result<ProxySocket<PipeClient>> {
+pub fn connect(buffer_size: usize) -> io::Result<ProxySocket<PipeClient>> {
     let username = env::var("USERNAME").unwrap();
     let pipe_name = format!("\\\\.\\pipe\\keepassxc\\{}\\org.keepassxc.KeePassXC.BrowserServer", username);
     let client = PipeClient::connect(pipe_name)?;
@@ -40,18 +41,46 @@ pub fn connect(buffer_size: u32) -> io::Result<ProxySocket<PipeClient>> {
 }
 
 #[cfg(not(windows))]
-pub fn connect(buffer_size: u32) -> io::Result<ProxySocket<UnixStream>> {
+/// Returns the directories where the socket could possible be located.
+///
+/// These directories should be tried in sequence, until one of them is found
+/// to contain the socket.
+fn get_socket_dirs() -> Vec<PathBuf> {
+    let mut dirs = Vec::new();
+
+    if !cfg!(target_os = "macos") {
+        if let Ok(dir) = env::var("XDG_RUNTIME_DIR")  {
+            let xdg_runtime_dir: PathBuf = dir.into();
+
+            // Sandbox-friendly path.
+            // Used in KeePassXC >= 2.7.2 and for all versions on Flatpak.
+            dirs.push(xdg_runtime_dir.join("app/org.keepassxc.KeePassXC/"));
+
+            // Legacy path.
+            // Used by KeePassXC < 2.7.2.
+            dirs.push(xdg_runtime_dir);
+        };
+    };
+
+    // Default for macOS, and final fallback for Linux.
+    dirs.push(env::temp_dir());
+
+    dirs
+}
+
+#[cfg(not(windows))]
+pub fn connect(buffer_size: usize) -> io::Result<ProxySocket<UnixStream>> {
     use std::time::Duration;
 
     let socket_name = "org.keepassxc.KeePassXC.BrowserServer";
-    let socket = if let Ok(dir) = if cfg!(target_os = "macos") {env::var("TMPDIR") } else { env::var("XDG_RUNTIME_DIR") } {
+    let dirs = get_socket_dirs();
-        format!("{}/{}", dir, socket_name)
+    let s = dirs
-    } else {
+        .iter()
-        format!("/tmp/{}", socket_name)
+        .find_map(|dir| UnixStream::connect(dir.join(socket_name)).ok())
-    };
+        .ok_or_else(|| io::Error::from(io::ErrorKind::NotFound))?;
-    let s = UnixStream::connect(socket)?;
+
-    socket::setsockopt(s.as_raw_fd(), SndBuf, &(buffer_size as usize)).expect("setsockopt for SndBuf failed");
+    socket::setsockopt(s.as_raw_fd(), SndBuf, &buffer_size).expect("setsockopt for SndBuf failed");
-    socket::setsockopt(s.as_raw_fd(), RcvBuf, &(buffer_size as usize)).expect("setsockopt for RcvBuf failed");
+    socket::setsockopt(s.as_raw_fd(), RcvBuf, &buffer_size).expect("setsockopt for RcvBuf failed");
     let timeout: Option<Duration> = Some(Duration::from_secs(1));
     s.set_read_timeout(timeout)?;
     Ok(ProxySocket { inner: s })