stratum-apps 0.9.0

//! A Noise-encrypted wrapper around a `TcpStream`, providing framed read/write I/O using the SV2
//! protocol and a stateful Noise handshake.
//!
//! This module provides `NoiseTcpStream`, which wraps a `TcpStream` and performs a Noise-based
//! authenticated key exchange, as the initiator ([`NoiseTcpStream::connect`]) or as the responder
//! ([`NoiseTcpStream::accept`]).
//!
//! After a successful handshake, the stream can be split into a `NoiseTcpReadHalf` and
//! `NoiseTcpWriteHalf`, which support frame-based encoding/decoding of SV2 messages with optional
//! non-blocking behavior.

use std::time::Duration;

use crate::network_helpers::Error;
use stratum_core::{
    codec_sv2::{
        Decoded, Decrypted, EncodableFrame, Handshake, HandshakeMessage, NoiseDecoder,
        NoiseEncoder, Transport, TransportDecryptState, TransportEncryptState,
        state::{ExpectsHandshakeMessage, InitiatorSent},
    },
    noise_sv2::{AeadError, INITIATOR_EXPECTED_HANDSHAKE_MESSAGE_SIZE, Initiator, Responder},
};
use tokio::net::{
    TcpStream,
    tcp::{OwnedReadHalf, OwnedWriteHalf},
};

use stratum_core::{codec_sv2::SerializedFrame, noise_sv2::ELLSWIFT_ENCODING_SIZE};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tracing::debug;

/// A Noise-secured duplex stream over TCP that wraps a `TcpStream`
/// and provides secure read/write capabilities using the Noise protocol.
///
/// This stream performs the full Noise handshake during construction
/// and returns a bidirectional encrypted stream split into read and write halves.
///
/// **Note:** This struct is **not cancellation-safe**.
/// If `read_frame()` or `write_frame()` is canceled mid-way,
/// internal state may be left in an inconsistent state, which can lead to
/// protocol errors or dropped frames.
pub struct NoiseTcpStream {
    reader: NoiseTcpReadHalf,
    writer: NoiseTcpWriteHalf,
}

/// The reading half of a `NoiseTcpStream`.
///
/// It buffers incoming encrypted bytes, attempts to decode full Noise frames,
/// and exposes a method to retrieve structured messages of type `Message`.
pub struct NoiseTcpReadHalf {
    reader: OwnedReadHalf,
    decoder: NoiseDecoder,
    state: Option<TransportDecryptState>,
}

/// The writing half of a `NoiseTcpStream`.
///
/// It accepts structured messages, encodes them via the Noise protocol,
/// and writes the result to the socket.
pub struct NoiseTcpWriteHalf {
    writer: OwnedWriteHalf,
    encoder: NoiseEncoder,
    state: TransportEncryptState,
}

impl NoiseTcpStream {
    /// Connects as the Noise initiator over the given TCP stream, performing the handshake.
    ///
    /// On success, returns a stream with encrypted communication channels.
    ///
    /// `timeout` applies to each individual handshake read. Prefer [`super::connect_with_noise`]
    /// for typical use, which applies a sensible default timeout automatically.
    pub async fn connect(
        stream: TcpStream,
        initiator: Box<Initiator>,
        timeout: Duration,
    ) -> Result<Self, Error> {
        let (mut reader, mut writer) = stream.into_split();
        let mut decoder = NoiseDecoder::new();
        let mut encoder = NoiseEncoder::new();
        let handshake = Handshake::initiator(initiator);

        let (first_msg, handshake) = handshake.step_0()?;
        send_handshake_frame(&mut writer, first_msg, &mut encoder).await?;
        debug!("First handshake message sent");

        let second_msg =
            receive_handshake_frame::<InitiatorSent>(&mut reader, &mut decoder, timeout).await?;
        debug!("Second handshake message received");
        let payload: [u8; INITIATOR_EXPECTED_HANDSHAKE_MESSAGE_SIZE] = second_msg
            .payload()
            .try_into()
            .map_err(|_| Error::HandshakeRemoteInvalidMessage)?;

        Ok(Self::from_transport(
            reader,
            writer,
            decoder,
            encoder,
            handshake.step_2(payload)?,
        ))
    }

    /// Accepts a connection as the Noise responder over the given TCP stream, performing the
    /// handshake.
    ///
    /// On success, returns a stream with encrypted communication channels.
    ///
    /// `timeout` applies to each individual handshake read. Prefer
    /// [`super::accept_noise_connection`] for typical use, which applies a sensible default
    /// timeout automatically.
    pub async fn accept(
        stream: TcpStream,
        responder: Box<Responder>,
        timeout: Duration,
    ) -> Result<Self, Error> {
        let (mut reader, mut writer) = stream.into_split();
        let mut decoder = NoiseDecoder::new();
        let mut encoder = NoiseEncoder::new();
        let handshake = Handshake::responder(responder);

        let first_msg =
            receive_handshake_frame::<Responder>(&mut reader, &mut decoder, timeout).await?;
        debug!("First handshake message received");
        let payload: [u8; ELLSWIFT_ENCODING_SIZE] = first_msg
            .payload()
            .try_into()
            .map_err(|_| Error::HandshakeRemoteInvalidMessage)?;

        let (second_msg, transport) = handshake.step_1(payload)?;
        send_handshake_frame(&mut writer, second_msg, &mut encoder).await?;
        debug!("Second handshake message sent");

        Ok(Self::from_transport(
            reader, writer, decoder, encoder, transport,
        ))
    }

    // Give each half only the direction it uses: the writer encrypts, the reader decrypts.
    fn from_transport(
        reader: OwnedReadHalf,
        writer: OwnedWriteHalf,
        decoder: NoiseDecoder,
        encoder: NoiseEncoder,
        transport: Transport,
    ) -> Self {
        let (encrypt_state, decrypt_state) = transport.split();

        Self {
            reader: NoiseTcpReadHalf {
                reader,
                decoder,
                state: Some(decrypt_state),
            },
            writer: NoiseTcpWriteHalf {
                writer,
                encoder,
                state: encrypt_state,
            },
        }
    }

    /// Consumes the stream and returns its reader and writer halves.
    pub fn into_split(self) -> (NoiseTcpReadHalf, NoiseTcpWriteHalf) {
        (self.reader, self.writer)
    }
}

impl NoiseTcpWriteHalf {
    /// Encrypts and writes a full message frame to the socket.
    ///
    /// Returns an error if the socket is closed or the message cannot be encoded.
    ///
    /// Not cancellation-safe: A canceled write may cause partial writes or state corruption.
    pub async fn write_frame<F: EncodableFrame>(&mut self, frame: F) -> Result<(), Error> {
        let buf = self.encoder.encode_transport(frame, &mut self.state)?;
        self.writer
            .write_all(buf.as_ref())
            .await
            .map_err(|_| Error::SocketClosed)?;
        Ok(())
    }

    /// Attempts to write a message without blocking.
    ///
    /// Returns:
    /// - `Ok(true)` if the entire frame was written successfully.
    /// - `Ok(false)` if the socket is not ready (would block).
    /// - `Err(_)` on socket or encoding errors.
    pub fn try_write_frame<F: EncodableFrame>(&mut self, frame: F) -> Result<bool, Error> {
        let buf = self.encoder.encode_transport(frame, &mut self.state)?;

        match self.writer.try_write(buf.as_ref()) {
            Ok(n) if n == buf.len() => Ok(true),
            Ok(_) => Err(Error::SocketClosed),
            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => Ok(false),
            Err(_) => Err(Error::SocketClosed),
        }
    }

    /// Gracefully shuts down the writing half of the stream.
    ///
    /// Returns an error if the shutdown fails.
    pub async fn shutdown(&mut self) -> Result<(), Error> {
        self.writer
            .shutdown()
            .await
            .map_err(|_| Error::SocketClosed)
    }
}

impl NoiseTcpReadHalf {
    /// Reads and decodes a complete frame from the socket.
    ///
    /// This method blocks until a full frame is read and decoded,
    /// handling `Incomplete` rounds from the codec automatically.
    ///
    /// Cancellation-safe: the decoder counts bytes only once a read has returned them, so a
    /// cancelled read leaves nothing behind.
    pub async fn read_frame(&mut self) -> Result<SerializedFrame, Error> {
        loop {
            match self.next_frame()? {
                Some(frame) => return Ok(frame),
                None => tokio::task::yield_now().await,
            }

            // An incomplete frame always leaves the decoder's read window non-empty, so a read of
            // 0 bytes means the peer closed the socket.
            let n = self
                .reader
                .read(self.decoder.read_buf())
                .await
                .map_err(|_| Error::SocketClosed)?;

            if n == 0 {
                return Err(Error::SocketClosed);
            }

            self.decoder.advance(n)?;
        }
    }

    /// Attempts to read and decode a frame without blocking.
    ///
    /// Returns:
    /// - `Ok(Some(frame))` if a full frame is successfully decoded.
    /// - `Ok(None)` if not enough data is available yet.
    /// - `Err(_)` on socket or decoding errors.
    pub fn try_read_frame(&mut self) -> Result<Option<SerializedFrame>, Error> {
        if let Some(frame) = self.next_frame()? {
            return Ok(Some(frame));
        }

        // An incomplete frame always leaves the decoder's read window non-empty, so `Ok(0)`
        // means the peer closed the socket.
        match self.reader.try_read(self.decoder.read_buf()) {
            Ok(0) => return Err(Error::SocketClosed),
            Ok(n) => self.decoder.advance(n)?,
            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => return Ok(None),
            Err(_) => return Err(Error::SocketClosed),
        }

        self.next_frame()
    }

    fn next_frame(&mut self) -> Result<Option<SerializedFrame>, Error> {
        let state = self
            .state
            .take()
            .ok_or(stratum_core::codec_sv2::Error::AeadError(AeadError))?;
        match self.decoder.next_transport_frame(state)? {
            Decrypted::Frame(frame, state) => {
                self.state = Some(state);
                Ok(Some(frame))
            }
            Decrypted::Incomplete(state) => {
                self.state = Some(state);
                Ok(None)
            }
        }
    }
}

async fn send_handshake_frame(
    writer: &mut OwnedWriteHalf,
    frame: HandshakeMessage,
    encoder: &mut NoiseEncoder,
) -> Result<(), Error> {
    let buffer = encoder.encode_handshake(frame);
    writer
        .write_all(buffer.as_ref())
        .await
        .map_err(|_| Error::SocketClosed)?;
    Ok(())
}

async fn receive_handshake_frame<R: ExpectsHandshakeMessage>(
    reader: &mut OwnedReadHalf,
    decoder: &mut NoiseDecoder,
    timeout: Duration,
) -> Result<HandshakeMessage, Error> {
    // The whole message shares one timeout, so a peer cannot stretch the handshake by sending
    // it a few bytes at a time.
    let receive = async {
        loop {
            match decoder.next_handshake_frame::<R>() {
                Ok(Decoded::Frame(frame)) => return Ok(frame),
                Ok(Decoded::Incomplete) => {
                    debug!("Waiting for more bytes during handshake");
                }
                Err(e) => return Err(Error::CodecError(e)),
            }

            let n = reader
                .read(decoder.read_buf())
                .await
                .map_err(|_| Error::SocketClosed)?;
            if n == 0 {
                return Err(Error::SocketClosed);
            }
            decoder.advance(n).map_err(Error::CodecError)?;
        }
    };
    tokio::time::timeout(timeout, receive)
        .await
        .map_err(|_| Error::HandshakeTimeout)?
}