lanspread/crates/lanspread-peer/src/peer.rs

use std::{convert::TryInto, path::Path};

use bytes::Bytes;
use lanspread_db::db::GameFileDescription;
use lanspread_utils::maybe_addr;
use s2n_quic::{
    connection,
    stream::{Error as StreamError, SendStream},
};
use tokio::{
    io::{AsyncReadExt, AsyncSeekExt},
    time::Instant,
};

use crate::{config::FILE_TRANSFER_BUFFER_SIZE, path_validation::validate_game_file_path};

async fn stream_file_bytes(
    tx: &mut SendStream,
    base_dir: &Path,
    relative_path: &str,
    offset: u64,
    length: Option<u64>,
) -> eyre::Result<()> {
    let remote_addr = maybe_addr!(tx.connection().remote_addr());

    // Validate the path to prevent directory traversal
    let validated_path = validate_game_file_path(base_dir, relative_path)?;
    log::debug!(
        "{remote_addr} streaming file bytes for peer: {}, offset: {offset}, length: {length:?}",
        validated_path.display()
    );

    let mut file = tokio::fs::File::open(&validated_path).await?;
    if offset > 0 {
        file.seek(std::io::SeekFrom::Start(offset)).await?;
    }

    let mut remaining = length.unwrap_or(u64::MAX);
    let expect_exact = length.is_some();
    let mut transfer_complete = matches!(length, Some(0));
    let mut total_bytes = 0u64;
    let mut last_total_bytes = 0u64;
    let started = Instant::now();
    let mut timestamp = Instant::now();
    let mut buf = vec![0u8; FILE_TRANSFER_BUFFER_SIZE];

    while remaining > 0 {
        let read_len = std::cmp::min(remaining, buf.len() as u64);
        let read_len: usize = read_len.try_into().unwrap_or(usize::MAX);
        if read_len == 0 {
            break;
        }

        let bytes_read = file.read(&mut buf[..read_len]).await?;
        if bytes_read == 0 {
            if !expect_exact {
                transfer_complete = true;
            }
            break;
        }

        tx.send(Bytes::copy_from_slice(&buf[..bytes_read])).await?;
        remaining = remaining.saturating_sub(bytes_read as u64);
        total_bytes += bytes_read as u64;

        if expect_exact && remaining == 0 {
            transfer_complete = true;
            break;
        }

        if last_total_bytes + 10_000_000 < total_bytes {
            let elapsed = timestamp.elapsed();
            let diff_bytes = total_bytes - last_total_bytes;

            if elapsed.as_secs_f64() >= 1.0 {
                #[allow(clippy::cast_precision_loss)]
                let mb_per_s = (diff_bytes as f64) / (elapsed.as_secs_f64() * 1_000_000.0);
                log::debug!(
                    "{remote_addr} sending file data: {}, MB/s: {mb_per_s:.2}",
                    validated_path.display()
                );
                last_total_bytes = total_bytes;
                timestamp = Instant::now();
            }
        }
    }

    let elapsed = started.elapsed();
    #[allow(clippy::cast_precision_loss)]
    let mib_per_s = if elapsed.as_secs_f64() > 0.0 {
        total_bytes as f64 / elapsed.as_secs_f64() / (1024.0 * 1024.0)
    } else {
        0.0
    };
    log::info!(
        "{remote_addr} finished streaming file bytes: {}, total_bytes: {total_bytes}, MiB/s: {mib_per_s:.2}",
        validated_path.display()
    );

    match tx.close().await {
        Ok(()) => {}
        Err(err) if transfer_complete && is_clean_remote_close(&err) => {
            log::debug!("{remote_addr} closed stream after transfer completion: {err}");
        }
        Err(err) => return Err(err.into()),
    }
    Ok(())
}

fn is_clean_remote_close(err: &StreamError) -> bool {
    matches!(
        err,
        StreamError::ConnectionError {
            error: connection::Error::Closed { .. },
            ..
        }
    )
}

pub async fn send_game_file_data(
    game_file_desc: &GameFileDescription,
    tx: &mut SendStream,
    game_dir: &Path,
) {
    if let Err(e) = stream_file_bytes(tx, game_dir, &game_file_desc.relative_path, 0, None).await {
        let remote_addr = maybe_addr!(tx.connection().remote_addr());
        log::error!(
            "{remote_addr} failed to stream file {}: {e}",
            game_file_desc.relative_path
        );
    }
}

pub async fn send_game_file_chunk(
    game_id: &str,
    relative_path: &str,
    offset: u64,
    length: u64,
    tx: &mut SendStream,
    game_dir: &Path,
) {
    if let Err(e) = stream_file_bytes(tx, game_dir, relative_path, offset, Some(length)).await {
        let remote_addr = maybe_addr!(tx.connection().remote_addr());
        log::error!(
            "{remote_addr} failed to stream chunk {game_id}/{relative_path} offset {offset} length {length}: {e}"
        );
    }
}