diff --git a/src/iroh/handler.rs b/src/iroh/handler.rs index efcaa65..5b8f505 100644 --- a/src/iroh/handler.rs +++ b/src/iroh/handler.rs @@ -47,10 +47,17 @@ impl AppaProtocol { send: &mut iroh::endpoint::SendStream, receive: &mut iroh::endpoint::RecvStream, ) -> Result<(), AcceptError> { - let request = tokio::time::timeout(CONTROL_REQUEST_TIMEOUT, read_request(receive)) - .await - .map_err(|_| protocol_error("Appa control request timed out"))??; - let response = match request { + 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 { + match request { ControlMessage::Request { folder_id, capability, @@ -58,28 +65,15 @@ impl AppaProtocol { requester_endpoint, request: request_kind, } => { - if requester_endpoint.id != connection.remote_id() { - return Err(protocol_error( - "Appa requester identity does not match its endpoint", - )); - } - 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, - }, - ); - response + self.handle_folder_request( + connection, + folder_id, + capability, + invite, + requester_endpoint, + request_kind, + ) + .await } ControlMessage::Announcement { folder_id, @@ -87,23 +81,62 @@ impl AppaProtocol { root_hash, sender_endpoint, } => { - if sender_endpoint.id != connection.remote_id() { - return Err(protocol_error( - "Appa announcer identity does not match its endpoint", - )); - } - 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"); - ControlMessage::AnnouncementAccepted + self.handle_announcement( + connection, + folder_id, + capability, + root_hash, + sender_endpoint, + ) + .await } - _ => return Err(protocol_error("unsupported Appa control message")), - }; - let payload = serde_json::to_vec(&response).map_err(protocol_error)?; - send.write_all(&payload).await.map_err(protocol_error)?; - send.finish()?; - Ok(()) + _ => Err(protocol_error("unsupported Appa control message")), + } + } + + async fn handle_folder_request( + &self, + connection: &Connection, + folder_id: FolderId, + capability: String, + invite: Option>, + requester_endpoint: EndpointAddr, + request_kind: ManifestRequestKind, + ) -> Result { + 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 { + 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) } } @@ -175,7 +208,15 @@ impl AppaProtocol { } } -async fn read_request( +async fn read_control_request( + receive: &mut iroh::endpoint::RecvStream, +) -> Result { + 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 { let payload = receive @@ -185,6 +226,29 @@ async fn read_request( 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())) }