matroskin 0.1.0

//! TCP connection establishment and salt/permission fetch.
//!
//! Handshake runs inside the actor task after a connect intent or reconnect.

use futures_util::{SinkExt, StreamExt};
use tokio::{
    net::TcpStream,
    time::{self, Duration},
};
use tokio_util::codec::Framed;
use tracing::{debug, info, instrument};

use super::state::Salt;
use crate::{
    account::Account,
    actor::codec::{WhatsMinerCodec, WhatsMinerFrame},
    command::{Command, get_device_info::GetDeviceInfo},
    error::{Error, Result},
    permission::Permission,
};

pub(super) const IO_TIMEOUT: Duration = Duration::from_millis(300);

/// Establish a TCP connection to a WhatsMiner device and fetch the session salt.
///
/// This function performs:
/// 1. TCP connect to `addr`.
/// 2. Set `TCP_NODELAY` for low-latency communication.
/// 3. Send `get.device.info` with `salt: true` to obtain the session salt and permissions.
///
/// Returns the `Framed` transport, the salt string, and the parsed permission.
pub(super) async fn connect_and_get_salt(
    addr: &str,
    account: Account,
) -> Result<(Framed<TcpStream, WhatsMinerCodec>, Salt, Permission)> {
    debug!(addr = %addr, "Connecting TCP stream.");
    let stream = time::timeout(IO_TIMEOUT, TcpStream::connect(addr))
        .await
        .map_err(|_| Error::Timeout)??;
    debug!(addr = %addr, "TCP stream connected. Setting no delay.");
    stream.set_nodelay(true)?;
    let mut framed = Framed::new(stream, WhatsMinerCodec);
    let (salt, permission) = get_actor_salt_and_permission(&mut framed, account).await?;
    Ok((framed, salt, permission))
}

/// Fetch the session salt and permissions from the device using `get.device.info`.
///
/// Salt and miner permissions are requested separately so the handshake does
/// not depend on support for combined section filters in a firmware revision.
#[instrument(level = "info", skip(framed))]
pub(super) async fn get_actor_salt_and_permission(
    framed: &mut Framed<TcpStream, WhatsMinerCodec>,
    account: Account,
) -> Result<(Salt, Permission)> {
    info!("Attempting to retrieve actor salt and permissions.");

    // Step 1: Request Salt
    let cmd_salt = GetDeviceInfo {
        salt: true,
        miner: false,
        power: false,
        network: false,
        system: false,
        error_code: false,
    };
    let mut encoded_salt = cmd_salt.encode(None)?;
    let message_salt = std::mem::take(&mut encoded_salt.0);
    time::timeout(
        IO_TIMEOUT,
        framed.send(WhatsMinerFrame::Data(bytes::Bytes::from(message_salt))),
    )
    .await
    .map_err(|_| Error::Timeout)??;

    let frame_salt = read_command_response_frame(framed, GetDeviceInfo::CMD_NAME).await?;

    let salt = {
        let data = GetDeviceInfo::parse_response(&frame_salt)?;
        let mut salt_redacted = data.msg.salt.ok_or(Error::SaltNotFound)?;
        let salt_str = std::mem::take(&mut salt_redacted.0);
        Salt::new(salt_str)
    };

    // Step 2: Request Permissions (Miner section)
    let cmd_miner = GetDeviceInfo {
        miner: true,
        salt: false,
        power: false,
        network: false,
        system: false,
        error_code: false,
    };
    let mut encoded_miner = cmd_miner.encode(None)?;
    let message_miner = std::mem::take(&mut encoded_miner.0);
    time::timeout(
        IO_TIMEOUT,
        framed.send(WhatsMinerFrame::Data(bytes::Bytes::from(message_miner))),
    )
    .await
    .map_err(|_| Error::Timeout)??;

    let frame_miner = read_command_response_frame(framed, GetDeviceInfo::CMD_NAME).await?;

    let permission = {
        let data = GetDeviceInfo::parse_response(&frame_miner)?;
        let permission_str = data
            .msg
            .miner
            .as_ref()
            .and_then(|m| m.permission.as_ref())
            .ok_or(Error::PermissionsNotFound)?;
        Permission::parse_for_account(permission_str, account)
    };

    Ok((salt, permission))
}

pub(super) async fn read_command_response_frame(
    framed: &mut Framed<TcpStream, WhatsMinerCodec>,
    expected_cmd: &str,
) -> Result<bytes::Bytes> {
    let deadline = time::Instant::now() + IO_TIMEOUT;
    loop {
        let frame = time::timeout_at(deadline, framed.next())
            .await
            .map_err(|_| Error::Timeout)?
            .ok_or(Error::ConnectionClosed)??;
        match frame {
            WhatsMinerFrame::Data(bytes) => {
                let envelope: crate::response::Envelope<'_> = serde_json::from_slice(&bytes)?;
                if envelope.desc == "get.miner.report" {
                    debug!(expected_cmd, "Drained unsolicited miner report");
                    continue;
                }
                if envelope.desc != expected_cmd {
                    return Err(Error::Protocol(format!(
                        "expected {expected_cmd}, received {}",
                        envelope.desc
                    )));
                }
                return Ok(bytes);
            }
            WhatsMinerFrame::Heartbeat => {
                debug!("Heartbeat received while awaiting data frame, continuing.");
                continue;
            }
        }
    }
}