Add test to ensure agent fails on dirty tables after first connection

Author: gememma

Diff for: mirrord/agent/iptables/src/lib.rs

@@ -214,7 +214,7 @@ where
     /// IP tables chain names, and they are in danger of conflicting with each-other
     ///
     /// returns `true` if iptables are clean and the agent can proceed, other returns `false`
-    pub async fn ensure_iptables_clean(ipt: IPT) -> IPTablesResult<bool> {
+    pub async fn ensure_iptables_clean(ipt: &IPT) -> IPTablesResult<bool> {
         let ipt = Arc::new(ipt);
         tracing::trace!("checking for leftover or existing iptable chains...");
         for name in [IPTABLE_PREROUTING, IPTABLE_MESH, IPTABLE_STANDARD] {
@@ -524,4 +524,72 @@ mod tests {
 
         assert!(ipt.cleanup().await.is_ok());
     }
+
+    /// Check that dirty ip tables fail the ['SafeIpTables::ensure_iptables_clean'] check
+    /// If there are any chains in the IP table with names used by the agent, the check should fail.
+    /// A fresh IP table, or one with only non-agent names, should pass the check
+    #[tokio::test]
+    async fn fail_on_dirty() {
+        let mut mock = MockIPTables::new();
+
+        // CASE 1: check that a fresh IP table passes cleanliness check
+        mock.expect_list_rules()
+            .with(eq("OUTPUT"))
+            .returning(|_| Ok(vec![]));
+
+        let expected_success = SafeIpTables::ensure_iptables_clean(&mock).await;
+        assert!(
+            expected_success.is_ok(),
+            "Fresh IP table should pass cleanliness check"
+        );
+
+        // CASE 2: check that an IP table with one of mirrord's static chain names will fail the
+        // cleanliness check
+        mock.expect_create_chain()
+            .with(eq(IPTABLE_PREROUTING))
+            .times(1)
+            .returning(|_| Ok(()));
+
+        mock.expect_insert_rule()
+            .with(
+                eq(IPTABLE_PREROUTING),
+                eq("-m tcp -p tcp --dport 69 -j REDIRECT --to-ports 420"),
+                eq(1),
+            )
+            .times(1)
+            .returning(|_, _, _| Ok(()));
+
+        let expected_failure = SafeIpTables::ensure_iptables_clean(&mock).await;
+        assert!(expected_failure.is_err(), "IP table that contains chain with name {IPTABLE_PREROUTING} should fail cleanliness check");
+
+        // CASE 3: check that an IP table with only rules that were not created by mirrord will pass
+        // the cleanliness check
+        mock.expect_remove_chain()
+            .with(eq(IPTABLE_PREROUTING))
+            .times(1)
+            .returning(|_| Ok(()));
+
+        // add a non-mirrord chain
+        let non_mirrord_chain = "MY_PET_CHAIN_NAME";
+
+        mock.expect_create_chain()
+            .with(eq(non_mirrord_chain))
+            .times(1)
+            .returning(|_| Ok(()));
+
+        mock.expect_insert_rule()
+            .with(
+                eq(non_mirrord_chain),
+                eq("-m tcp -p tcp --dport 69 -j REDIRECT --to-ports 420"),
+                eq(1),
+            )
+            .times(1)
+            .returning(|_, _, _| Ok(()));
+
+        let expected_success = SafeIpTables::ensure_iptables_clean(&mock).await;
+        assert!(
+            expected_success.is_ok(),
+            "IP table that contains no chains with mirrord name should pass cleanliness check"
+        );
+    }
 }

Diff for: mirrord/agent/src/entrypoint.rs

@@ -587,7 +587,7 @@ async fn start_agent(args: Args) -> AgentResult<()> {
             } else {
                 new_iptables()
             };
-            SafeIpTables::ensure_iptables_clean(IPTablesWrapper::from(iptables))
+            SafeIpTables::ensure_iptables_clean(&IPTablesWrapper::from(iptables))
                 .await
                 .map_err(|error| AgentError::IPTablesSetupError(error.into()))
         })