geth/crates/geth-node/src/wire.rs

//! Shared Iroh stream helpers.

use crate::NodeError;
use std::time::Duration;

pub(crate) const PEER_CONTROL_LINE_MAX: usize = 16 * 1024 * 1024;
pub(crate) const WIRE_REQUEST_LINE_MAX: usize = 16 * 1024 * 1024;
pub(crate) const STREAM_HANDSHAKE_LINE_MAX: usize = 64 * 1024;
pub(crate) const IROH_LINE_READ_TIMEOUT: Duration = Duration::from_secs(30);

pub(crate) async fn read_iroh_line(
    recv: &mut iroh::endpoint::RecvStream,
    max_len: usize,
) -> Result<String, NodeError> {
    tokio::time::timeout(IROH_LINE_READ_TIMEOUT, read_iroh_line_inner(recv, max_len))
        .await
        .map_err(|_| {
            NodeError::IrohPeer(format!(
                "iroh line read timed out after {} seconds",
                IROH_LINE_READ_TIMEOUT.as_secs()
            ))
        })?
}

async fn read_iroh_line_inner(
    recv: &mut iroh::endpoint::RecvStream,
    max_len: usize,
) -> Result<String, NodeError> {
    let mut bytes = Vec::new();
    let mut byte = [0_u8; 1];
    loop {
        let Some(n) = recv
            .read(&mut byte)
            .await
            .map_err(|error| NodeError::IrohPeer(error.to_string()))?
        else {
            if bytes.is_empty() {
                return Err(NodeError::IrohPeer(
                    "iroh stream closed before response line".to_owned(),
                ));
            }
            break;
        };
        if n == 0 {
            continue;
        }
        if push_line_byte(&mut bytes, byte[0], max_len)? {
            break;
        }
    }
    String::from_utf8(bytes).map_err(|error| NodeError::IrohPeer(error.to_string()))
}

fn push_line_byte(bytes: &mut Vec<u8>, byte: u8, max_len: usize) -> Result<bool, NodeError> {
    if bytes.len() >= max_len {
        return Err(NodeError::IrohPeer(format!(
            "iroh line exceeded {max_len} bytes"
        )));
    }
    bytes.push(byte);
    Ok(byte == b'\n')
}

pub(crate) async fn finish_iroh_send(
    send: &mut iroh::endpoint::SendStream,
) -> Result<(), NodeError> {
    send.finish()
        .map_err(|error| NodeError::IrohPeer(error.to_string()))?;
    send.stopped()
        .await
        .map_err(|error| NodeError::IrohPeer(error.to_string()))?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn bounded_line_rejects_byte_past_limit_without_appending() {
        let mut bytes = b"abc".to_vec();
        let error = push_line_byte(&mut bytes, b'd', 3).expect_err("line should exceed limit");

        assert_eq!(bytes, b"abc");
        assert!(error.to_string().contains("iroh line exceeded 3 bytes"));
    }

    #[test]
    fn bounded_line_accepts_newline_within_limit() {
        let mut bytes = b"abc".to_vec();

        assert!(push_line_byte(&mut bytes, b'\n', 4).expect("push newline"));
        assert_eq!(bytes, b"abc\n");
    }
}