wooly-vault/src/challenge/c2.rs

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::config::Config;
use crate::has_won;
use crate::vault::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",
        )]
    }

    fn solution() -> String {
        String::from("Connect by TCP, send 1337 as bytes (not text).")
    }

    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}");
                    }
                }
            });
        }
    }
}