diff --git a/crates/paralegal-policy/src/context.rs b/crates/paralegal-policy/src/context.rs index 6d57631307..d207c2caf5 100644 --- a/crates/paralegal-policy/src/context.rs +++ b/crates/paralegal-policy/src/context.rs @@ -29,6 +29,17 @@ pub type FunctionId = DefId; type MarkerIndex = HashMap>; type FlowsTo = HashMap; +/// 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, 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, 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)) }) }) diff --git a/crates/paralegal-policy/src/flows_to.rs b/crates/paralegal-policy/src/flows_to.rs index cb6df7c199..2c0a53c532 100644 --- a/crates/paralegal-policy/src/flows_to.rs +++ b/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,7 +256,6 @@ fn test_flows_to() { }) .unwrap() }; - let get_callsite = |name| { let name = Identifier::new_intern(name); ctx.desc().controllers[&controller] @@ -205,10 +263,27 @@ fn test_flows_to() { .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 + )); } diff --git a/crates/paralegal-policy/tests/test-crate/src/lib.rs b/crates/paralegal-policy/tests/test-crate/src/lib.rs index 65787c5ba4..ffabe2636a 100644 --- a/crates/paralegal-policy/tests/test-crate/src/lib.rs +++ b/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 diff --git a/crates/paralegal-spdg/src/lib.rs b/crates/paralegal-spdg/src/lib.rs index bcb9a4dbd5..10e807994d 100644 --- a/crates/paralegal-spdg/src/lib.rs +++ b/crates/paralegal-spdg/src/lib.rs @@ -459,6 +459,34 @@ define_index_type! { pub struct CallSiteOrDataSinkIndex for CallSiteOrDataSink = u32; } +impl From for CallSiteOrDataSink { + fn from(cs: CallSite) -> Self { + CallSiteOrDataSink::CallSite(cs) + } +} + +impl From 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() } diff --git a/guide/deletion-policy/src/main.rs b/guide/deletion-policy/src/main.rs index 3114827c00..365021d8ee 100644 --- a/guide/deletion-policy/src/main.rs +++ b/guide/deletion-policy/src/main.rs @@ -28,11 +28,18 @@ fn deletion_policy(ctx: Arc) -> 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::>(); + let delete_sinks_borrowed = delete_sinks.iter().collect::>(); user_data_types.iter().all(|&t| { let sources = ctx.srcs_with_type(deleter, t).collect::>(); - 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"); diff --git a/props/plume/src/main.rs b/props/plume/src/main.rs index 6b40e348a8..ed01f52ced 100644 --- a/props/plume/src/main.rs +++ b/props/plume/src/main.rs @@ -55,11 +55,18 @@ fn check(ctx: Arc) -> Result<()> { .find(|(deleter_id, deleter)| { let delete_sinks = ctx .marked_sinks(deleter.data_sinks(), marker!(to_delete)) + .map(|s| s.clone().into()) .collect::>(); + let delete_sinks_borrowed = delete_sinks.iter().collect::>(); user_data_types.iter().all(|&t| { let sources = ctx.srcs_with_type(deleter, t).collect::>(); - 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() { diff --git a/props/websubmit/src/main.rs b/props/websubmit/src/main.rs index 61ebaa5add..2663887cca 100644 --- a/props/websubmit/src/main.rs +++ b/props/websubmit/src/main.rs @@ -25,8 +25,13 @@ impl DeletionProp { .collect::>(); 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::>(); 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; } }