Adjust middleware to modify responses

diff --git a/src/middleware/find_mailboxes_compatibility.rs b/src/middleware/find_mailboxes_compatibility.rs
index 6a9b34120a8e42700ed7bb81d6d39a4b34103750..a1ddda20358d23c444e8db6513fb376058fb9937 100644 (file)
--- a/src/middleware/find_mailboxes_compatibility.rs
+++ b/src/middleware/find_mailboxes_compatibility.rs
@@ -1,16 +1,68 @@
 use log::debug;
 
-/// `MailDrop` can't find folders to sync because it sends FIND MAILBOXES /*
+use super::Middleware;
+
+/// `MailDrop` can't find folders to sync because it implements IMAPv3 and
+/// sends FIND MAILBOXES /*, which does not exist in IMAPv4.
 /// which is not understood by modern servers. It instead replaces it with
 /// a LIST command.
-pub fn middleware(input: &[u8]) -> Vec<u8> {
-    let command = String::from_utf8_lossy(input);
-    if command.contains("FIND MAILBOXES /*") {
-        if let Some(tag) = command.split("FIND MAILBOXES /*").next() {
-            let replacement = format!("{} LIST \"\" \"*\"\r\n", tag.trim());
-            debug!("### {replacement}");
-            return replacement.into_bytes();
+pub struct FindMailboxesCompatibility {
+    tags: Vec<String>,
+}
+
+impl FindMailboxesCompatibility {
+    pub fn new() -> Self {
+        FindMailboxesCompatibility { tags: vec![] }
+    }
+}
+
+impl Middleware for FindMailboxesCompatibility {
+    fn client_message(&mut self, input: &[u8]) -> Vec<u8> {
+        let command = String::from_utf8_lossy(input);
+        if command.contains("FIND MAILBOXES /*") {
+            if let Some(tag) = command.split("FIND MAILBOXES /*").next() {
+                // We'll need to convert the LIST to a FIND
+                self.tags.push(tag.to_string());
+                let replacement = format!("{} LIST \"\" \"*\"\r\n", tag.trim());
+                debug!("### {replacement}");
+                return replacement.into_bytes();
+            }
+        }
+        input.to_vec()
+    }
+
+    fn server_message(&mut self, input: &[u8]) -> Vec<u8> {
+        let command = String::from_utf8_lossy(input);
+        let contains_ok_completed = self
+            .tags
+            .iter()
+            .any(|tag| command.contains(&format!("{} OK Completed", tag)));
+
+        // We want to only modify responses that were a result of a MAILBOX call.
+        if !contains_ok_completed {
+            return input.to_vec();
         }
+
+        let lines: Vec<String> = command
+            .lines()
+            .filter_map(|line| {
+                // The IMAPv3 spec specifically says INBOX is excluded from MAILBOX
+                if line.starts_with("* LIST") && line.trim_end().ends_with("\"/\" INBOX") {
+                    return None;
+                }
+
+                // Transform IMAPv4 "* LIST" lines to IMAPv3 "* MAILBOX"
+                if line.starts_with("* LIST") {
+                    if let Some(last_slash_pos) = line.rfind('/') {
+                        let mailbox_name = line[(last_slash_pos + 1)..].trim();
+                        return Some(format!("* MAILBOX {}\r", mailbox_name));
+                    }
+                }
+
+                return Some(line.to_string());
+            })
+            .collect();
+
+        return lines.join("\n").into_bytes();
     }
-    input.to_vec()
 }

diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs
index da7e64ad7b774e6e6ec7b707c0e6a6d698dd0bb7..b9685bf8c0d6a02778104db76b1b1a6dfdd543cc 100644 (file)
--- a/src/middleware/mod.rs
+++ b/src/middleware/mod.rs
@@ -1,8 +1,16 @@
 mod find_mailboxes_compatibility;
 
-use find_mailboxes_compatibility::middleware as find_mailboxes_compatibility_middleware;
+use find_mailboxes_compatibility::FindMailboxesCompatibility;
 
-type Middleware = fn(&[u8]) -> Vec<u8>;
+use std::sync::{Arc, Mutex};
 
-pub const CLIENT_MIDDLEWARE: [Middleware; 1] = [find_mailboxes_compatibility_middleware];
-pub const SERVER_MIDDLEWARE: [Middleware; 0] = [];
+pub trait Middleware: Sync + Send {
+    fn client_message(&mut self, input: &[u8]) -> Vec<u8>;
+    fn server_message(&mut self, input: &[u8]) -> Vec<u8>;
+}
+
+pub fn get_middleware() -> Arc<Mutex<Vec<Box<dyn Middleware>>>> {
+    Arc::new(Mutex::new(vec![
+        Box::new(FindMailboxesCompatibility::new()),
+    ]))
+}

diff --git a/src/proxy.rs b/src/proxy.rs
index 33de44039615959fdc360b56767d601bbbba58d7..781ce2e1b0dbc9297357cada9b6c4a76bfb4c30e 100644 (file)
--- a/src/proxy.rs
+++ b/src/proxy.rs
@@ -42,7 +42,7 @@ use std::thread::{sleep, spawn, JoinHandle};
 use std::time::Duration;
 
 use crate::configuration::Proxy;
-use crate::middleware::{CLIENT_MIDDLEWARE, SERVER_MIDDLEWARE};
+use crate::middleware::get_middleware;
 
 /// A proxy server that listens for plaintext connections and forwards them
 /// via TLS.
@@ -117,8 +117,8 @@ fn run_proxy(configuration: &Arc<Proxy>, running: &Arc<AtomicBool>) {
 
     while running.load(Ordering::SeqCst) {
         match listener.accept() {
-            Ok((stream, addr)) => {
-                info!("New {} connection from {}", configuration.protocol, addr);
+            Ok((stream, address)) => {
+                info!("New {} connection from {}", configuration.protocol, address);
 
                 let configuration_clone = Arc::clone(configuration);
                 let handle = spawn(move || {
@@ -159,6 +159,9 @@ fn handle_client(client_stream: TcpStream, configuration: &Arc<Proxy>) {
         return;
     }
 
+    let available_middleware = get_middleware();
+    let available_middleware_clone = Arc::clone(&available_middleware);
+
     let connector = match TlsConnector::new() {
         Ok(c) => c,
         Err(e) => {
@@ -167,14 +170,14 @@ fn handle_client(client_stream: TcpStream, configuration: &Arc<Proxy>) {
         }
     };
 
-    let remote_addr = format!(
+    let remote_address = format!(
         "{}:{}",
         configuration.remote_host, configuration.remote_port
     );
-    let tcp_stream = match TcpStream::connect(&remote_addr) {
+    let tcp_stream = match TcpStream::connect(&remote_address) {
         Ok(stream) => stream,
         Err(e) => {
-            error!("Failed to connect to {}: {}", remote_addr, e);
+            error!("Failed to connect to {}: {}", remote_address, e);
             return;
         }
     };
@@ -229,11 +232,13 @@ fn handle_client(client_stream: TcpStream, configuration: &Arc<Proxy>) {
 
             let mut command = buffer[..bytes_read].to_vec();
 
-            for middleware in CLIENT_MIDDLEWARE {
-                command = middleware(&command);
+            if let Ok(mut guard) = available_middleware.lock() {
+                for middleware in guard.iter_mut() {
+                    command = middleware.client_message(&command);
+                }
             }
 
-            let debug_str = String::from_utf8_lossy(&command)
+            let debug_str = String::from_utf8_lossy(&buffer[..bytes_read])
                 .replace('\n', "\\n")
                 .replace('\r', "\\r")
                 .replace('\t', "\\t");
@@ -288,11 +293,13 @@ fn handle_client(client_stream: TcpStream, configuration: &Arc<Proxy>) {
 
             let mut command = buffer[..bytes_read].to_vec();
 
-            for middleware in SERVER_MIDDLEWARE {
-                command = middleware(&command);
+            if let Ok(mut guard) = available_middleware_clone.lock() {
+                for middleware in guard.iter_mut() {
+                    command = middleware.server_message(&command);
+                }
             }
 
-            let debug_str = String::from_utf8_lossy(&command)
+            let debug_str = String::from_utf8_lossy(&buffer[..bytes_read])
                 .replace('\n', "\\n")
                 .replace('\r', "\\r")
                 .replace('\t', "\\t");