striker / src / net.rs

use std::{
    io::{Read, Write},
    net::{SocketAddr, TcpStream},
    sync::{Arc, LazyLock, Mutex, OnceLock},
    time::Duration,
};

use async_channel::{Receiver, Sender};
use async_io::{Async, Timer};
use bevy::{
    app::{Plugin, Startup},
    ecs::{resource::Resource, system::Commands},
    tasks::IoTaskPool,
};
use futures_concurrency::future::Race;
use snow::{Builder, TransportState, params::NoiseParams};

use crate::constants::NOISE_PSK;

static PARAMS: LazyLock<NoiseParams> =
    LazyLock::new(|| "Noise_XXpsk3_25519_ChaChaPoly_BLAKE2s".parse().unwrap());
static LOCAL_PRIVATE_KEY: OnceLock<snow::Keypair> = OnceLock::new();

pub enum StrikeUpdateState {
    Disconnected,
    Connecting,
    Connected,
    Updating(striker_proto::StrikerResponse),
}

#[derive(Debug, Resource)]
pub struct StrikeUpdates(pub Receiver<StrikeUpdateState>);

#[derive(Debug, Resource)]
pub struct StrikeRequests(pub Sender<striker_proto::StrikerRequest>);

#[derive(Debug, Resource)]
pub struct StrikeActions(pub Sender<StrikeAction>);

pub enum StrikeAction {
    Connect(SocketAddr),
    Disconnect,
}

pub fn setup_strike_connection(mut commands: Commands) {
    let io = IoTaskPool::get();

    let (signal_tx, signal_rx) = async_channel::bounded(2);
    let (req_tx, req_rx) = async_channel::bounded(1);
    let (resp_tx, resp_rx) = async_channel::bounded(64);

    io.spawn(async move {
        let mut read_buf = vec![0u8; 4096];
        let mut enc_buf = vec![0u8; 4096];
        let mut write_buf = vec![0u8; 4096];

        while let Ok(StrikeAction::Connect(addr)) = signal_rx.recv().await {
            let net_fut = async {
                loop {
                    resp_tx.send(StrikeUpdateState::Connecting).await.ok();
                    let Ok(stream) = Async::<TcpStream>::connect(addr).await else {
                        Timer::after(Duration::from_secs(1)).await;
                        continue;
                    };

                    resp_tx.send(StrikeUpdateState::Connected).await.ok();
                    stream.write_with(|s| s.set_nodelay(true)).await.ok();

                    let Ok(transport) = noise_handshake(&stream).await else {
                        resp_tx.send(StrikeUpdateState::Disconnected).await.ok();
                        continue;
                    };

                    let transport_access = Arc::new(Mutex::new(transport));

                    let read_fut = async {
                        while let Ok(buh) = recv(&stream).await {
                            if let Ok(decrypted) = {
                                transport_access
                                    .lock()
                                    .unwrap()
                                    .read_message(&buh, &mut read_buf)
                            } && let Ok(data) =
                                striker_proto::receive_response(&mut read_buf[..decrypted])
                                && resp_tx
                                    .send(StrikeUpdateState::Updating(data))
                                    .await
                                    .is_err()
                            {
                                break;
                            }
                        }
                    };

                    let write_fut = async {
                        while let Ok(req) = req_rx.recv().await {
                            if let Ok(payload) = striker_proto::send_request(req, &mut write_buf)
                                && let Ok(encrypted) = {
                                    transport_access
                                        .lock()
                                        .unwrap()
                                        .write_message(payload, &mut enc_buf)
                                }
                                && send(&stream, &enc_buf[..encrypted]).await.is_err()
                            {
                                break;
                            };
                        }
                    };

                    (read_fut, write_fut).race().await;

                    stream
                        .write_with(|s| s.shutdown(std::net::Shutdown::Both))
                        .await
                        .ok();

                    break;
                }
            };

            let cancel_fut = async {
                while signal_rx
                    .recv()
                    .await
                    .is_ok_and(|strike| !matches!(strike, StrikeAction::Disconnect))
                {
                }
            };

            (net_fut, cancel_fut).race().await;
            resp_tx.send(StrikeUpdateState::Disconnected).await.ok();
        }
    })
    .detach();

    commands.insert_resource(StrikeActions(signal_tx));
    commands.insert_resource(StrikeUpdates(resp_rx));
    commands.insert_resource(StrikeRequests(req_tx));
}

pub struct NetPlugin;

impl Plugin for NetPlugin {
    fn build(&self, app: &mut bevy::app::App) {
        app.add_systems(Startup, setup_strike_connection);
    }
}

async fn noise_handshake(tcp: &Async<TcpStream>) -> color_eyre::Result<TransportState> {
    let builder = Builder::new(PARAMS.clone());
    let static_key = LOCAL_PRIVATE_KEY.get_or_init(|| builder.generate_keypair().unwrap());

    let mut noise = builder
        .local_private_key(&static_key.private)?
        .psk(3, &NOISE_PSK)?
        .build_initiator()?;

    let mut payload = vec![0u8; 2048];

    // -> e
    let len = noise.write_message(&[], &mut payload)?;

    send(tcp, &payload[..len]).await?;

    // <- e, ee, s, es
    noise.read_message(&recv(tcp).await?, &mut payload)?;

    // -> s, se
    let len = noise.write_message(&[], &mut payload)?;

    send(tcp, &payload[..len]).await?;

    let transport = noise.into_transport_mode()?;

    Ok(transport)
}

async fn recv(stream: &Async<TcpStream>) -> std::io::Result<Vec<u8>> {
    let mut msg_len_buf = [0_u8; 2];

    stream
        .read_with(|mut stream| {
            stream.read_exact(&mut msg_len_buf)?;
            let msg_len = usize::from(u16::from_be_bytes(msg_len_buf));
            let mut msg = vec![0_u8; msg_len];
            stream.read_exact(&mut msg[..])?;

            Ok(msg)
        })
        .await
}

/// Hyper-basic stream transport sender. 16-bit BE size followed by payload.
async fn send(stream: &Async<TcpStream>, buf: &[u8]) -> std::io::Result<()> {
    stream
        .write_with(|mut stream| {
            let len = u16::try_from(buf.len()).expect("message too large");
            stream.write_all(&len.to_be_bytes())?;
            stream.write_all(buf)
        })
        .await
}