hydra_sync/

protocol.rs

use crate::client::ServerMetrics;
use crate::crypto::{NONCE_LEN, TAG_LEN, decrypt_into, encrypt_into, generate_x25519_keypair};
use anyhow::{Result, bail};
use bytes::BytesMut;
use rand::Rng;
use sha3::{Digest, Sha3_256};
use std::fmt::Display;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use x25519_dalek::PublicKey;

/// Represents the role of a client in the Hydra protocol, either as an Observer, Producer, or Consumer.
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Role {
    Producer,
    Consumer,
    Observer,
}

impl Role {
    #[inline(always)]
    pub fn to_u8(self) -> u8 {
        match self {
            Role::Producer => 0x00,
            Role::Consumer => 0x01,
            Role::Observer => 0x02,
        }
    }
    #[inline(always)]
    pub fn from_u8(val: u8) -> Result<Self> {
        match val {
            0x00 => Ok(Self::Producer),
            0x01 => Ok(Self::Consumer),
            0x02 => Ok(Self::Observer),
            _ => bail!("Unknown role: {:#04x}", val),
        }
    }
}

impl Display for Role {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Role::Observer => write!(f, "Observer"),
            Role::Producer => write!(f, "Producer"),
            Role::Consumer => write!(f, "Consumer"),
        }
    }
}

const HANDSHAKE_SALT_LEN: usize = 128;
const HANDSHAKE_INFO: &[u8] =
    concat!("hydra-sync transport key/v", env!("CARGO_PKG_VERSION")).as_bytes();

/// Perform X25519 key exchange handshake on client side and return the derived shared secret key
pub async fn perform_client_handshake<R: AsyncReadExt + Unpin, W: AsyncWriteExt + Unpin>(
    reader: &mut R,
    writer: &mut W,
) -> Result<[u8; 32]> {
    let (secret, client_pub) = generate_x25519_keypair()?;

    let mut client_nonce = [0u8; HANDSHAKE_SALT_LEN];
    rand::rng().fill_bytes(&mut client_nonce);

    // write nonce + 32 bytes key
    writer.write_all(&client_nonce).await?;
    writer.write_all(client_pub.as_bytes()).await?;
    writer.flush().await?;

    let mut server_nonce = [0u8; HANDSHAKE_SALT_LEN];
    reader.read_exact(&mut server_nonce).await?;
    let mut server_pub_bytes = [0u8; 32];
    reader.read_exact(&mut server_pub_bytes).await?;
    let server_pub = PublicKey::from(server_pub_bytes);

    let shared_secret = secret.diffie_hellman(&server_pub);
    derive_encryption_key(
        shared_secret.as_bytes(),
        client_pub.as_bytes(),
        &client_nonce,
        server_pub.as_bytes(),
        &server_nonce,
    )
}

/// Perform X25519 key exchange handshake on server side and return the derived shared secret key
pub async fn perform_server_handshake<R: AsyncReadExt + Unpin, W: AsyncWriteExt + Unpin>(
    reader: &mut R,
    writer: &mut W,
) -> Result<[u8; 32]> {
    let (secret, server_pub) = generate_x25519_keypair()?;

    // read client nonce + 32 bytes key
    let mut client_nonce = [0u8; HANDSHAKE_SALT_LEN];
    reader.read_exact(&mut client_nonce).await?;
    let mut client_pub_bytes = [0u8; 32];
    reader.read_exact(&mut client_pub_bytes).await?;
    let client_pub = PublicKey::from(client_pub_bytes);

    let mut server_nonce = [0u8; HANDSHAKE_SALT_LEN];
    rand::rng().fill_bytes(&mut server_nonce);

    writer.write_all(&server_nonce).await?;
    writer.write_all(server_pub.as_bytes()).await?;
    writer.flush().await?;

    let shared_secret = secret.diffie_hellman(&client_pub);
    derive_encryption_key(
        shared_secret.as_bytes(),
        client_pub.as_bytes(),
        &client_nonce,
        server_pub.as_bytes(),
        &server_nonce,
    )
}

#[inline(always)]
/// Derives encryption key by hashing the concatenation of client/server public keys, nonce, handshake info, and shared secret using SHA3-256
fn derive_encryption_key(
    shared_secret: &[u8],
    client_pub: &[u8; 32],
    client_nonce: &[u8; HANDSHAKE_SALT_LEN],
    server_pub: &[u8; 32],
    server_nonce: &[u8; HANDSHAKE_SALT_LEN],
) -> Result<[u8; 32]> {
    let mut transcript = vec![0u8; 32 + HANDSHAKE_SALT_LEN + 32 + HANDSHAKE_SALT_LEN + 32];
    transcript.extend_from_slice(client_pub);
    transcript.extend_from_slice(client_nonce);
    transcript.extend_from_slice(server_pub);
    transcript.extend_from_slice(server_nonce);
    transcript.extend_from_slice(shared_secret);
    transcript.extend_from_slice(HANDSHAKE_INFO);
    Ok(Sha3_256::digest(transcript).into())
}

/// Asserts that the size of a type is equal to `size`.
#[macro_export]
macro_rules! assert_size {
    ($ty:ty, $expected:expr $(,)?) => {
        const _: () = assert!(
            ::core::mem::size_of::<$ty>() == $expected,
            concat!(
                "size of `",
                stringify!($ty),
                "` is not ",
                stringify!($expected),
                " bytes"
            ),
        );
    };
}

// ---

/// Header sent by the client to the server during the join/role determine process.
#[repr(C)]
pub struct JoinHeader {
    pub role: u8,
    // for Role::Producer/Consumer, this is session_uuid
    // for Role::Observer, this is observe_token
    pub uuid_or_token: [u8; 64],
}

impl JoinHeader {
    #[inline(always)]
    pub fn as_bytes(&self) -> &[u8] {
        assert_size!(JoinHeader, 65);
        unsafe { std::slice::from_raw_parts(self as *const _ as *const u8, size_of::<Self>()) }
    }

    #[inline(always)]
    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() < size_of::<Self>() {
            bail!(
                "Invalid JoinHeader length: expected {}, got {}",
                size_of::<Self>(),
                bytes.len()
            );
        }
        let session_uuid: [u8; 64] = bytes[1..65].try_into()?;
        Ok(Self {
            role: bytes[0],
            uuid_or_token: session_uuid,
        })
    }
}

pub async fn read_join_header<R: AsyncReadExt + Unpin>(
    reader: &mut R,
    transport_key: &[u8; 32],
    decrypt_buf: &mut BytesMut,
) -> Result<JoinHeader> {
    let read_len = NONCE_LEN + size_of::<JoinHeader>() + TAG_LEN;
    let decrypted_buf = read_decrypt_packet(reader, read_len, transport_key, decrypt_buf).await?;
    JoinHeader::from_bytes(decrypted_buf)
}

pub async fn write_join_header<W: AsyncWriteExt + Unpin>(
    writer: &mut W,
    role: Role,
    uuid_or_token: [u8; 64],
    transport_key: &[u8; 32],
    encrypt_buf: &mut BytesMut,
) -> Result<()> {
    let join_header = JoinHeader {
        role: role.to_u8(),
        uuid_or_token,
    };
    write_encrypt_packet(writer, join_header.as_bytes(), transport_key, encrypt_buf).await
}

// ---

// #[repr(C)]
// pub struct PacketHeader {
//     pub seq: u128,
//     pub timestamp: u128,
// }
// pub const PACKET_HEADER_SIZE_OVERHEAD: usize = size_of::<PacketHeader>();
// impl PacketHeader {
//     #[inline(always)]
//     pub fn as_bytes(&self) -> &[u8] {
//         assert_size!(PacketHeader, 32);
//         unsafe { std::slice::from_raw_parts(self as *const _ as *const u8, size_of::<Self>()) }
//     }
//
//     #[inline(always)]
//     pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
//         if bytes.len() < size_of::<Self>() {
//             bail!(
//                 "Invalid PacketHeader length: expected {}, got {}",
//                 size_of::<Self>(),
//                 bytes.len()
//             );
//         }
//         let seq = u128::from_le_bytes(bytes[0..16].try_into()?);
//         let timestamp = u128::from_le_bytes(bytes[16..32].try_into()?);
//         Ok(Self { seq, timestamp })
//     }
// }

#[repr(u8)]
#[derive(Debug)]
pub enum StatusCode {
    Success,
    ErrInvalidRole,
    ErrSessionNotFound,
    ErrSessionAlreadyOccupied,
    ErrInvalidToken,
    ErrInternalServerError,
}

impl StatusCode {
    #[inline(always)]
    pub fn to_u8(&self) -> u8 {
        match self {
            StatusCode::Success => 0x00,
            StatusCode::ErrInvalidRole => 0x01,
            StatusCode::ErrSessionNotFound => 0x02,
            StatusCode::ErrSessionAlreadyOccupied => 0x03,
            StatusCode::ErrInvalidToken => 0x04,
            StatusCode::ErrInternalServerError => 0x05,
        }
    }
    #[inline(always)]
    pub fn from_u8(val: u8) -> Result<Self> {
        match val {
            0x00 => Ok(StatusCode::Success),
            0x01 => Ok(StatusCode::ErrInvalidRole),
            0x02 => Ok(StatusCode::ErrSessionNotFound),
            0x03 => Ok(StatusCode::ErrSessionAlreadyOccupied),
            0x04 => Ok(StatusCode::ErrInvalidToken),
            0x05 => Ok(StatusCode::ErrInternalServerError),
            _ => bail!("Unknown Status code: {:#04x}", val),
        }
    }
}

// ---

#[inline]
pub async fn read_status_code<R: AsyncReadExt + Unpin>(
    reader: &mut R,
    session_key: &[u8; 32],
    scratch_buf: &mut BytesMut,
) -> Result<StatusCode> {
    let read_len = NONCE_LEN + 1 + TAG_LEN;
    let status_code_bytes = read_decrypt_packet(reader, read_len, session_key, scratch_buf).await?;
    StatusCode::from_u8(status_code_bytes[0])
}

#[inline]
pub async fn write_status_code<W: AsyncWriteExt + Unpin>(
    writer: &mut W,
    status_code: StatusCode,
    session_key: &[u8; 32],
    scratch_buf: &mut BytesMut,
) -> Result<()> {
    let status_code_bytes = [status_code.to_u8()];
    write_encrypt_packet(writer, &status_code_bytes, session_key, scratch_buf).await
}

// ---

#[inline]
/// Decrypt & read the server's read/write length from the provided reader.
pub async fn read_server_read_write_len<R: AsyncReadExt + Unpin>(
    reader: &mut R,
    session_key: &[u8; 32],
    scratch_buf: &mut BytesMut,
) -> Result<u64> {
    let read_len = NONCE_LEN + size_of::<u64>() + TAG_LEN;
    let capacity_bytes = read_decrypt_packet(reader, read_len, session_key, scratch_buf).await?;
    let read_write_capacity = u64::from_le_bytes(capacity_bytes.try_into()?);
    Ok(read_write_capacity)
}

#[inline]
/// Encrypt & write the server's read/write length to the provided writer.
pub async fn write_server_read_write_len<W: AsyncWriteExt + Unpin>(
    writer: &mut W,
    read_write_capacity: u64,
    session_key: &[u8; 32],
    scratch_buf: &mut BytesMut,
) -> Result<()> {
    let capacity_bytes: [u8; 8] = read_write_capacity.to_le_bytes();
    write_encrypt_packet(writer, &capacity_bytes, session_key, scratch_buf).await
}

// ---

impl ServerMetrics {
    #[inline(always)]
    pub fn as_bytes(&self) -> &[u8] {
        assert_size!(ServerMetrics, 32);
        unsafe { std::slice::from_raw_parts(self as *const _ as *const u8, size_of::<Self>()) }
    }

    #[inline(always)]
    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() < size_of::<Self>() {
            bail!(
                "Invalid ServerMetrics length: expected {}, got {}",
                size_of::<Self>(),
                bytes.len()
            );
        }
        let uptime_hrs = f64::from_le_bytes(bytes[0..8].try_into()?);
        let total_sessions = u64::from_le_bytes(bytes[8..16].try_into()?);
        let active_sessions = u64::from_le_bytes(bytes[16..24].try_into()?);
        let total_network_bandwidth = u64::from_le_bytes(bytes[24..32].try_into()?);
        Ok(Self {
            uptime_hrs,
            total_sessions,
            active_sessions,
            total_network_bandwidth,
        })
    }
}

#[inline]
pub async fn read_server_metrics<R: AsyncReadExt + Unpin>(
    reader: &mut R,
    transport_key: &[u8; 32],
    scratch_buf: &mut BytesMut,
) -> Result<ServerMetrics> {
    let read_len = NONCE_LEN + size_of::<ServerMetrics>() + TAG_LEN;
    let metrics_bytes = read_decrypt_packet(reader, read_len, transport_key, scratch_buf).await?;
    ServerMetrics::from_bytes(metrics_bytes)
}

#[inline]
pub async fn write_server_metrics<W: AsyncWriteExt + Unpin>(
    writer: &mut W,
    metrics: &ServerMetrics,
    session_key: &[u8; 32],
    scratch_buf: &mut BytesMut,
) -> Result<()> {
    write_encrypt_packet(writer, metrics.as_bytes(), session_key, scratch_buf).await
}

// ---

/// Decrypt & read an encrypted packet from the provided reader, and returns a reference to the decrypted data `(plaintext)`.
/// The `read_len` is the total `length of the encrypted packet`, which must be > `NONCE_LEN` + `TAG_LEN`.
pub async fn read_decrypt_packet<'a, R: AsyncReadExt + Unpin>(
    reader: &mut R,
    read_len: usize,
    encrypt_key: &[u8; 32],
    scratch_buf: &'a mut BytesMut,
) -> Result<&'a [u8]> {
    if read_len < NONCE_LEN + TAG_LEN {
        bail!("Data length too short for decryption: {}", read_len);
    }

    let plaintext_len = read_len - NONCE_LEN - TAG_LEN;
    let required_len = read_len + plaintext_len;

    // should accommodate both encrypted & decrypted chunk
    if scratch_buf.len() < required_len {
        scratch_buf.resize(required_len, 0);
    }

    let (encrypted_chunk, decrypted_chunk) = scratch_buf.split_at_mut(read_len);
    let encrypted_buf = &mut encrypted_chunk[..read_len];
    reader.read_exact(encrypted_buf).await?;

    let decrypted_buf = &mut decrypted_chunk[..plaintext_len];
    decrypt_into(encrypted_buf, decrypted_buf, encrypt_key)?;

    Ok(decrypted_buf)
}

/// Encrypts the given data and writes it to the provided writer.
/// The encrypted packet's len is `NONCE_LEN` + `data.len()` + `TAG_LEN`.
pub async fn write_encrypt_packet<W: AsyncWriteExt + Unpin>(
    writer: &mut W,
    data: &[u8],
    encrypt_key: &[u8; 32],
    scratch_buf: &mut BytesMut,
) -> Result<()> {
    let plaintext_len = data.len();
    let ciphertext_len = NONCE_LEN + plaintext_len + TAG_LEN;

    // alloc if less
    if scratch_buf.len() < ciphertext_len {
        scratch_buf.resize(ciphertext_len, 0);
    }

    // encrypt that son of a bitch
    encrypt_into(data, &mut scratch_buf[..ciphertext_len], encrypt_key)?;
    writer.write_all(&scratch_buf[..ciphertext_len]).await?;
    writer.flush().await?;
    Ok(())
}