From fa1ece4a22104583a123ae996974ebe74365b70b Mon Sep 17 00:00:00 2001
From: lzoghbi <zoghbi.lydia@gmail.com>
Date: Sun, 4 Aug 2024 04:55:07 -0700
Subject: [PATCH] Extension: keep track of safe pub fns when recalculating pub
 caller-checked to update chain

---
 lang_server/src/server.rs |  1 +
 lang_server/src/util.rs   |  5 +++++
 src/audit_chain.rs        | 14 +++++++++++++-
 src/audit_file.rs         | 20 ++++++++++++++++++++
 4 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/lang_server/src/server.rs b/lang_server/src/server.rs
index 2106b6f..1bf559d 100644
--- a/lang_server/src/server.rs
+++ b/lang_server/src/server.rs
@@ -163,6 +163,7 @@ fn runner(
                                     curr_effect,
                                 );
                             }
+                            af.recalc_pub_caller_checked(&scan_res.pub_fns);
                             af.save_to_file(audit_file_path.clone())?;
                         }
 
diff --git a/lang_server/src/util.rs b/lang_server/src/util.rs
index 3aada49..07ae3d5 100644
--- a/lang_server/src/util.rs
+++ b/lang_server/src/util.rs
@@ -64,6 +64,11 @@ pub fn get_all_chain_effects(
             anyhow!("Couldn't find audit chain manifest at {}", chain_manifest.display())
         })?;
 
+    // Check if any dependency sinks were removed by other audits
+    // and update chain before loading all existing effects
+    let removed_sinks = chain.collect_all_safe_sinks()?;
+    chain.remove_cross_crate_effects(removed_sinks, &chain.root_crate()?)?;
+
     for crate_id in chain.to_owned().all_crates() {
         if let Some(af) = chain.read_audit_file(crate_id)? {
             for (effect_instance, audit_tree) in &af.audit_trees {
diff --git a/src/audit_chain.rs b/src/audit_chain.rs
index 277cf8d..ee9864d 100644
--- a/src/audit_chain.rs
+++ b/src/audit_chain.rs
@@ -121,6 +121,18 @@ impl AuditChain {
         }
     }
 
+    pub fn collect_all_safe_sinks(&mut self) -> Result<HashSet<CanonicalPath>> {
+        let mut safe_sinks = HashSet::new();
+        for (crate_id, (af_path, _)) in &self.crate_policies {
+            let audit_file = AuditFile::read_audit_file(af_path.clone())?.context(
+                format!("Can't find an associated audit for crate `{}`", crate_id),
+            )?;
+            safe_sinks.extend(audit_file.safe_pub_fns());
+        }
+
+        Ok(safe_sinks)
+    }
+
     /// Returns the full package name with version if there is exactly one
     /// package matching the input, or none otherwise
     pub fn resolve_crate_id(&self, crate_name: &str) -> Option<CrateId> {
@@ -179,7 +191,7 @@ impl AuditChain {
         if let Ok(l) = Lockfile::load(&crate_path) {
             Ok(l)
         } else {
-            println!("Lockfile missing: generating new lockfile");
+            info!("Lockfile missing: generating new lockfile");
             let config = GlobalContext::default()?;
             crate_path.pop();
             crate_path.push("Cargo.toml");
diff --git a/src/audit_file.rs b/src/audit_file.rs
index 5482511..27294fd 100644
--- a/src/audit_file.rs
+++ b/src/audit_file.rs
@@ -361,6 +361,26 @@ impl AuditFile {
             pub_caller_checked.into_iter().filter(|(_, v)| !v.is_empty()).collect();
     }
 
+    /// Mirror of `recalc_pub_caller_checked`, but retains safe public functions.
+    /// Useful when auditing through the IDE extension, where we use 
+    /// the same audit file both when auditing chain and a single crate.
+    pub fn recalc_pub_cc_with_safe(&mut self, pub_fns: &HashSet<CanonicalPath>) {
+        // NOTE: initialize everything at the start so we don't have to check for
+        //       entries and clone keys every time
+        let mut pub_caller_checked =
+            HashMap::from_iter(pub_fns.iter().map(|p| (p.clone(), HashSet::new())));
+        for (effect, tree) in self.audit_trees.iter() {
+            AuditFile::recalc_pub_caller_checked_tree(
+                effect,
+                tree,
+                &mut pub_caller_checked,
+                pub_fns,
+            );
+        }
+
+        self.pub_caller_checked = pub_caller_checked;
+    }
+
     /// Returns the list of all safe public functions (these include all the
     /// public functions which have been removed since the last audit update).
     pub fn safe_pub_fns(&self) -> HashSet<CanonicalPath> {