Add message store and refactor main into bridge startup

Step 5: in-memory MessageStore mapping IMAP seq numbers to ProtonMail IDs.
Oldest-first ordering, mark/expunge, range queries, subject/sender search.

Refactor main.rs: extract unlock_key_pool(), build SharedState (Arc<Mutex>),
load initial inbox, then wait for Ctrl-C. Ready to plug in IMAP/SMTP servers.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Cargo.lock

@@ -4015,6 +4015,7 @@ dependencies = [
  "libc",
  "mio",
  "pin-project-lite",
+ "signal-hook-registry",
  "socket2",
  "tokio-macros",
  "windows-sys 0.61.2",

proton-bridge/Cargo.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-tokio = { version = "1", features = ["net", "io-util", "rt-multi-thread", "macros"] }
+tokio = { version = "1", features = ["net", "io-util", "rt-multi-thread", "macros", "signal"] }
 reqwest = { version = "0.12", default-features = false, features = ["json", "native-tls"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"

proton-bridge/src/main.rs

@@ -3,51 +3,108 @@ mod auth;
 mod config;
 mod crypto;
 mod srp;
+mod store;
+
+use std::sync::Arc;
 
 use api::{ApiClient, LABEL_INBOX};
+use store::MessageStore;
+use tokio::sync::Mutex;
+
+// ── Shared bridge state ───────────────────────────────────────────────────────
+
+/// State shared between the IMAP and SMTP server tasks.
+pub struct BridgeState {
+    pub store: MessageStore,
+    pub key_pool: Vec<crypto::PrivateKey>,
+    /// The underlying reqwest client (stored so we can create ApiClient on demand).
+    pub http_client: reqwest::Client,
+    pub session: auth::Session,
+    pub config: config::Config,
+}
+
+pub type SharedState = Arc<Mutex<BridgeState>>;
+
+// ── Entry point ───────────────────────────────────────────────────────────────
 
 #[tokio::main]
 async fn main() {
     env_logger::init();
+
     let config = match config::Config::load() {
         Ok(c) => c,
-        Err(e) => { eprintln!("Failed to load bridge.toml: {}", e); std::process::exit(1); }
+        Err(e) => {
+            eprintln!("Failed to load bridge.toml: {e}");
+            std::process::exit(1);
+        }
     };
 
     let client = match auth::build_client() {
         Ok(c) => c,
-        Err(e) => { eprintln!("Failed to build HTTP client: {}", e); std::process::exit(1); }
+        Err(e) => {
+            eprintln!("Failed to build HTTP client: {e}");
+            std::process::exit(1);
+        }
     };
 
     let session = match auth::authenticate(&client, &config.proton).await {
         Ok(s) => s,
-        Err(e) => { eprintln!("Authentication failed: {}", e); std::process::exit(1); }
+        Err(e) => {
+            eprintln!("Authentication failed: {e}");
+            std::process::exit(1);
+        }
     };
-    println!("Session UID: {}", session.uid);
+    println!("Authenticated as {}", config.proton.username);
 
     let api = ApiClient::new(&client, &session);
 
-    // ── Step 3: list inbox ────────────────────────────────────────────────────
+    // Unlock user + address keys.
-    let (messages, total) = match api.list_messages(LABEL_INBOX, 0, 10).await {
+    let key_pool = match unlock_key_pool(&api, &config).await {
+        Ok(k) => k,
+        Err(e) => {
+            eprintln!("{e}");
+            std::process::exit(1);
+        }
+    };
+    println!("{} key(s) in pool", key_pool.len());
+
+    // Load the initial inbox into the message store.
+    let (messages, total) = match api.list_messages(LABEL_INBOX, 0, 50).await {
         Ok(r) => r,
-        Err(e) => { eprintln!("list_messages failed: {}", e); std::process::exit(1); }
+        Err(e) => {
+            eprintln!("list_messages failed: {e}");
+            std::process::exit(1);
+        }
     };
-    println!("\nInbox ({total} total):");
+    let mut store = MessageStore::new();
-    for m in &messages {
+    store.load_all(messages);
-        println!(
+    println!("Inbox: {} messages ({total} total)", store.count());
-            "  [{}] {:50} — {}",
-            if m.unread == 1 { "UNREAD" } else { "  read" },
-            m.subject,
-            m.sender.address,
-        );
-    }
 
-    // ── Step 4: decrypt first message body ───────────────────────────────────
+    let _state: SharedState = Arc::new(Mutex::new(BridgeState {
-    let first = match messages.first() {
+        store,
-        Some(m) => m,
+        key_pool,
-        None => { println!("\nInbox is empty — nothing to decrypt."); return; }
+        http_client: client,
-    };
+        session,
+        config: config.clone(),
+    }));
 
+    println!(
+        "Bridge ready — IMAP :{}, SMTP :{}  (Ctrl-C to stop)",
+        config.bridge.imap_port, config.bridge.smtp_port
+    );
+
+    tokio::signal::ctrl_c()
+        .await
+        .expect("failed to listen for ctrl-c");
+    println!("Shutting down.");
+}
+
+// ── Key pool setup ────────────────────────────────────────────────────────────
+
+async fn unlock_key_pool(
+    api: &ApiClient<'_>,
+    config: &config::Config,
+) -> Result<Vec<crypto::PrivateKey>, String> {
     // Fetch user keys, address keys, and bcrypt salts in parallel.
     let (user_keys_res, addresses_res, key_salts_res) = tokio::join!(
         api.get_user_keys(),
@@ -55,154 +112,100 @@ async fn main() {
         api.get_key_salts(),
     );
 
-    let user_keys = user_keys_res.unwrap_or_else(|e| { eprintln!("get_user_keys: {e}"); vec![] });
+    let user_keys = user_keys_res.map_err(|e| format!("get_user_keys: {e}"))?;
-    let addresses = addresses_res.unwrap_or_else(|e| { eprintln!("get_addresses: {e}"); vec![] });
+    let addresses = addresses_res.map_err(|e| format!("get_addresses: {e}"))?;
 
-    // key_salts requires the 'locked' OAuth scope.  The server may revoke it even
-    // when our cached session still claims to have it.  If we get 403, purge the
-    // session so the next run does a full SRP+TOTP login that re-grants 'locked'.
     let key_salts = match key_salts_res {
         Ok(salts) => salts,
         Err(e) => {
-            eprintln!("get_key_salts: {e}");
             if e.contains("403") {
                 auth::Session::delete();
-                eprintln!("Session 'locked' scope expired server-side — session.json deleted.");
+                return Err(format!(
-                eprintln!("Re-run; you will be prompted for your TOTP code to regain 'locked' scope.");
+                    "Session 'locked' scope expired — session.json deleted. \
-                std::process::exit(1);
+                     Re-run to enter TOTP and regain 'locked' scope. ({e})"
+                ));
             }
+            eprintln!("get_key_salts: {e} — continuing without salt");
             vec![]
         }
     };
 
-    // In two-password mode the mailbox password unlocks the keys.
+    let password = config
-    let password = config.proton.mailbox_password
+        .proton
+        .mailbox_password
         .as_deref()
         .unwrap_or(&config.proton.password);
 
-    // Derive passphrase for a key ID via bcrypt(password, keySalt).
+    // Derive the passphrase for a key (bcrypt-31 or raw password fallback).
-    // Falls back to raw password when salt is absent.
     let passphrase_for = |key_id: &str| -> String {
-        let salt = key_salts.iter()
+        let salt = key_salts
+            .iter()
             .find(|s| s.id == key_id)
             .and_then(|s| s.key_salt.as_deref())
             .unwrap_or("");
-        crypto::derive_key_passphrase(password, salt)
+        crypto::derive_key_passphrase(password, salt).unwrap_or_else(|_| password.to_string())
-            .unwrap_or_else(|_| password.to_string())
     };
 
-    // Phase 1: unlock user keys — try multiple passphrase candidates so we can
+    // Phase 1: unlock user keys.
-    // identify which derivation is correct.
+    let mut user_private_keys: Vec<crypto::PrivateKey> = Vec::new();
-    let mut user_private_keys: Vec<(String, crypto::PrivateKey)> = Vec::new();
     for uk in &user_keys {
-        if uk.active == 1 {
+        if uk.active != 1 {
-            let salt = key_salts.iter()
+            continue;
-                .find(|s| s.id == uk.id)
+        }
-                .and_then(|s| s.key_salt.as_deref())
+        let salt = key_salts
-                .unwrap_or("");
+            .iter()
-
+            .find(|s| s.id == uk.id)
-            // Derive key passphrase per ProtonMail's go-proton-api SaltForKey:
+            .and_then(|s| s.key_salt.as_deref())
-            //   bcrypt($2y$, cost=10, raw_salt_16_bytes, password+\0)[last 31 chars]
+            .unwrap_or("");
-            eprintln!("  [key-salt] raw API value: {:?} (len={})", salt, salt.len());
+        let pp = crypto::derive_key_passphrase(password, salt)
-
+            .unwrap_or_else(|_| password.to_string());
-            let pp = crypto::derive_key_passphrase(password, salt)
+        match crypto::PrivateKey::unlock(&uk.private_key, &pp) {
-                .unwrap_or_else(|_| password.to_string());
+            Ok(k) => {
-            eprintln!("  [pp] len={} last8=...{}", pp.len(),
+                println!("user key {} unlocked", &uk.id[..8.min(uk.id.len())]);
-                pp.chars().rev().take(8).collect::<String>().chars().rev().collect::<String>());
+                user_private_keys.push(k);
-
-            let mut candidates: Vec<(&str, String)> = vec![
-                ("bcrypt-31",    pp),
-                ("raw-password", password.to_string()),
-            ];
-
-            // Manual override — set user_key_passphrase in bridge.toml to bypass derivation.
-            if let Some(ref override_pp) = config.proton.user_key_passphrase {
-                eprintln!("  [override] using user_key_passphrase from config");
-                candidates.push(("config-override", override_pp.clone()));
-            }
-
-            // Use the first candidate to store the parsed key structure.
-            // (unlock() only parses the key; actual passphrase is tested during decrypt.)
-            let first_pp = candidates[0].1.clone();
-            match crypto::PrivateKey::unlock(&uk.private_key, &first_pp) {
-                Ok(k) => {
-                    println!("user key {} parsed OK", &uk.id[..8]);
-                    // Run the independent S2K+CFB+SHA1 diagnostic for all candidates.
-                    let diag_candidates: Vec<(&str, &str)> = candidates
-                        .iter()
-                        .map(|(l, p)| (*l, p.as_str()))
-                        .collect();
-                    crypto::diagnose_subkey_passphrase(&k, &diag_candidates);
-                    // Store all candidate passphrases so token decrypt can try each.
-                    for (label, pp) in &candidates {
-                        match crypto::PrivateKey::unlock(&uk.private_key, pp) {
-                            Ok(k2) => user_private_keys.push((label.to_string(), k2)),
-                            Err(e) => eprintln!("user key {} [{}] FAILED: {}", &uk.id[..8], label, e),
-                        }
-                    }
-                }
-                Err(e) => eprintln!("user key {} parse FAILED: {}", &uk.id[..8], e),
             }
+            Err(e) => eprintln!("user key {} FAILED: {e}", &uk.id[..8.min(uk.id.len())]),
         }
     }
 
-    // Phase 2: unlock address keys.
+    // Phase 2: unlock address keys via token (newer accounts) or shared passphrase.
-    // Newer accounts: the address key's passphrase is stored as a PGP-encrypted
-    //   Token (encrypted to the user key).  Decrypt the token, then unlock.
-    // Older accounts: address keys share the same bcrypt passphrase as user keys.
     let mut key_pool: Vec<crypto::PrivateKey> = Vec::new();
     for addr in &addresses {
         for ak in &addr.keys {
-            if ak.active == 1 {
+            if ak.active != 1 {
-                let pp = if let Some(token) = &ak.token {
+                continue;
-                    let mut token_pp = None;
+            }
-                    for (label, uk) in &user_private_keys {
+            let pp = if let Some(token) = &ak.token {
-                        match crypto::decrypt_token(token, uk) {
+                let mut token_pp = None;
-                            Ok(p) => {
+                for uk in &user_private_keys {
-                                println!("token decrypted using [{}] passphrase!", label);
+                    if let Ok(p) = crypto::decrypt_token(token, uk) {
-                                token_pp = Some(p);
+                        token_pp = Some(p);
-                                break;
+                        break;
-                            }
-                            Err(e) => eprintln!("decrypt_token [{}] ({}): {e}", label, addr.email),
-                        }
                     }
-                    token_pp.unwrap_or_else(|| passphrase_for(&ak.id))
-                } else {
-                    passphrase_for(&ak.id)
-                };
-                match crypto::PrivateKey::unlock(&ak.private_key, &pp) {
-                    Ok(k) => {
-                        println!("address key {} ({}) unlocked", &ak.id[..8], addr.email);
-                        key_pool.push(k);
-                    }
-                    Err(e) => eprintln!("address key {} ({}): {e}", &ak.id[..8], addr.email),
                 }
+                token_pp.unwrap_or_else(|| passphrase_for(&ak.id))
+            } else {
+                passphrase_for(&ak.id)
+            };
+            match crypto::PrivateKey::unlock(&ak.private_key, &pp) {
+                Ok(k) => {
+                    println!(
+                        "address key {} ({}) unlocked",
+                        &ak.id[..8.min(ak.id.len())],
+                        addr.email
+                    );
+                    key_pool.push(k);
+                }
+                Err(e) => eprintln!(
+                    "address key {} ({}): {e}",
+                    &ak.id[..8.min(ak.id.len())],
+                    addr.email
+                ),
             }
         }
     }
-    // Include user keys as fallback (messages may be encrypted to the user key).
-    key_pool.extend(user_private_keys.into_iter().map(|(_, k)| k));
-    println!("\n{} key(s) in pool ({} address(es))", key_pool.len(), addresses.len());
 
-    // Fetch and decrypt the first message.
+    // Include user keys as fallback.
-    let full_msg = match api.get_message(&first.id).await {
+    key_pool.extend(user_private_keys);
-        Ok(m) => m,
+    Ok(key_pool)
-        Err(e) => { eprintln!("get_message failed: {}", e); std::process::exit(1); }
-    };
-
-    crypto::dump_key_pool(&key_pool);
-    crypto::dump_message_pkesk(&full_msg.body);
-
-    let key_refs: Vec<&crypto::PrivateKey> = key_pool.iter().collect();
-    match crypto::decrypt_body(&full_msg.body, &key_refs) {
-        Ok(plaintext) => {
-            println!("\n── Decrypted body ({}) ──────────────────", full_msg.mime_type);
-            let preview: String = plaintext.chars().take(2000).collect();
-            println!("{preview}");
-            if plaintext.len() > 2000 {
-                println!("… ({} chars total)", plaintext.len());
-            }
-        }
-        Err(e) => eprintln!("decrypt_body failed: {}", e),
-    }
 }

proton-bridge/src/store.rs

@@ -0,0 +1,104 @@
+/// Step 5: In-memory message store.
+///
+/// Maps IMAP sequence numbers (1-based, contiguous) ↔ ProtonMail message IDs.
+/// Messages are stored oldest-first so that seq 1 = oldest, seq N = newest.
+/// The API returns messages newest-first, so `load_all` reverses the slice.
+///
+/// After EXPUNGE the Vec automatically renumbers: the element that was at
+/// index i is now at some new index j, and its seq number is j + 1.
+use std::collections::HashSet;
+
+use crate::api::MessageMeta;
+
+pub struct MessageStore {
+    /// Oldest-first. `messages[i]` has IMAP sequence number `i + 1`.
+    messages: Vec<MessageMeta>,
+    /// Proton IDs of messages marked `\Deleted`, waiting for EXPUNGE.
+    deleted_pending: HashSet<String>,
+}
+
+impl MessageStore {
+    pub fn new() -> Self {
+        Self {
+            messages: Vec::new(),
+            deleted_pending: HashSet::new(),
+        }
+    }
+
+    /// Replace the store contents with a fresh API response (newest-first).
+    /// Reverses to oldest-first and clears any pending deletions.
+    pub fn load_all(&mut self, mut messages: Vec<MessageMeta>) {
+        messages.reverse(); // API is newest-first; store oldest-first
+        self.messages = messages;
+        self.deleted_pending.clear();
+    }
+
+    /// Number of messages currently in the store (the IMAP EXISTS count).
+    pub fn count(&self) -> u32 {
+        self.messages.len() as u32
+    }
+
+    /// Look up the ProtonMail message ID for a given IMAP sequence number.
+    pub fn proton_id(&self, seq: u32) -> Option<&str> {
+        self.messages
+            .get(seq.wrapping_sub(1) as usize)
+            .map(|m| m.id.as_str())
+    }
+
+    /// Look up message metadata for a given IMAP sequence number.
+    pub fn get_meta(&self, seq: u32) -> Option<&MessageMeta> {
+        self.messages.get(seq.wrapping_sub(1) as usize)
+    }
+
+    /// Return all (seq, meta) pairs in the inclusive range [from, to].
+    /// Clamps to the actual message count.
+    pub fn get_range(&self, from: u32, to: u32) -> Vec<(u32, &MessageMeta)> {
+        let max = self.count();
+        let from = from.max(1);
+        let to = to.min(max);
+        if from > to {
+            return vec![];
+        }
+        (from..=to)
+            .filter_map(|seq| self.get_meta(seq).map(|m| (seq, m)))
+            .collect()
+    }
+
+    /// Mark a message for deletion (STORE … +FLAGS (\Deleted)).
+    pub fn mark_deleted(&mut self, seq: u32) {
+        if let Some(m) = self.get_meta(seq) {
+            self.deleted_pending.insert(m.id.clone());
+        }
+    }
+
+    /// Remove all pending-deleted messages and return their ProtonMail IDs.
+    /// After this call the Vec indices automatically provide new seq numbers.
+    pub fn expunge(&mut self) -> Vec<String> {
+        let removed: Vec<String> = self
+            .messages
+            .iter()
+            .filter(|m| self.deleted_pending.contains(&m.id))
+            .map(|m| m.id.clone())
+            .collect();
+        self.messages
+            .retain(|m| !self.deleted_pending.contains(&m.id));
+        self.deleted_pending.clear();
+        removed
+    }
+
+    /// Return the seq numbers of messages whose subject, sender name, or
+    /// sender address contains `query` (case-insensitive substring match).
+    pub fn search_subject_or_from(&self, query: &str) -> Vec<u32> {
+        let q = query.to_lowercase();
+        self.messages
+            .iter()
+            .enumerate()
+            .filter(|(_, m)| {
+                m.subject.to_lowercase().contains(&q)
+                    || m.sender.address.to_lowercase().contains(&q)
+                    || m.sender.name.to_lowercase().contains(&q)
+            })
+            .map(|(i, _)| i as u32 + 1)
+            .collect()
+    }
+}