From 320625ecaff99c3f0f07f780463fac9ab8c7c906 Mon Sep 17 00:00:00 2001
From: "Christoph J. Scherr" <Christoph.Scherr@newtec.de>
Date: Fri, 6 Sep 2024 12:01:48 +0200
Subject: [PATCH] feat: challenge 2 + better hints

---
 src/challenge/c1.rs  |  6 ++--
 src/challenge/c2.rs  | 80 ++++++++++++++++++++++++++++++++++++++++++++
 src/challenge/mod.rs | 20 +++++++----
 src/lib.rs           |  2 +-
 src/main.rs          | 10 ++----
 5 files changed, 99 insertions(+), 19 deletions(-)
 create mode 100644 src/challenge/c2.rs

diff --git a/src/challenge/c1.rs b/src/challenge/c1.rs
index e5ed0d8..c2ca865 100644
--- a/src/challenge/c1.rs
+++ b/src/challenge/c1.rs
@@ -19,8 +19,8 @@ impl Challenge for C1 {
     fn new(config: Config, vault: VaultRef) -> Self {
         Self { config, vault }
     }
-    fn hint() -> String {
-        String::from("TCP connect to 1337")
+    fn hints() -> Vec<String> {
+        vec![String::from("TCP connect to 1337")]
     }
     async fn serve(self) -> anyhow::Result<()> {
         info!("serving challenge 1");
@@ -35,7 +35,7 @@ impl Challenge for C1 {
                     continue;
                 }
             };
-            has_won(addr);
+            has_won(&addr);
             tokio::spawn(async move {
                 if let Err(e) = stream.write_all(vault.secret().as_bytes()).await {
                     warn!("could not write to peer {addr}: {e}");
diff --git a/src/challenge/c2.rs b/src/challenge/c2.rs
new file mode 100644
index 0000000..95b7c3f
--- /dev/null
+++ b/src/challenge/c2.rs
@@ -0,0 +1,80 @@
+use anyhow::Result;
+use async_trait::async_trait;
+use libpt::log::{info, warn};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::{TcpListener, TcpStream};
+
+use super::Challenge;
+use crate::has_won;
+use crate::vault::{Config, VaultRef};
+
+pub struct C2 {
+    config: Config,
+    vault: VaultRef,
+}
+
+impl C2 {
+    async fn win(
+        vault: &VaultRef,
+        stream: &mut TcpStream,
+        addr: &std::net::SocketAddr,
+    ) -> Result<()> {
+        has_won(addr);
+        if let Err(e) = stream.write_all(vault.secret().as_bytes()).await {
+            warn!("could not write to peer {addr}: {e}");
+            return Err(e.into());
+        };
+        if let Err(e) = stream.shutdown().await {
+            warn!("could end connection to peer {addr}: {e}");
+            return Err(e.into());
+        };
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl Challenge for C2 {
+    fn new(config: Config, vault: VaultRef) -> Self {
+        Self { config, vault }
+    }
+    fn hints() -> Vec<String> {
+        vec![String::from(
+            "TCP connect to 1337 and give me a special u16",
+        )]
+    }
+    async fn serve(self) -> anyhow::Result<()> {
+        info!("serving challenge 2");
+        let listener = TcpListener::bind(self.config.addr).await?;
+
+        loop {
+            let vault = self.vault.clone();
+            let (mut stream, addr) = match listener.accept().await {
+                Ok(s) => s,
+                Err(err) => {
+                    warn!("could not accept tcp stream: {err:?}");
+                    continue;
+                }
+            };
+            info!("new peer: {addr}");
+            tokio::spawn(async move {
+                let mut buf: u16;
+                loop {
+                    match stream.read_u16().await {
+                        Err(e) => {
+                            warn!("could read from {addr}: {e}");
+                            return;
+                        }
+                        Ok(u) => buf = u,
+                    }
+                    if buf == 1337 {
+                        if let Err(e) = Self::win(&vault, &mut stream, &addr).await {
+                            warn!("could not let {addr} win: {e}");
+                        }
+                    } else {
+                        info!("peer wrote crap: {buf}");
+                    }
+                }
+            });
+        }
+    }
+}
diff --git a/src/challenge/mod.rs b/src/challenge/mod.rs
index 039efa3..89cb531 100644
--- a/src/challenge/mod.rs
+++ b/src/challenge/mod.rs
@@ -3,9 +3,8 @@ use async_trait::async_trait;
 
 use crate::vault::{Config, VaultRef};
 
-use self::c1::C1;
-
 pub mod c1;
+pub mod c2;
 
 #[async_trait]
 pub trait Challenge
@@ -13,12 +12,19 @@ where
     Self: Sized,
 {
     fn new(config: Config, vault: VaultRef) -> Self;
-    fn hint() -> String;
+    fn hints() -> Vec<String>;
     async fn serve(self) -> anyhow::Result<()>;
-    fn get_challenge(idx: u16, config: Config, vault: VaultRef) -> anyhow::Result<impl Challenge> {
-        match idx {
-            1 => Ok(C1::new(config, vault)),
-            _ => Err(anyhow!("no challenge with that index exists")),
+}
+
+pub async fn select_and_start(index: u16, config: Config, vault: VaultRef) -> anyhow::Result<()> {
+    match index {
+        1 => c1::C1::new(config, vault).serve().await?,
+        2 => c2::C2::new(config, vault).serve().await?,
+        _ => {
+            return Err(anyhow!(
+                "no challenge with index {index} does currently exist"
+            ))
         }
     }
+    Ok(())
 }
diff --git a/src/lib.rs b/src/lib.rs
index 12acbd1..b5db9de 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,6 @@ pub mod challenge;
 pub mod vault;
 
 #[inline]
-pub(crate) fn has_won(addr: std::net::SocketAddr) {
+pub(crate) fn has_won(addr: &std::net::SocketAddr) {
     libpt::log::info!("Sending the secret to {addr}")
 }
diff --git a/src/main.rs b/src/main.rs
index fb1d7d3..428db12 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,7 +1,6 @@
 use anyhow::anyhow;
 
-use wooly_vault::challenge::c1::C1;
-use wooly_vault::challenge::Challenge;
+use wooly_vault::challenge::select_and_start;
 use wooly_vault::vault::{Config, Vault};
 
 #[tokio::main(flavor = "current_thread")]
@@ -22,12 +21,7 @@ async fn main() -> anyhow::Result<()> {
 
     println!("What challenge to serve?");
     let i = select_challenge()?;
-    match i {
-        1 => {
-            C1::new(conf, v).serve().await?;
-        }
-        _ => return Err(anyhow!("no challenge with index {i} does currently exist")),
-    }
+    select_and_start(i, conf, v).await?;
 
     Ok(())
 }