harlequinn 0.1.1

use std::time::Duration;

use {
    futures::StreamExt,
    quinn::{Connection, IncomingBiStreams, RecvStream, SendStream},
    tokio::time::timeout,
};

const HANDSHAKE_TIMEOUT: Duration = Duration::from_millis(1000);

pub async fn verify_handshake_server(
    connection: &Connection,
    protocol_checksum: u32,
) -> Option<(SendStream, RecvStream)> {
    let checksum_data = protocol_checksum.to_be_bytes();

    // On the server we need to open a stream
    let (mut send, mut recv) = connection.open_bi().await.ok()?;

    // Send the handshake message the client can identify the application by
    // TODO: User defined string hashed
    send.write_all(&checksum_data).await.ok()?;

    // The client should send us back another checksum
    // Use a timeout, because at this point we don't have a channel yet that can get backed up to
    // detect slowloris attacks.
    let mut client_checksum_data = [0u8; 4];
    timeout(
        HANDSHAKE_TIMEOUT,
        recv.read_exact(&mut client_checksum_data),
    )
    .await
    // Timeout error
    .ok()?
    // Read error
    .ok()?;
    let client_checksum = u32::from_be_bytes(client_checksum_data);

    // Check to make sure it is the same
    if protocol_checksum != client_checksum {
        return None;
    }

    Some((send, recv))
}

pub async fn verify_handshake_client(
    bi_streams: &mut IncomingBiStreams,
    protocol_checksum: u32,
) -> Option<(SendStream, RecvStream)> {
    // On the client we need to await a stream to be opened
    let stream_result = timeout(HANDSHAKE_TIMEOUT, bi_streams.next()).await.ok()?;

    let (mut send, mut recv) = stream_result?.ok()?;

    // Now that we have a stream, wait for the server to send us its checksum
    let mut server_checksum_data = [0u8; 4];
    timeout(
        HANDSHAKE_TIMEOUT,
        recv.read_exact(&mut server_checksum_data),
    )
    .await
    // Timeout error
    .ok()?
    // Read error
    .ok()?;
    let server_checksum = u32::from_be_bytes(server_checksum_data);

    // Verify the checksum
    if protocol_checksum != server_checksum {
        return None;
    }

    // Send back our confirmation
    send.write_all(&server_checksum_data).await.ok()?;

    Some((send, recv))
}