Relay all messages regardless of order

src/main.rs

@@ -1,70 +1,75 @@
 extern crate byteorder;
-extern crate nix;
 #[cfg(windows)]
 extern crate named_pipe;
+extern crate nix;
 
+use byteorder::{ByteOrder, NativeEndian, WriteBytesExt};
 use std::convert::TryInto;
-use std::io::{self, Read, Write};
+use std::io::{stdin, stdout, Error, ErrorKind, Read, Result, Write};
 use std::thread;
-use byteorder::{ByteOrder, NativeEndian, WriteBytesExt};
 
 mod proxy_socket;
 
 use proxy_socket::ProxySocket;
 
-const BUFFER_SIZE: usize = 1024 ^ 2;	 // 1024 ^ 2 is the maximum
+// > The maximum size of a single message from the application is 1 MB.
+//
+// From: https://developer.mozilla.org/en-US/docs/Mozilla/Add-ons/WebExtensions/Native_messaging#app_side
+const BUFFER_SIZE: usize = 1024 * 1024;
 
-fn valid_length(length: usize) -> bool {
-    length > 0 && length <= BUFFER_SIZE
-}
-
-fn read_header() -> usize {
-    let stdin = io::stdin();
-    let mut buf = vec![0; 4];
-    let mut handle = stdin.lock();
+/// Reads from stdin and writes to the socket.
+/// Returns on error.
+fn stdin_to_socket<T: Read + Write>(socket: &mut ProxySocket<T>) -> Result<()> {
+    let mut handle = stdin().lock();
+    let mut len = vec![0; std::mem::size_of::<u32>()];
 
-    handle.read_exact(&mut buf).unwrap();
-    NativeEndian::read_u32(&buf).try_into().unwrap()
-}
+    loop {
+        handle.read_exact(&mut len)?;
+        let length: usize = NativeEndian::read_u32(&len)
+            .try_into()
+            .map_err(|err| Error::new(ErrorKind::InvalidData, err))?;
 
-fn read_body<T: Read + Write>(length: usize, socket: &mut ProxySocket<T>) {
-    let mut buffer = vec![0; length];
-    let stdin = io::stdin();
-    let mut handle = stdin.lock();
+        let mut buffer = vec![0; length];
+        handle.read_exact(&mut buffer)?;
 
-    if handle.read_exact(&mut buffer).is_ok() && valid_length(length) {
-        socket.write_all(&buffer).unwrap();
-        socket.flush().unwrap();
-        read_response(socket);
+        socket.write_all(&buffer)?;
+        socket.flush()?;
     }
 }
 
-fn read_response<T: Read>(socket: &mut ProxySocket<T>) {
-    let mut buf = vec![0; BUFFER_SIZE];
-    if let Ok(len) = socket.read(&mut buf) {
-        write_response(&buf[0..len]);
-    }
-}
+/// Reads from the socket and writes to stdout.
+/// Returns on error.
+fn socket_to_stdout<T: Read + Write>(socket: &mut ProxySocket<T>) -> Result<()> {
+    let mut out = stdout().lock();
+    let mut buf = [0; BUFFER_SIZE];
 
-fn write_response(buf: &[u8]) {
-    let stdout = io::stdout();
-    let mut out = stdout.lock();
+    loop {
+        if let Ok(len) = socket.read(&mut buf) {
+            // If a message is larger than the maximum, ignore it entirely. These are disallowed
+            // by the browser anyway, so sending one would be a protocol violation.
+            if len <= BUFFER_SIZE {
+                out.write_u32::<NativeEndian>(len as u32)?;
+                out.write_all(&buf[..len])?;
+                out.flush()?;
+            };
+        } else {
+            // TOOD: is the socket is closed, we should try to reconnect.
 
-    out.write_u32::<NativeEndian>(buf.len() as u32).unwrap();
-    out.write_all(buf).unwrap();
-    out.flush().unwrap();
+            return Err(Error::from(ErrorKind::BrokenPipe));
+        }
+    }
 }
 
-fn main() {
-    let mut socket = proxy_socket::connect(BUFFER_SIZE).unwrap();
+fn main() -> Result<()> {
+    let mut socket = proxy_socket::connect(BUFFER_SIZE)?;
+    let mut socket_clone = socket.try_clone()?;
 
-    // Start thread for user input reading
-    let ui = thread::spawn(move || {
-        loop {
-            let length = read_header();
-            read_body(length, &mut socket);
-        }
-    });
+    thread::spawn(move || socket_to_stdout(&mut socket_clone).unwrap());
+
+    // If stdin is closed, that means that Firefox has exited, so we exit too.
+    // If the socket is closed, this will (eventually) fail too, however, this can later be
+    // refactored to reconnect the underlying ProxySocket.
+    stdin_to_socket(&mut socket).unwrap();
 
-    let _ui_res = ui.join().unwrap();
+    Ok(())
 }

src/proxy_socket.rs

@@ -2,12 +2,14 @@ use std::env;
 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;
+use {
+    nix::sys::socket::{
+        setsockopt,
+        sockopt::{RcvBuf, SndBuf},
+    },
+    std::os::unix::{io::AsRawFd, net::UnixStream},
+    std::path::PathBuf,
+};
 
 #[cfg(windows)]
 use named_pipe::PipeClient;
@@ -16,6 +18,20 @@ pub struct ProxySocket<T> {
     inner: T,
 }
 
+#[cfg(not(windows))]
+impl ProxySocket<UnixStream> {
+    pub(crate) fn try_clone(&self) -> io::Result<Self> {
+        let inner = self.inner.try_clone()?;
+        Ok(Self { inner })
+    }
+}
+#[cfg(windows)]
+impl ProxySocket<PipeClient> {
+    pub(crate) fn try_clone(&self) -> io::Result<Self> {
+        todo!();
+    }
+}
+
 impl<R: Read> Read for ProxySocket<R> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
         self.inner.read(buf)
@@ -33,9 +49,12 @@ impl<W: Write> Write for ProxySocket<W> {
 }
 
 #[cfg(windows)]
-pub fn connect(buffer_size: usize) -> 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 pipe_name = format!(
+        "\\\\.\\pipe\\keepassxc\\{}\\org.keepassxc.KeePassXC.BrowserServer",
+        username
+    );
     let client = PipeClient::connect(pipe_name)?;
     Ok(ProxySocket { inner: client })
 }
@@ -49,7 +68,7 @@ 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")  {
+        if let Ok(dir) = env::var("XDG_RUNTIME_DIR") {
             let xdg_runtime_dir: PathBuf = dir.into();
 
             // Sandbox-friendly path.
@@ -70,18 +89,18 @@ fn get_socket_dirs() -> Vec<PathBuf> {
 
 #[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 dirs = get_socket_dirs();
     let s = dirs
         .iter()
         .find_map(|dir| UnixStream::connect(dir.join(socket_name)).ok())
         .ok_or_else(|| io::Error::from(io::ErrorKind::NotFound))?;
 
-    socket::setsockopt(s.as_raw_fd(), SndBuf, &buffer_size).expect("setsockopt for SndBuf 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)?;
+    setsockopt(s.as_raw_fd(), SndBuf, &buffer_size)?;
+    setsockopt(s.as_raw_fd(), RcvBuf, &buffer_size)?;
+
+    // Make sure reads are blocking.
+    s.set_nonblocking(false)?;
+
     Ok(ProxySocket { inner: s })
 }