Something went wrong. Try again.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272use std::{ collections::BTreeMap, sync::{Arc, Mutex},};
use iroh::{ EndpointAddr, endpoint::Connection, protocol::{AcceptError, ProtocolHandler},};use tokio::sync::RwLock;
use crate::{ domain::FolderId, iroh::{ CONTROL_REQUEST_TIMEOUT, DiscoveredPeer, HostedFolder, MAX_AUDIT_EVENTS_PER_RESPONSE, MAX_CONTROL_MESSAGE_BYTES, recover_mutex, }, protocol::{ControlMessage, ManifestRequestKind, ManifestSummary, validate_invite},};
#[derive(Clone, Debug)]pub(super) struct AppaProtocol { pub(super) folders: Arc<RwLock<BTreeMap<FolderId, HostedFolder>>>, pub(super) discovered_peers: Arc<Mutex<BTreeMap<(FolderId, iroh::EndpointId), DiscoveredPeer>>>, pub(super) announced_folders: Arc<Mutex<BTreeMap<FolderId, String>>>,}
impl ProtocolHandler for AppaProtocol { async fn accept(&self, connection: Connection) -> Result<(), AcceptError> { while let Ok((mut send, mut receive)) = connection.accept_bi().await { if let Err(error) = self .handle_request(&connection, &mut send, &mut receive) .await { tracing::warn!(peer = %connection.remote_id(), error = %error, "Rejected Appa control request"); } } Ok(()) }}
impl AppaProtocol { async fn handle_request( &self, connection: &Connection, send: &mut iroh::endpoint::SendStream, receive: &mut iroh::endpoint::RecvStream, ) -> Result<(), AcceptError> { let request = read_control_request(receive).await?; let response = self.request_response(connection, request).await?; send_response(send, &response).await }
async fn request_response( &self, connection: &Connection, request: ControlMessage, ) -> Result<ControlMessage, AcceptError> { match request { ControlMessage::Request { folder_id, capability, invite, requester_endpoint, request: request_kind, } => { self.handle_folder_request( connection, folder_id, capability, invite, requester_endpoint, request_kind, ) .await } ControlMessage::Announcement { folder_id, capability, root_hash, sender_endpoint, } => { self.handle_announcement( connection, folder_id, capability, root_hash, sender_endpoint, ) .await } _ => Err(protocol_error("unsupported Appa control message")), } }
async fn handle_folder_request( &self, connection: &Connection, folder_id: FolderId, capability: String, invite: Option<Box<crate::protocol::Invite>>, requester_endpoint: EndpointAddr, request_kind: ManifestRequestKind, ) -> Result<ControlMessage, AcceptError> { verify_remote_identity(connection, &requester_endpoint, "requester")?; let response = self .authorize_request( folder_id, &capability, invite.as_deref(), &requester_endpoint, request_kind, ) .await?; recover_mutex(&self.discovered_peers).insert( (folder_id, requester_endpoint.id), DiscoveredPeer { folder_id, endpoint: requester_endpoint, }, ); Ok(response) }
async fn handle_announcement( &self, connection: &Connection, folder_id: FolderId, capability: String, root_hash: String, sender_endpoint: EndpointAddr, ) -> Result<ControlMessage, AcceptError> { verify_remote_identity(connection, &sender_endpoint, "announcer")?; self.authorize_announcement(folder_id, &capability, &sender_endpoint) .await?; recover_mutex(&self.announced_folders).insert(folder_id, root_hash); tracing::debug!(peer = %connection.remote_id(), folder = %folder_id, "Received folder change announcement"); Ok(ControlMessage::AnnouncementAccepted) }}
impl AppaProtocol { async fn authorize_announcement( &self, folder_id: FolderId, capability: &str, sender: &EndpointAddr, ) -> Result<(), AcceptError> { let folders = self.folders.read().await; let folder = folders .get(&folder_id) .ok_or_else(|| protocol_error("folder is not hosted by this Appa node"))?; if capability != folder.capability || !folder.roster.can_publish(&sender.id.to_string()) { return Err(protocol_error("unauthorized Appa folder announcement")); } Ok(()) }}
impl AppaProtocol { async fn authorize_request( &self, folder_id: FolderId, capability: &str, invite: Option<&crate::protocol::Invite>, requester_endpoint: &EndpointAddr, request_kind: ManifestRequestKind, ) -> Result<ControlMessage, AcceptError> { let folders = self.folders.read().await; let folder = folders .get(&folder_id) .ok_or_else(|| protocol_error("folder is not hosted by this Appa node"))?; let is_current_capability = capability == folder.capability; let is_active_member = folder .active_member_ids .contains(&requester_endpoint.id.to_string()); if !is_current_capability { return Err(protocol_error("invalid Appa folder capability")); } if !is_active_member && !invite.is_some_and(|invite| { validate_invite(invite, time::OffsetDateTime::now_utc()).is_ok() && invite.folder_id == folder_id && invite.capability == folder.capability }) { return Err(protocol_error("requester lacks a valid Appa invitation")); } match request_kind { ManifestRequestKind::Summary => Ok(ControlMessage::Summary { summary: ManifestSummary { folder_id, root_hash: folder.merkle_index.root_hash().to_owned(), }, roster: folder.roster.clone(), peer_endpoints: folder.peer_endpoints.clone(), capability: folder.capability.clone(), }), ManifestRequestKind::MerkleNode { prefix } => Ok(ControlMessage::MerkleNode { node: folder .merkle_index .node(&prefix) .map_err(protocol_error)? .ok_or_else(|| protocol_error("unknown Merkle prefix"))?, }), ManifestRequestKind::AuditEvents { after_sequences } => { Ok(ControlMessage::AuditEvents { events: folder .audit_events .iter() .filter(|event| { event.sequence > after_sequences .get(&event.author_device_id) .copied() .unwrap_or_default() }) .take(MAX_AUDIT_EVENTS_PER_RESPONSE) .cloned() .collect(), }) } } }}
async fn read_control_request( receive: &mut iroh::endpoint::RecvStream,) -> Result<ControlMessage, AcceptError> { tokio::time::timeout(CONTROL_REQUEST_TIMEOUT, read_request_payload(receive)) .await .map_err(|_| protocol_error("Appa control request timed out"))?}
async fn read_request_payload( receive: &mut iroh::endpoint::RecvStream,) -> Result<ControlMessage, AcceptError> { let payload = receive .read_to_end(MAX_CONTROL_MESSAGE_BYTES) .await .map_err(protocol_error)?; serde_json::from_slice(&payload).map_err(protocol_error)}
async fn send_response( send: &mut iroh::endpoint::SendStream, response: &ControlMessage,) -> Result<(), AcceptError> { let payload = serde_json::to_vec(response).map_err(protocol_error)?; send.write_all(&payload).await.map_err(protocol_error)?; send.finish()?; Ok(())}
fn verify_remote_identity( connection: &Connection, endpoint: &EndpointAddr, peer_role: &str,) -> Result<(), AcceptError> { if endpoint.id == connection.remote_id() { return Ok(()); } Err(protocol_error(format!( "Appa {peer_role} identity does not match its endpoint" )))}
fn protocol_error(error: impl std::fmt::Display) -> AcceptError { AcceptError::from_err(std::io::Error::other(error.to_string()))}