fix(storage): prevent migration FK-cascade data loss; surface decrypt failures

- Disable foreign-key enforcement for the whole migration run (outside the
  per-migration transaction, per SQLite's documented table-rebuild procedure)
  and run `foreign_key_check` before re-enabling. The `PRAGMA foreign_keys=OFF`
  inside migrations 005/006/012 was a no-op within the runner's transaction, so
  their `DROP TABLE` rebuilds were cascade-deleting oauth_tokens, auth codes,
  and the workspace-binding↔FeatureSet junction. Regression-tested.
- `decrypt_input_values` now distinguishes a legacy plaintext row from a real
  decrypt failure (wrong key / tampered ciphertext) and propagates an error
  instead of silently launching a server with all input values missing.
- Wrap workspace-binding create/update (parent row + junction rewrite) in a
  transaction so a partial write can't orphan the binding.

Signed-off-by: Mohammod Al Amin Ashik <maa.ashik00@gmail.com>

Author: its-mash

crates/mcpmux-storage/src/database.rs

@@ -178,6 +178,23 @@ impl Database {
             MIGRATIONS.last().map(|m| m.version).unwrap_or(0)
         );
 
+        // Disable foreign-key enforcement for the duration of the migration
+        // run, following SQLite's documented "other kinds of table schema
+        // changes" procedure (https://sqlite.org/lang_altertable.html): the
+        // table-rebuild migrations (005/006 rebuild inbound_clients, 012
+        // rebuilds workspace_bindings) `DROP TABLE` the parent, and SQLite's
+        // implicit pre-DROP DELETE FIRES `ON DELETE CASCADE` into child rows
+        // (oauth_tokens, oauth_authorization_codes, the FS junction) — wiping
+        // them. The `PRAGMA foreign_keys=OFF` inside those migration files is
+        // a NO-OP because it runs inside the per-migration transaction below;
+        // the pragma only takes effect OUTSIDE a transaction, so it must be
+        // set here, before any `BEGIN`. The connection-level setting persists
+        // across the per-migration transactions until we restore it.
+        let fk_was_on = self.foreign_keys_enabled();
+        if fk_was_on {
+            self.conn.pragma_update(None, "foreign_keys", "OFF")?;
+        }
+
         // Run all migrations that haven't been applied yet
         for migration in MIGRATIONS {
             if migration.version > current_version {
@@ -196,6 +213,10 @@ impl Database {
                         migration.name,
                         e
                     );
+                    // Best-effort restore of FK enforcement before bailing.
+                    if fk_was_on {
+                        let _ = self.conn.pragma_update(None, "foreign_keys", "ON");
+                    }
                     return Err(anyhow::anyhow!(
                         "Failed to run migration {} ({}): {}",
                         migration.version,
@@ -219,6 +240,45 @@ impl Database {
             }
         }
 
+        // Re-enable FK enforcement and verify the migrated schema is still
+        // referentially consistent (catches any genuine orphan a migration
+        // might have introduced while enforcement was off).
+        if fk_was_on {
+            self.foreign_key_check()?;
+            self.conn.pragma_update(None, "foreign_keys", "ON")?;
+        }
+
+        Ok(())
+    }
+
+    /// Whether foreign-key enforcement is currently enabled on the connection.
+    fn foreign_keys_enabled(&self) -> bool {
+        self.conn
+            .query_row("PRAGMA foreign_keys", [], |row| row.get::<_, i64>(0))
+            .map(|v| v != 0)
+            .unwrap_or(false)
+    }
+
+    /// Run `PRAGMA foreign_key_check` and fail if any violation is reported.
+    /// Called after a migration run completes with FK enforcement temporarily
+    /// disabled, so a buggy migration can't silently leave dangling rows.
+    fn foreign_key_check(&self) -> Result<()> {
+        let mut stmt = self.conn.prepare("PRAGMA foreign_key_check")?;
+        let mut rows = stmt.query([])?;
+        let mut violations: Vec<String> = Vec::new();
+        while let Some(row) = rows.next()? {
+            // Columns: table, rowid, referred_table, fk_index
+            let table: String = row.get(0).unwrap_or_default();
+            let referred: String = row.get(2).unwrap_or_default();
+            violations.push(format!("{table} -> {referred}"));
+        }
+        if !violations.is_empty() {
+            return Err(anyhow::anyhow!(
+                "Post-migration foreign_key_check found {} violation(s): {}",
+                violations.len(),
+                violations.join(", ")
+            ));
+        }
         Ok(())
     }
 
@@ -316,6 +376,140 @@ mod tests {
         assert!(count > 0, "Tables should be created");
     }
 
+    /// After a full migration run on a fresh DB, FK enforcement must be back
+    /// ON and the schema must be referentially consistent — the contract of
+    /// the runner's disable-during-migration / re-enable-after logic.
+    #[test]
+    fn fresh_db_has_fk_enabled_and_no_orphans() {
+        let db = Database::open_in_memory().unwrap();
+        let fk: i64 = db
+            .connection()
+            .query_row("PRAGMA foreign_keys", [], |r| r.get(0))
+            .unwrap();
+        assert_eq!(
+            fk, 1,
+            "runner must re-enable FK enforcement after migrations"
+        );
+        let violations: i64 = db
+            .connection()
+            .query_row("SELECT count(*) FROM pragma_foreign_key_check", [], |r| {
+                r.get(0)
+            })
+            .unwrap();
+        assert_eq!(violations, 0, "migrated schema must have no FK violations");
+    }
+
+    /// The exact mechanic the table-rebuild migrations (005/006/012) depend
+    /// on: DROPping a parent referenced by a child via `ON DELETE CASCADE`
+    /// must NOT wipe the child — which only holds when FK enforcement is OFF
+    /// during the rebuild (the runner disables it outside the transaction;
+    /// the `PRAGMA foreign_keys=OFF` *inside* a migration file is a no-op).
+    #[test]
+    fn dropping_parent_with_fk_off_preserves_cascade_child() {
+        use rusqlite::Connection;
+        let conn = Connection::open_in_memory().unwrap();
+        conn.pragma_update(None, "foreign_keys", "ON").unwrap();
+        conn.execute_batch(
+            "CREATE TABLE parent(id TEXT PRIMARY KEY);
+             CREATE TABLE child(id TEXT PRIMARY KEY,
+                 pid TEXT REFERENCES parent(id) ON DELETE CASCADE);
+             INSERT INTO parent VALUES('p1');
+             INSERT INTO child VALUES('c1','p1');",
+        )
+        .unwrap();
+
+        // Disable FK enforcement the way the runner does, then rebuild parent.
+        conn.pragma_update(None, "foreign_keys", "OFF").unwrap();
+        conn.execute_batch(
+            "CREATE TABLE parent_new(id TEXT PRIMARY KEY);
+             INSERT INTO parent_new SELECT * FROM parent;
+             DROP TABLE parent;
+             ALTER TABLE parent_new RENAME TO parent;",
+        )
+        .unwrap();
+        conn.pragma_update(None, "foreign_keys", "ON").unwrap();
+
+        let n: i64 = conn
+            .query_row("SELECT count(*) FROM child", [], |r| r.get(0))
+            .unwrap();
+        assert_eq!(
+            n, 1,
+            "child must survive a parent table-rebuild when FK enforcement is off"
+        );
+    }
+
+    /// End-to-end regression for the 005/006 data-loss bug: an
+    /// `oauth_tokens` row seeded BEFORE the inbound_clients rebuild
+    /// migrations must survive the upgrade. With FK enforcement left on
+    /// (pre-fix), `DROP TABLE inbound_clients` cascade-deleted every token.
+    #[test]
+    fn migration_upgrade_preserves_oauth_tokens() {
+        use rusqlite::{params, Connection};
+
+        let conn = Connection::open_in_memory().unwrap();
+        conn.pragma_update(None, "foreign_keys", "ON").unwrap();
+        let db = Database { conn };
+        db.ensure_migrations_table().unwrap();
+
+        // Apply migrations up to v4 (last version before the first
+        // inbound_clients rebuild at 005), recording each as applied.
+        const PRE_REBUILD: i64 = 4;
+        for m in MIGRATIONS.iter().filter(|m| m.version <= PRE_REBUILD) {
+            db.conn.execute_batch(m.sql).unwrap();
+            db.conn
+                .execute(
+                    "INSERT OR REPLACE INTO schema_migrations (version, name, applied_at) \
+                     VALUES (?1, ?2, datetime('now'))",
+                    params![m.version, m.name],
+                )
+                .unwrap();
+        }
+
+        // Seed an inbound client + an issued OAuth token (child via CASCADE).
+        db.conn
+            .execute(
+                "INSERT INTO inbound_clients
+                   (client_id, registration_type, client_name, redirect_uris,
+                    grant_types, response_types, token_endpoint_auth_method,
+                    created_at, updated_at)
+                 VALUES ('client-1','dcr','Test','[\"http://127.0.0.1/cb\"]',
+                    '[\"authorization_code\"]','[\"code\"]','none',
+                    datetime('now'), datetime('now'))",
+                [],
+            )
+            .unwrap();
+        db.conn
+            .execute(
+                "INSERT INTO oauth_tokens
+                   (id, client_id, token_type, token_hash, created_at)
+                 VALUES ('tok-1','client-1','access','deadbeef', datetime('now'))",
+                [],
+            )
+            .unwrap();
+
+        // Run the remaining migrations (005..=latest) through the real runner.
+        db.run_migrations().unwrap();
+
+        let tokens: i64 = db
+            .conn
+            .query_row("SELECT count(*) FROM oauth_tokens", [], |r| r.get(0))
+            .unwrap();
+        assert_eq!(
+            tokens, 1,
+            "oauth token must survive the inbound_clients rebuild migrations"
+        );
+        // And the client row itself was preserved through the rebuilds.
+        let clients: i64 = db
+            .conn
+            .query_row(
+                "SELECT count(*) FROM inbound_clients WHERE client_id='client-1'",
+                [],
+                |r| r.get(0),
+            )
+            .unwrap();
+        assert_eq!(clients, 1, "inbound client must survive the rebuilds");
+    }
+
     #[test]
     fn test_persistent_database() {
         let temp_dir = TempDir::new().unwrap();

crates/mcpmux-storage/src/repositories/installed_server_repository.rs

@@ -53,17 +53,37 @@ impl SqliteInstalledServerRepository {
     }
 
     /// Decrypt input values from storage.
-    /// Falls back to plaintext JSON for backward compatibility with unencrypted data.
-    fn decrypt_input_values(&self, stored: Option<String>) -> HashMap<String, String> {
+    ///
+    /// Three cases, kept distinct so a real failure can't masquerade as an
+    /// empty config (which would silently launch a server with all its
+    /// secrets missing):
+    ///   * `None` / empty column → no input values (`Ok(empty)`).
+    ///   * Decrypts cleanly → parse the plaintext JSON (a parse failure here
+    ///     is corruption → error).
+    ///   * Decrypt fails → it may be a legacy *unencrypted* row, so try a
+    ///     plaintext-JSON parse; if THAT also fails the data is neither
+    ///     decryptable nor valid plaintext (wrong master key or tampered
+    ///     ciphertext) → propagate a hard error rather than returning empty.
+    fn decrypt_input_values(&self, stored: Option<String>) -> Result<HashMap<String, String>> {
         let Some(data) = stored else {
-            return HashMap::new();
+            return Ok(HashMap::new());
         };
-        // Try decrypting first (new encrypted format)
+        if data.trim().is_empty() {
+            return Ok(HashMap::new());
+        }
+        // Try decrypting first (new encrypted format).
         if let Ok(json) = self.encryptor.decrypt(&data) {
-            return serde_json::from_str(&json).unwrap_or_default();
+            return serde_json::from_str(&json)
+                .map_err(|e| anyhow::anyhow!("Corrupt decrypted input values (not JSON): {}", e));
         }
-        // Fallback: try parsing as plaintext JSON (backward compat)
-        serde_json::from_str(&data).unwrap_or_default()
+        // Fallback: legacy unencrypted row stored as plaintext JSON.
+        serde_json::from_str(&data).map_err(|e| {
+            anyhow::anyhow!(
+                "Failed to decrypt input values and data is not valid plaintext JSON \
+                 (wrong master key or tampered ciphertext): {}",
+                e
+            )
+        })
     }
 
     /// Parse a datetime string to DateTime<Utc>.
@@ -154,14 +174,17 @@ impl SqliteInstalledServerRepository {
     }
 
     /// Build InstalledServer from extracted row data (needs &self for decryption).
-    fn build_server(&self, row: RawServerRow) -> InstalledServer {
-        InstalledServer {
+    fn build_server(&self, row: RawServerRow) -> Result<InstalledServer> {
+        let input_values = self
+            .decrypt_input_values(row.input_values)
+            .map_err(|e| anyhow::anyhow!("server {}: {}", row.server_id, e))?;
+        Ok(InstalledServer {
             id: Uuid::parse_str(&row.id).unwrap_or_else(|_| Uuid::new_v4()),
             space_id: row.space_id,
             server_id: row.server_id,
             server_name: row.server_name,
             cached_definition: row.cached_definition,
-            input_values: self.decrypt_input_values(row.input_values),
+            input_values,
             enabled: row.enabled,
             env_overrides: Self::parse_json_map(row.env_overrides),
             args_append: Self::parse_json_vec(row.args_append),
@@ -170,7 +193,7 @@ impl SqliteInstalledServerRepository {
             source: Self::parse_source(row.source),
             created_at: Self::parse_datetime(&row.created_at),
             updated_at: Self::parse_datetime(&row.updated_at),
-        }
+        })
     }
 }
 
@@ -189,7 +212,7 @@ impl InstalledServerRepository for SqliteInstalledServerRepository {
             .query_map([], Self::extract_row)?
             .collect::<Result<Vec<_>, _>>()?;
 
-        Ok(rows.into_iter().map(|r| self.build_server(r)).collect())
+        rows.into_iter().map(|r| self.build_server(r)).collect()
     }
 
     async fn list_for_space(&self, space_id: &str) -> Result<Vec<InstalledServer>> {
@@ -205,7 +228,7 @@ impl InstalledServerRepository for SqliteInstalledServerRepository {
             .query_map([space_id], Self::extract_row)?
             .collect::<Result<Vec<_>, _>>()?;
 
-        Ok(rows.into_iter().map(|r| self.build_server(r)).collect())
+        rows.into_iter().map(|r| self.build_server(r)).collect()
     }
 
     async fn list_by_source_file(
@@ -227,7 +250,7 @@ impl InstalledServerRepository for SqliteInstalledServerRepository {
             .query_map([&source_prefix], Self::extract_row)?
             .collect::<Result<Vec<_>, _>>()?;
 
-        Ok(rows.into_iter().map(|r| self.build_server(r)).collect())
+        rows.into_iter().map(|r| self.build_server(r)).collect()
     }
 
     async fn get(&self, id: &Uuid) -> Result<Option<InstalledServer>> {
@@ -243,7 +266,7 @@ impl InstalledServerRepository for SqliteInstalledServerRepository {
             .query_row([id.to_string()], Self::extract_row)
             .optional()?;
 
-        Ok(row.map(|r| self.build_server(r)))
+        row.map(|r| self.build_server(r)).transpose()
     }
 
     async fn get_by_server_id(
@@ -263,7 +286,7 @@ impl InstalledServerRepository for SqliteInstalledServerRepository {
             .query_row([space_id, server_id], Self::extract_row)
             .optional()?;
 
-        Ok(row.map(|r| self.build_server(r)))
+        row.map(|r| self.build_server(r)).transpose()
     }
 
     async fn install(&self, server: &InstalledServer) -> Result<()> {
@@ -350,7 +373,7 @@ impl InstalledServerRepository for SqliteInstalledServerRepository {
             .query_map([space_id], Self::extract_row)?
             .collect::<Result<Vec<_>, _>>()?;
 
-        Ok(rows.into_iter().map(|r| self.build_server(r)).collect())
+        rows.into_iter().map(|r| self.build_server(r)).collect()
     }
 
     async fn list_enabled_all(&self) -> Result<Vec<InstalledServer>> {
@@ -366,7 +389,7 @@ impl InstalledServerRepository for SqliteInstalledServerRepository {
             .query_map([], Self::extract_row)?
             .collect::<Result<Vec<_>, _>>()?;
 
-        Ok(rows.into_iter().map(|r| self.build_server(r)).collect())
+        rows.into_iter().map(|r| self.build_server(r)).collect()
     }
 
     async fn set_enabled(&self, id: &Uuid, enabled: bool) -> Result<()> {