Add TCP command (#284)

* Add tcp command

* Add binary support

* Add shutdown

puffin/Cargo.toml

@@ -20,7 +20,7 @@ libafl = "0.10.1"
 libafl_targets = "0.10.1"
 
 # Logging
-log = "0.4.17"
+log = { version = "0.4.17" }
 log4rs = "1.2.0"
 
 # Other Dependencies

puffin/src/cli.rs

@@ -22,7 +22,7 @@ use crate::{
         start, FuzzerConfig,
     },
     graphviz::write_graphviz,
-    log::create_stdout_config,
+    log::create_stderr_config,
     protocol::{ProtocolBehavior, ProtocolMessage},
     put::PutOptions,
     put_registry::PutRegistry,
@@ -68,14 +68,23 @@ fn create_app() -> Command {
             Command::new("binary-attack")
                 .about("Serializes a trace as much as possible and output its")
                 .arg(arg!(<input> "The file which stores a trace"))
-                .arg(arg!(<output> "The file to write serialized data to"))
+                .arg(arg!(<output> "The file to write serialized data to")),
+            Command::new("tcp")
+                .about("Executes a trace against a TCP client/server")
+                .arg(arg!(<input> "The file which stores a trace"))
+                .arg(arg!(-c --cwd [p] "The current working directory for the binary"))
+                .arg(arg!(-b --binary [p] "The program to start"))
+                .arg(arg!(-a --args [a] "The args of the program"))
+                .arg(arg!(-t --host [h] "The host to connect to, or the server host"))
+                .arg(arg!(-p --port [n] "The client port to connect to, or the server port")
+                    .value_parser(value_parser!(u16).range(1..)))
         ])
 }
 
 pub fn main<PB: ProtocolBehavior + Clone + 'static>(
     put_registry: &'static PutRegistry<PB>,
 ) -> ExitCode {
-    let handle = match log4rs::init_config(create_stdout_config(LevelFilter::Info)) {
+    let handle = match log4rs::init_config(create_stderr_config(LevelFilter::Info)) {
         Ok(handle) => handle,
         Err(err) => {
             error!("Failed to init logging: {:?}", err);
@@ -206,6 +215,50 @@ pub fn main<PB: ProtocolBehavior + Clone + 'static>(
             error!("Failed to create trace output: {:?}", err);
             return ExitCode::FAILURE;
         }
+    } else if let Some(matches) = matches.subcommand_matches("tcp") {
+        let input: &String = matches.get_one("input").unwrap();
+        let prog: Option<&String> = matches.get_one("binary");
+        let args: Option<&String> = matches.get_one("args");
+        let cwd: Option<&String> = matches.get_one("cwd");
+        let default_host = "127.0.0.1".to_string();
+        let host: &String = matches.get_one("host").unwrap_or(&default_host);
+        let port = matches
+            .get_one::<u16>("port")
+            .unwrap_or(&44338u16)
+            .to_string();
+
+        let trace = Trace::<PB::Matcher>::from_file(input).unwrap();
+        let ctx = TraceContext::new(put_registry, default_put_options().clone());
+
+        let mut options = vec![("port", port.as_str()), ("host", &host)];
+
+        if let Some(prog) = prog {
+            options.push(("prog", &prog))
+        }
+
+        if let Some(args) = args {
+            options.push(("args", &args))
+        }
+
+        if let Some(cwd) = cwd {
+            options.push(("cwd", &cwd))
+        }
+
+        let put = PutDescriptor {
+            name: PutName(['T', 'C', 'P', '_', '_', '_', '_', '_', '_', '_']),
+            options: PutOptions::from_slice_vec(options),
+        };
+
+        let server = trace.descriptors[0].name;
+        let mut context = trace
+            .execute_with_non_default_puts(&put_registry, &[(server, put)])
+            .unwrap();
+
+        let server = AgentName::first();
+        let shutdown = context.find_agent_mut(server).unwrap().put_mut().shutdown();
+        info!("{}", shutdown);
+
+        return ExitCode::SUCCESS;
     } else {
         let experiment_path = if let Some(matches) = matches.subcommand_matches("experiment") {
             let title: &String = matches.get_one("title").unwrap();
@@ -348,6 +401,11 @@ use nix::{
     unistd::{fork, ForkResult},
 };
 
+use crate::{
+    agent::AgentName,
+    put::{PutDescriptor, PutName},
+};
+
 pub fn expect_crash<R>(func: R)
 where
     R: FnOnce(),

puffin/src/log.rs

@@ -24,7 +24,7 @@ fn create_config(
         .unwrap()
 }
 
-pub fn create_stdout_config(default_level: LevelFilter) -> Config {
+pub fn create_stderr_config(default_level: LevelFilter) -> Config {
     let stderr = ConsoleAppender::builder()
         .target(log4rs::append::console::Target::Stderr)
         .encoder(Box::new(PatternEncoder::new(

tlspuffin/Cargo.toml

@@ -211,7 +211,7 @@ introspection = ["puffin/introspection"]
 
 puffin = { path = "../puffin" }
 
-log = "0.4.17"
+log = { version = "0.4.17" }
 itertools = "0.10.3"
 smallvec = "1.8.1"
 libc = { version = "0.2.126" }

tlspuffin/src/tcp/mod.rs

@@ -11,7 +11,7 @@ use std::{
     time::Duration,
 };
 
-use log::error;
+use log::{error, info, warn};
 use puffin::{
     agent::{AgentDescriptor, AgentName, AgentType},
     error::Error,
@@ -43,28 +43,45 @@ pub fn new_tcp_factory() -> Box<dyn Factory<TLSProtocolBehavior>> {
 
             let options = &put_descriptor.options;
 
-            let args = options
-                .get_option("args")
-                .ok_or_else(|| Error::Agent("Unable to find args".to_string()))?
-                .to_owned();
-            let prog = options
-                .get_option("prog")
-                .ok_or_else(|| Error::Agent("Unable to find prog".to_string()))?
-                .to_owned();
-            let cwd = options
-                .get_option("cwd")
-                .map(|cwd| Some(cwd.to_owned()))
-                .unwrap_or_default();
-
-            if agent_descriptor.typ == AgentType::Client {
-                let mut server = TcpServerPut::new(agent_descriptor, &put_descriptor)?;
-                server.set_process(TLSProcess::new(&prog, &args, cwd.as_ref()));
-                Ok(Box::new(server))
+            if options.get_option("args").is_some() {
+                info!("Trace contains TCP running information we shall reuse.");
+                let args = options
+                    .get_option("args")
+                    .ok_or_else(|| {
+                        Error::Agent(format!(
+                            "{} // {:?}",
+                            "Unable to find args".to_string(),
+                            put_descriptor
+                        ))
+                    })?
+                    .to_owned();
+                let prog = options
+                    .get_option("prog")
+                    .ok_or_else(|| Error::Agent("Unable to find prog".to_string()))?
+                    .to_owned();
+                let cwd = options
+                    .get_option("cwd")
+                    .map(|cwd| Some(cwd.to_owned()))
+                    .unwrap_or_default();
+                if agent_descriptor.typ == AgentType::Client {
+                    let mut server = TcpServerPut::new(agent_descriptor, &put_descriptor)?;
+                    server.set_process(TLSProcess::new(&prog, &args, cwd.as_ref()));
+                    Ok(Box::new(server))
+                } else {
+                    let process = TLSProcess::new(&prog, &args, cwd);
+                    let mut client = TcpClientPut::new(agent_descriptor, &put_descriptor)?;
+                    client.set_process(process);
+                    Ok(Box::new(client))
+                }
             } else {
-                let process = TLSProcess::new(&prog, &args, cwd);
-                let mut client = TcpClientPut::new(agent_descriptor, &put_descriptor)?;
-                client.set_process(process);
-                Ok(Box::new(client))
+                info!("Trace contains no TCP running information so we fall back to external TCP client and servers.");
+                if agent_descriptor.typ == AgentType::Client {
+                    let server = TcpServerPut::new(agent_descriptor, &put_descriptor)?;
+                    Ok(Box::new(server))
+                } else {
+                    let client = TcpClientPut::new(agent_descriptor, &put_descriptor)?;
+                    Ok(Box::new(client))
+                }
             }
         }
 
@@ -316,7 +333,14 @@ fn addr_from_config(put_descriptor: &PutDescriptor) -> Result<SocketAddr, AddrPa
     let port = options
         .get_option("port")
         .and_then(|value| u16::from_str(value).ok())
-        .expect("Failed to parse port option");
+        .unwrap_or_else(|| {
+            let port = 44338;
+            warn!(
+                "Failed to parse port option (maybe you executed a trace that was not produced in \
+            TCP mode?). We anyway fall back to port {port}."
+            );
+            port
+        });
 
     Ok(SocketAddr::new(IpAddr::from_str(host)?, port))
 }