p2p-beaver / src / message_protocol.rs

/* SPDX Id: AGPL-3.0-or-later */

use iroh::{
    endpoint::Connection,
    protocol::{AcceptError, ProtocolHandler},
};
use log::{error, info};

use crate::{PeerEvent, packet::BasePacket, state::SharedState};

#[derive(Debug)]
pub(crate) struct MessageProtocol {
    state: SharedState,
}

impl MessageProtocol {
    pub(crate) fn new(state: SharedState) -> Self {
        Self { state }
    }
}

/// Message protocol.
/// Both sides can send messages back and forth with
/// no pre-determined pattern.
impl ProtocolHandler for MessageProtocol {
    async fn accept(&self, connection: Connection) -> Result<(), AcceptError> {
        let remote_id = connection.remote_id();
        println!(
            "accepted connection on {:?} from {:?}",
            String::from_utf8_lossy(connection.alpn()),
            remote_id
        );

        // 1. Accept the bidirectional stream.
        let (sender, mut receiver) = connection.accept_bi().await?;

        // 2. Store the sender in the state for that remote ID.
        self.state
            .lock()
            .await
            .set_message_sender(&remote_id, sender);

        // 3. Relay all the incoming messages.
        loop {
            match BasePacket::recv(&mut receiver).await {
                Ok(payload) => {
                    self.state
                        .lock()
                        .await
                        .notify(PeerEvent::Message(remote_id, payload));
                }
                Err(err) => {
                    error!("Error reading message packet: {err}");
                    break;
                }
            }
        }

        // 4. Done with that connection?
        self.state.lock().await.remove_message_sender(&remote_id);

        Ok(())
    }
}