Update flows_to primitive to take an EdgeType parameter enum of Data,…

… Ctrl, or DataOrControl (#80)

## What Changed?

See title.

## Why Does It Need To?

More flexible flows_to API

## Checklist

- [x] Above description has been filled out so that upon quash merge we
have a
  good record of what changed.
- [x] New functions, methods, types are documented. Old documentation is
updated
  if necessary
- [x] Documentation in Notion has been updated
- [x] Tests for new behaviors are provided
  - [x] New test suites (if any) ave been added to the CI tests (in
`.github/workflows/rust.yml`) either as compiler test or integration
test.
*Or* justification for their omission from CI has been provided in this
PR
    description.

crates/paralegal-policy/src/context.rs

@@ -29,6 +29,17 @@ pub type FunctionId = DefId;
 type MarkerIndex = HashMap<Marker, Vec<(DefId, MarkerRefinement)>>;
 type FlowsTo = HashMap<ControllerId, CtrlFlowsTo>;
 
+/// Enum for identifying an edge type (data, control or both)
+#[derive(Clone, Copy)]
+pub enum EdgeType {
+    /// Only consider dataflow edges
+    Data,
+    /// Only consider control flow edges
+    Control,
+    /// Consider both types of edges
+    DataAndControl,
+}
+
 /// Interface for defining policies.
 ///
 /// Holds a PDG ([`Self::desc`]) and defines basic queries like
@@ -153,27 +164,48 @@ impl Context {
             .collect()
     }
 
-    /// Returns true if `src` has a data-flow to `sink` in the controller `ctrl_id`
-    pub fn data_flows_to(&self, ctrl_id: ControllerId, src: &DataSource, sink: &DataSink) -> bool {
-        let ctrl_flows = &self.flows_to[&ctrl_id];
-        ctrl_flows
-            .data_flows_to
-            .row_set(&src.to_index(&ctrl_flows.sources))
-            .contains(CallSiteOrDataSink::DataSink(sink.clone()))
-    }
-
-    /// Returns true if `src` has a data+ctrl-flow to `sink` in the controller `ctrl_id`
+    /// Returns whether a node or set of nodes that match some condition flows to a node or set of nodes that match another condition through the configured edge type.
     pub fn flows_to(
         &self,
-        ctrl_id: ControllerId,
+        ctrl_id: Option<ControllerId>,
         src: &DataSource,
         sink: &CallSiteOrDataSink,
+        edge_type: EdgeType,
     ) -> bool {
-        let ctrl_flows = &self.flows_to[&ctrl_id];
-        ctrl_flows
-            .flows_to
-            .row_set(&src.to_index(&ctrl_flows.sources))
-            .contains(sink)
+        let ctrl_flow_ids = match &ctrl_id {
+            Some(id) => vec![id],
+            None => self.flows_to.keys().collect(),
+        };
+
+        match edge_type {
+            EdgeType::Data => ctrl_flow_ids.iter().any(|cf_id| {
+                self.flows_to[cf_id]
+                    .data_flows_to
+                    .row_set(&src.to_index(&self.flows_to[cf_id].sources))
+                    .contains(sink)
+            }),
+            EdgeType::DataAndControl => ctrl_flow_ids.iter().any(|cf_id| {
+                self.flows_to[cf_id]
+                    .flows_to
+                    .row_set(&src.to_index(&self.flows_to[cf_id].sources))
+                    .contains(sink)
+            }),
+            EdgeType::Control => {
+                let cs = match sink {
+                    CallSiteOrDataSink::CallSite(cs) => cs,
+                    CallSiteOrDataSink::DataSink(ds) => match ds {
+                        DataSink::Argument { function, .. } => function,
+                        DataSink::Return => return false,
+                    },
+                };
+                ctrl_flow_ids.iter().any(|cf_id| {
+                    match self.desc.controllers[cf_id].ctrl_flow.get(src) {
+                        Some(callsites) => callsites.iter().contains(cs),
+                        None => false,
+                    }
+                })
+            }
+        }
     }
 
     /// Returns an iterator over all objects marked with `marker`.
@@ -321,17 +353,18 @@ impl Context {
             })
     }
 
-    /// Return an example pair for a data flow from an source from `from` to a sink
+    /// Return an example pair for a flow from an source from `from` to a sink
     /// in `to` if any exist.
-    pub fn any_data_flows<'a>(
+    pub fn any_flows<'a>(
         &self,
-        ctrl_id: ControllerId,
+        ctrl_id: Option<ControllerId>,
         from: &[&'a DataSource],
-        to: &[&'a DataSink],
-    ) -> Option<(&'a DataSource, &'a DataSink)> {
+        to: &[&'a CallSiteOrDataSink],
+        edge_type: EdgeType,
+    ) -> Option<(&'a DataSource, &'a CallSiteOrDataSink)> {
         from.iter().find_map(|&src| {
             to.iter().find_map(|&sink| {
-                self.data_flows_to(ctrl_id, src, sink)
+                self.flows_to(ctrl_id, src, sink, edge_type)
                     .then_some((src, sink))
             })
         })

crates/paralegal-policy/src/flows_to.rs

@@ -125,7 +125,7 @@ impl CtrlFlowsTo {
         for (src, callsites) in &ctrl.ctrl_flow.0 {
             let src = src.to_index(&sources);
             for cs in callsites {
-                let new_call_site = CallSiteOrDataSink::CallSite(cs.clone());
+                let new_call_site: CallSiteOrDataSink = cs.clone().into();
                 flows_to.insert(src, &new_call_site);
                 // initialize with flows from the DataSource to all of the CallSite's DataSinks
                 for sink in cs_to_sink.row_set(&sinks.index(&new_call_site)).iter() {
@@ -172,10 +172,69 @@ fn test_data_flows_to() {
             })
             .unwrap()
     };
-    let sink1 = get_sink("sink1");
-    let sink2 = get_sink("sink2");
-    assert!(ctx.data_flows_to(controller, &src, sink1));
-    assert!(!ctx.data_flows_to(controller, &src, sink2));
+    let sink1: CallSiteOrDataSink = get_sink("sink1").clone().into();
+    let sink2: CallSiteOrDataSink = get_sink("sink2").clone().into();
+    assert!(ctx.flows_to(
+        Some(controller.clone()),
+        &src,
+        &sink1,
+        crate::EdgeType::Data
+    ));
+    assert!(!ctx.flows_to(
+        Some(controller.clone()),
+        &src,
+        &sink2,
+        crate::EdgeType::Data
+    ));
+}
+
+#[test]
+fn test_ctrl_flows_to() {
+    use paralegal_spdg::Identifier;
+    let ctx = crate::test_utils::test_ctx();
+    let controller = ctx.find_by_name("controller_ctrl").unwrap();
+    let src_a = DataSource::Argument(0);
+    let src_b = DataSource::Argument(1);
+    let src_c = DataSource::Argument(2);
+    let get_callsite = |name| {
+        let name = Identifier::new_intern(name);
+        ctx.desc().controllers[&controller]
+            .call_sites()
+            .find(|callsite| ctx.desc().def_info[&callsite.function].name == name)
+            .unwrap()
+    };
+    let cs1: CallSiteOrDataSink = get_callsite("sink1").clone().into();
+    let cs2: CallSiteOrDataSink = get_callsite("sink2").clone().into();
+    assert!(ctx.flows_to(
+        Some(controller.clone()),
+        &src_a,
+        &cs1,
+        crate::EdgeType::Control
+    ));
+    assert!(ctx.flows_to(
+        Some(controller.clone()),
+        &src_c,
+        &cs2,
+        crate::EdgeType::Control
+    ));
+    assert!(ctx.flows_to(
+        Some(controller.clone()),
+        &src_a,
+        &cs2,
+        crate::EdgeType::Control
+    ));
+    assert!(!ctx.flows_to(
+        Some(controller.clone()),
+        &src_b,
+        &cs1,
+        crate::EdgeType::Control
+    ));
+    assert!(!ctx.flows_to(
+        Some(controller.clone()),
+        &src_b,
+        &cs2,
+        crate::EdgeType::Control
+    ));
 }
 
 #[test]
@@ -197,18 +256,34 @@ fn test_flows_to() {
             })
             .unwrap()
     };
-
     let get_callsite = |name| {
         let name = Identifier::new_intern(name);
         ctx.desc().controllers[&controller]
             .call_sites()
             .find(|callsite| ctx.desc().def_info[&callsite.function].name == name)
             .unwrap()
     };
-    let sink = CallSiteOrDataSink::DataSink(get_datasink("sink1").clone());
-    let cs = CallSiteOrDataSink::CallSite(get_callsite("sink1").clone());
+    let sink: CallSiteOrDataSink = get_datasink("sink1").clone().into();
+    let cs: CallSiteOrDataSink = get_callsite("sink1").clone().into();
     // a flows to the sink1 callsite (by ctrl flow)
-    assert!(ctx.flows_to(controller, &src_a, &cs));
+    assert!(ctx.flows_to(
+        Some(controller.clone()),
+        &src_a,
+        &cs,
+        crate::EdgeType::DataAndControl
+    ));
+    assert!(!ctx.flows_to(Some(controller.clone()), &src_a, &cs, crate::EdgeType::Data));
     // b flows to the sink1 datasink (by data flow)
-    assert!(ctx.flows_to(controller, &src_b, &sink));
+    assert!(ctx.flows_to(
+        Some(controller.clone()),
+        &src_b,
+        &sink,
+        crate::EdgeType::DataAndControl
+    ));
+    assert!(ctx.flows_to(
+        Some(controller.clone()),
+        &src_b,
+        &sink,
+        crate::EdgeType::Data
+    ));
 }

crates/paralegal-policy/tests/test-crate/src/lib.rs

@@ -28,6 +28,16 @@ fn controller_data_ctrl(a: Foo, b: Foo) {
     }
 }
 
+#[paralegal_flow::analyze]
+fn controller_ctrl(a: bool, b: Foo, c: bool, d: Foo) {
+    if a {
+        sink1(b);
+        if c {
+            sink2(d);
+        }
+    }
+}
+
 #[paralegal_flow::marker(start, return)]
 fn create() -> u32 {
     9

crates/paralegal-spdg/src/lib.rs

@@ -459,6 +459,34 @@ define_index_type! {
     pub struct CallSiteOrDataSinkIndex for CallSiteOrDataSink = u32;
 }
 
+impl From<CallSite> for CallSiteOrDataSink {
+    fn from(cs: CallSite) -> Self {
+        CallSiteOrDataSink::CallSite(cs)
+    }
+}
+
+impl From<DataSink> for CallSiteOrDataSink {
+    fn from(ds: DataSink) -> Self {
+        CallSiteOrDataSink::DataSink(ds)
+    }
+}
+
+impl CallSiteOrDataSink {
+    pub fn as_call_site(&self) -> Option<&CallSite> {
+        match self {
+            Self::CallSite(cs) => Some(cs),
+            _ => None,
+        }
+    }
+
+    pub fn as_data_sink(&self) -> Option<&DataSink> {
+        match self {
+            Self::DataSink(ds) => Some(ds),
+            _ => None,
+        }
+    }
+}
+
 /// Create a hash for this object that is no longer than six hex digits
 ///
 /// The intent for this is to be used as a pre- or postfix to make a non-unique
@@ -695,21 +723,18 @@ impl Ctrl {
             .0
             .values()
             .flatten()
-            .flat_map(|s| {
-                let as_or = CallSiteOrDataSink::DataSink(s.clone());
-                match s {
-                    DataSink::Argument { function, .. } => {
-                        vec![CallSiteOrDataSink::CallSite(function.clone()), as_or]
-                    }
-                    _ => vec![as_or],
+            .flat_map(|s| match s {
+                DataSink::Argument { function, .. } => {
+                    vec![function.clone().into(), s.clone().into()]
                 }
+                _ => vec![s.clone().into()],
             })
             .chain(
                 self.ctrl_flow
                     .0
                     .values()
                     .flatten()
-                    .map(|cs| CallSiteOrDataSink::CallSite(cs.clone())),
+                    .map(|cs| cs.clone().into()),
             )
             .dedup()
     }

guide/deletion-policy/src/main.rs

@@ -28,11 +28,18 @@ fn deletion_policy(ctx: Arc<Context>) -> Result<()> {
     let found = ctx.all_controllers().any(|(deleter_id, deleter)| {
         let delete_sinks = ctx
             .marked_sinks(deleter.data_sinks(), Marker::new_intern("deletes"))
+            .map(|s| s.clone().into())
             .collect::<Vec<_>>();
+        let delete_sinks_borrowed = delete_sinks.iter().collect::<Vec<_>>();
         user_data_types.iter().all(|&t| {
             let sources = ctx.srcs_with_type(deleter, t).collect::<Vec<_>>();
-            ctx.any_data_flows(deleter_id, &sources, &delete_sinks)
-                .is_some()
+            ctx.any_flows(
+                Some(deleter_id),
+                &sources,
+                &delete_sinks_borrowed,
+                paralegal_policy::EdgeType::Data,
+            )
+            .is_some()
         })
     });
     assert_error!(ctx, found, "Could not find a controller deleting all types");

props/plume/src/main.rs

@@ -55,11 +55,18 @@ fn check(ctx: Arc<Context>) -> Result<()> {
             .find(|(deleter_id, deleter)| {
                 let delete_sinks = ctx
                     .marked_sinks(deleter.data_sinks(), marker!(to_delete))
+                    .map(|s| s.clone().into())
                     .collect::<Vec<_>>();
+                let delete_sinks_borrowed = delete_sinks.iter().collect::<Vec<_>>();
                 user_data_types.iter().all(|&t| {
                     let sources = ctx.srcs_with_type(deleter, t).collect::<Vec<_>>();
-                    ctx.any_data_flows(*deleter_id, &sources, &delete_sinks)
-                        .is_some()
+                    ctx.any_flows(
+                        Some(*deleter_id),
+                        &sources,
+                        &delete_sinks_borrowed,
+                        paralegal_policy::EdgeType::Data,
+                    )
+                    .is_some()
                 })
             });
     if found.is_none() {

props/websubmit/src/main.rs

@@ -25,8 +25,13 @@ impl DeletionProp {
                 .collect::<Vec<_>>();
 
             for t_src in t_srcs {
-                for store in &store_cs {
-                    if self.cx.data_flows_to(*c_id, t_src, store) {
+                for &store in &store_cs {
+                    if self.cx.flows_to(
+                        Some(*c_id),
+                        t_src,
+                        &store.clone().into(),
+                        paralegal_policy::EdgeType::Data,
+                    ) {
                         return true;
                     }
                 }
@@ -51,8 +56,13 @@ impl DeletionProp {
                     .collect::<Vec<_>>();
 
                 for t_src in &t_srcs {
-                    for delete in &delete_cs {
-                        if self.cx.data_flows_to(*c_id, t_src, delete) {
+                    for &delete in &delete_cs {
+                        if self.cx.flows_to(
+                            Some(*c_id),
+                            t_src,
+                            &delete.clone().into(),
+                            paralegal_policy::EdgeType::Data,
+                        ) {
                             return true;
                         }
                     }