Relay all messages regardless of order

The existing implementation expects alternating messages from each
source. If either side delivered two consecutive messages, these were
not handled properly and communication ended up out-of-sync.

This new approach runs two threads, each one handling data in one
direction.

I replaced the BUFFER_SIZE because the caret (^) symbol means bitwise or
in rust, so:

    1024 ^ 2 = 1026

The read timeout has been removed, since we now actually WANT to block
indefinitely if there is no traffic in a given direction.

Fixes: #13
Fixes: #14

Author: WhyNotHugo

src/main.rs

@@ -12,75 +12,54 @@ mod proxy_socket;
 
 use proxy_socket::ProxySocket;
 
-const BUFFER_SIZE: usize = 1024 ^ 2; // 1024 ^ 2 is the maximum
-
-fn valid_length(length: usize) -> bool {
-    length > 0 && length <= BUFFER_SIZE
-}
-
-// Read a header (message size) from stdin (e.g.: from the browser).
-fn read_header() -> Result<usize> {
-    let stdin = stdin();
-    let mut buf = vec![0; 4];
-    let mut handle = stdin.lock();
-
-    handle.read_exact(&mut buf)?;
+// > 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;
 
-    NativeEndian::read_u32(&buf)
-        .try_into()
-        .map_err(|err| Error::new(ErrorKind::InvalidData, err))
-}
+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 a whole request/response cycle
-//
-// Read a message body from stdin (e.g.: from the browser), and echo it back to the browser's
-// socket. Then await a response from the socket and relay that back to the browser.
-fn read_body<T: Read + Write>(length: usize, socket: &mut ProxySocket<T>) -> Result<()> {
-    let mut buffer = vec![0; length];
-    let stdin = stdin();
-    let mut handle = stdin.lock();
+    loop {
+        handle.read_exact(&mut len)?;
+        let length: usize = NativeEndian::read_u32(&len)
+            .try_into()
+            .map_err(|err| Error::new(ErrorKind::InvalidData, err))?;
 
-    handle.read_exact(&mut buffer)?;
+        let mut buffer = vec![0; length];
+        handle.read_exact(&mut buffer)?;
 
-    if valid_length(length) {
         socket.write_all(&buffer)?;
         socket.flush()?;
-        read_response(socket)?;
     }
-
-    Ok(())
 }
 
-// Read a response (from KP's socket) and echo it back to the browser.
-fn read_response<T: Read>(socket: &mut ProxySocket<T>) -> Result<()>{
-    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 when the socket is closed.
+fn socket_to_stdout<T: Read + Write>(socket: &mut ProxySocket<T>) -> Result<()> {
+    let mut out = stdout().lock();
+    let mut buf = [0; BUFFER_SIZE];
+
+    while 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()?;
+        }
     }
 
     Ok(())
 }
 
-// Write a response to stdout (e.g.: to the browser).
-fn write_response(buf: &[u8]) -> Result<()> {
-    let stdout = stdout();
-    let mut out = stdout.lock();
+fn main() -> Result<()> {
+    let mut socket = proxy_socket::connect(BUFFER_SIZE)?;
+    let mut socket_clone = socket.try_clone()?;
 
-    out.write_u32::<NativeEndian>(buf.len() as u32)?;
-    out.write_all(buf)?;
-    out.flush()?;
+    thread::spawn(move || stdin_to_socket(&mut socket));
+    socket_to_stdout(&mut socket_clone)?;
 
     Ok(())
 }
-
-fn main() {
-    let mut socket = proxy_socket::connect(BUFFER_SIZE).unwrap();
-
-    // Start thread for user input reading
-    let ui = thread::spawn(move || loop {
-        let length = read_header().unwrap();
-        read_body(length, &mut socket).unwrap();
-    });
-
-    let _ui_res = ui.join().unwrap();
-}

src/proxy_socket.rs

@@ -1,13 +1,13 @@
 use std::env;
 use std::io::{self, Read, Write};
 
+use nix::sys::socket;
+use nix::sys::socket::sockopt::RcvBuf;
+use nix::sys::socket::sockopt::SndBuf;
 #[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;
 
 #[cfg(windows)]
 use named_pipe::PipeClient;
@@ -16,6 +16,15 @@ pub struct ProxySocket<T> {
     inner: T,
 }
 
+impl ProxySocket<UnixStream> {
+    pub(crate) fn try_clone(&self) -> io::Result<Self> {
+        let inner = self.inner.try_clone()?;
+        Ok(Self { inner })
+    }
+}
+
+// TODO: the window Pipeclient
+
 impl<R: Read> Read for ProxySocket<R> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
         self.inner.read(buf)
@@ -35,7 +44,10 @@ impl<W: Write> Write for ProxySocket<W> {
 #[cfg(windows)]
 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 })
 }
@@ -70,8 +82,6 @@ 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
@@ -82,7 +92,8 @@ pub fn connect(buffer_size: usize) -> io::Result<ProxySocket<UnixStream>> {
     socket::setsockopt(s.as_raw_fd(), SndBuf, &buffer_size)?;
     socket::setsockopt(s.as_raw_fd(), RcvBuf, &buffer_size)?;
 
-    let timeout: Option<Duration> = Some(Duration::from_secs(1));
-    s.set_read_timeout(timeout)?;
+    // Make sure reads are blocking.
+    s.set_nonblocking(false)?;
+
     Ok(ProxySocket { inner: s })
 }