refactor(puffin): replace variants of Trace::execute() by builder

tlspuffin · May 28, 2024 · a04a9b5 · a04a9b5
1 parent bcc522b
commit a04a9b5
Show file tree

Hide file tree

Showing 9 changed files with 62 additions and 84 deletions.
diff --git a/puffin/src/algebra/mod.rs b/puffin/src/algebra/mod.rs
@@ -554,7 +554,6 @@ mod tests {
         algebra::{
             atoms::Variable, dynamic_function::TypeShape, signature::Signature, AnyMatcher, Term,
         },
-        put::PutOptions,
         put_registry::{Factory, PutRegistry},
         term,
         trace::{Knowledge, TraceContext},

diff --git a/puffin/src/cli.rs b/puffin/src/cli.rs
@@ -312,9 +312,11 @@ where
         };
 
         let server = trace.descriptors[0].name;
-        let mut context = trace
-            .execute_with_non_default_puts(&put_registry, &[(server, put)])
-            .unwrap();
+        let mut context = TraceContext::builder(&put_registry)
+            .set_put(server, put)
+            .build();
+
+        trace.execute(&mut context).unwrap();
 
         let server = AgentName::first();
         let shutdown = context.find_agent_mut(server).unwrap().put_mut().shutdown();

diff --git a/puffin/src/trace.rs b/puffin/src/trace.rs
@@ -455,39 +455,6 @@ impl<M: Matcher> Trace<M> {
         Ok(())
     }
 
-    pub fn execute_deterministic<PB>(
-        &self,
-        put_registry: &PutRegistry<PB>,
-        default_put_options: PutOptions,
-    ) -> Result<TraceContext<PB>, Error>
-    where
-        PB: ProtocolBehavior<Matcher = M>,
-    {
-        let mut ctx = TraceContext::builder(put_registry)
-            .set_default_put_options(default_put_options)
-            .set_deterministic(true)
-            .build();
-
-        self.execute(&mut ctx)?;
-        Ok(ctx)
-    }
-
-    pub fn execute_with_non_default_puts<PB>(
-        &self,
-        put_registry: &PutRegistry<PB>,
-        descriptors: &[(AgentName, PutDescriptor)],
-    ) -> Result<TraceContext<PB>, Error>
-    where
-        PB: ProtocolBehavior<Matcher = M>,
-    {
-        let mut ctx = TraceContext::builder(put_registry)
-            .set_puts(descriptors)
-            .build();
-
-        self.execute(&mut ctx)?;
-        Ok(ctx)
-    }
-
     pub fn serialize_postcard(&self) -> Result<Vec<u8>, postcard::Error> {
         postcard::to_allocvec(&self)
     }

diff --git a/sshpuffin/src/ssh/seeds.rs b/sshpuffin/src/ssh/seeds.rs
@@ -166,7 +166,7 @@ pub fn seed_successful(client: AgentName, server: AgentName) -> Trace<AnyMatcher
 
 #[cfg(test)]
 mod tests {
-    use puffin::{agent::AgentName, put::PutOptions};
+    use puffin::{agent::AgentName, trace::TraceContext};
     use test_log::test;
 
     use crate::{libssh::ssh::set_log_level, ssh::seeds::seed_successful, ssh_registry};
@@ -176,9 +176,11 @@ mod tests {
         set_log_level(100);
         let client = AgentName::first();
         let trace = seed_successful(client, client.next());
-        let context = trace
-            .execute_deterministic(&ssh_registry(), PutOptions::default())
-            .unwrap();
+        let mut context = TraceContext::builder(&ssh_registry())
+            .set_deterministic(true)
+            .build();
+
+        trace.execute(&mut context).unwrap();
 
         assert!(context
             .find_agent(client)

diff --git a/tlspuffin/src/integration_tests/mutations.rs b/tlspuffin/src/integration_tests/mutations.rs
@@ -12,7 +12,6 @@ use puffin::{
         mutators::{MutationResult, Mutator},
         state::StdState,
     },
-    put::PutOptions,
     trace::{Action, Step, Trace, TraceContext},
 };
 

diff --git a/tlspuffin/src/tcp/mod.rs b/tlspuffin/src/tcp/mod.rs
@@ -733,6 +733,7 @@ mod tests {
         agent::{AgentName, TLSVersion},
         put::PutDescriptor,
         put_registry::TCP_PUT,
+        trace::TraceContext,
     };
     use test_log::test;
 
@@ -761,12 +762,12 @@ mod tests {
         let trace = seed_session_resumption_dhe_full.build_trace();
         let initial_server = trace.prior_traces[0].descriptors[0].name;
         let server = trace.descriptors[0].name;
-        let mut context = trace
-            .execute_with_non_default_puts(
-                &put_registry,
-                &[(initial_server, put.clone()), (server, put)],
-            )
-            .unwrap();
+        let mut context = TraceContext::builder(&put_registry)
+            .set_put(initial_server, put.clone())
+            .set_put(server, put)
+            .build();
+
+        trace.execute(&mut context).unwrap();
 
         let server = AgentName::first().next();
         let shutdown = context.find_agent_mut(server).unwrap().put_mut().shutdown();
@@ -787,9 +788,11 @@ mod tests {
         let put_registry = tls_registry();
         let trace = seed_client_attacker_full.build_trace();
         let server = trace.descriptors[0].name;
-        let mut context = trace
-            .execute_with_non_default_puts(&put_registry, &[(server, put)])
-            .unwrap();
+        let mut context = TraceContext::builder(&put_registry)
+            .set_put(server, put)
+            .build();
+
+        trace.execute(&mut context).unwrap();
 
         let server = AgentName::first();
         let shutdown = context.find_agent_mut(server).unwrap().put_mut().shutdown();
@@ -821,12 +824,12 @@ mod tests {
         let descriptors = &trace.descriptors;
         let client_name = descriptors[0].name;
         let server_name = descriptors[1].name;
-        let mut context = trace
-            .execute_with_non_default_puts(
-                &put_registry,
-                &[(client_name, client), (server_name, server)],
-            )
-            .unwrap();
+        let mut context = TraceContext::builder(&put_registry)
+            .set_put(client_name, client)
+            .set_put(server_name, server)
+            .build();
+
+        trace.execute(&mut context).unwrap();
 
         let client = AgentName::first();
         let shutdown = context.find_agent_mut(client).unwrap().put_mut().shutdown();
@@ -863,12 +866,12 @@ mod tests {
         let descriptors = &trace.descriptors;
         let client_name = descriptors[0].name;
         let server_name = descriptors[1].name;
-        let mut context = trace
-            .execute_with_non_default_puts(
-                &put_registry,
-                &[(client_name, client), (server_name, server)],
-            )
-            .unwrap();
+        let mut context = TraceContext::builder(&put_registry)
+            .set_put(client_name, client)
+            .set_put(server_name, server)
+            .build();
+
+        trace.execute(&mut context).unwrap();
 
         let client = AgentName::first();
         let shutdown = context.find_agent_mut(client).unwrap().put_mut().shutdown();

diff --git a/tlspuffin/src/test_utils.rs b/tlspuffin/src/test_utils.rs
@@ -4,7 +4,7 @@ use log::info;
 use puffin::{
     execution::{forked_execution, ExecutionStatus},
     put::PutOptions,
-    trace::Trace,
+    trace::{Trace, TraceContext},
 };
 
 use crate::{put_registry::tls_registry, query::TlsQueryMatcher};
@@ -36,10 +36,13 @@ pub fn expect_trace_crash(
         .map(|_| {
             forked_execution(
                 || {
-                    // Ignore Rust errors
-                    let _ = trace
-                        .clone()
-                        .execute_deterministic(&tls_registry(), default_put_options.clone());
+                    let mut context = TraceContext::builder(&tls_registry())
+                        .set_default_put_options(default_put_options.clone())
+                        .set_deterministic(true)
+                        .build();
+
+                    // NOTE: we ignore Rust errors because we expect a crash
+                    let _ = trace.clone().execute(&mut context);
                 },
                 timeout,
             )

diff --git a/tlspuffin/src/tls/trace_helper.rs b/tlspuffin/src/tls/trace_helper.rs
@@ -3,7 +3,6 @@ use std::fs::File;
 use puffin::{
     agent::AgentName,
     libafl::inputs::Input,
-    put::PutOptions,
     trace::{Trace, TraceContext},
 };
 
@@ -22,9 +21,12 @@ pub trait TraceExecutor<A> {
 
 impl<A, H: TraceHelper<A>> TraceExecutor<A> for H {
     fn execute_trace(self) -> TraceContext<TLSProtocolBehavior> {
-        self.build_trace()
-            .execute_deterministic(&tls_registry(), PutOptions::default())
-            .unwrap()
+        let mut context = TraceContext::builder(&tls_registry())
+            .set_deterministic(true)
+            .build();
+
+        self.build_trace().execute(&mut context).unwrap();
+        context
     }
 
     fn store_to_seeds(self) {

diff --git a/tlspuffin/src/tls/vulnerabilities.rs b/tlspuffin/src/tls/vulnerabilities.rs
@@ -1296,6 +1296,7 @@ pub mod tests {
             agent::{AgentName, TLSVersion},
             put::PutDescriptor,
             put_registry::TCP_PUT,
+            trace::TraceContext,
         };
         use test_log::test;
 
@@ -1329,12 +1330,12 @@ pub mod tests {
             let descriptors = &trace.descriptors;
             let client_name = descriptors[0].name;
             let server_name = descriptors[1].name;
-            let mut context = trace
-                .execute_with_non_default_puts(
-                    &put_registry,
-                    &[(client_name, client), (server_name, server)],
-                )
-                .unwrap();
+            let mut context = TraceContext::builder(&put_registry)
+                .set_put(client_name, client)
+                .set_put(server_name, server)
+                .build();
+
+            trace.execute(&mut context).unwrap();
 
             let client = AgentName::first();
             let shutdown = context.find_agent_mut(client).unwrap().put_mut().shutdown();
@@ -1356,12 +1357,12 @@ pub mod tests {
             let trace = seed_cve_2022_39173_full.build_trace();
             let initial_server = trace.prior_traces[0].descriptors[0].name;
             let server = trace.descriptors[0].name;
-            let mut context = trace
-                .execute_with_non_default_puts(
-                    &put_registry,
-                    &[(initial_server, put.clone()), (server, put)],
-                )
-                .unwrap();
+            let mut context = TraceContext::builder(&put_registry)
+                .set_put(initial_server, put.clone())
+                .set_put(server, put)
+                .build();
+
+            trace.execute(&mut context).unwrap();
 
             let server = AgentName::first().next();
             let shutdown = context.find_agent_mut(server).unwrap().put_mut().shutdown();