cuprate/p2p/monero-p2p/src/handles.rs

//! Connection Handles.
//!
//! This module contains the [`ConnectionHandle`] which allows banning a peer, disconnecting a peer and
//! checking if the peer is still connected.
use std::{
    sync::{Arc, OnceLock},
    time::Duration,
};

use futures::SinkExt;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use tokio_util::sync::CancellationToken;

/// A [`ConnectionHandle`] builder.
#[derive(Default, Debug)]
pub struct HandleBuilder {
    permit: Option<OwnedSemaphorePermit>,
}

impl HandleBuilder {
    /// Create a new builder.
    pub fn new() -> Self {
        Self { permit: None }
    }

    /// Sets the permit for this connection.
    ///
    /// This must be called at least once.
    pub fn with_permit(mut self, permit: OwnedSemaphorePermit) -> Self {
        self.permit = Some(permit);
        self
    }

    /// Builds the [`ConnectionGuard`] which should be handed to the connection task and the [`ConnectionHandle`].
    ///
    /// This will panic if a permit was not set [`HandleBuilder::with_permit`]
    pub fn build(self) -> (ConnectionGuard, ConnectionHandle) {
        let token = CancellationToken::new();

        (
            ConnectionGuard {
                token: token.clone(),
                permit: self.permit.expect("connection permit was not set!"),
            },
            ConnectionHandle {
                token: token.clone(),
                ban: Arc::new(OnceLock::new()),
            },
        )
    }
}

/// A struct representing the time a peer should be banned for.
#[derive(Debug, Copy, Clone)]
pub struct BanPeer(pub Duration);

/// A struct given to the connection task.
pub struct ConnectionGuard {
    token: CancellationToken,
    permit: OwnedSemaphorePermit,
}

impl ConnectionGuard {
    /// Checks if we should close the connection.
    pub fn should_shutdown(&self) -> bool {
        self.token.is_cancelled()
    }
    /// Tell the corresponding [`ConnectionHandle`]s that this connection is closed.
    ///
    /// This will be called on [`Drop::drop`].
    pub fn connection_closed(&self) {
        self.token.cancel()
    }
}

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        self.token.cancel()
    }
}

/// A handle given to a task that needs to ban, disconnect, check if the peer should be banned or check
/// the peer is still connected.
#[derive(Debug, Clone)]
pub struct ConnectionHandle {
    token: CancellationToken,
    ban: Arc<OnceLock<BanPeer>>,
}

impl ConnectionHandle {
    /// Bans the peer for the given `duration`.
    pub fn ban_peer(&self, duration: Duration) {
        let _ = self.ban.set(BanPeer(duration));
    }
    /// Checks if this connection is closed.
    pub fn is_closed(&self) -> bool {
        self.token.is_cancelled()
    }
    /// Returns if this peer has been banned and the [`Duration`] of that ban.
    pub fn check_should_ban(&mut self) -> Option<BanPeer> {
        self.ban.get().copied()
    }
    /// Sends the signal to the connection task to disconnect.
    pub fn send_close_signal(&self) {
        self.token.cancel()
    }
}