use std::collections::HashMap;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;
use iroh::endpoint::{Connection, RecvStream};
use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf};
use tokio::sync::{oneshot, RwLock};
use super::codec::{self, Pong};
type ProbeKey = (String, u64, u64);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct IrohConnectionProbe {
pub transport_stable_id: u64,
pub responsive: bool,
}
#[derive(Debug)]
pub struct PrefixedRecvStream {
prefix: Vec<u8>,
prefix_offset: usize,
inner: RecvStream,
}
impl PrefixedRecvStream {
fn new(prefix: Vec<u8>, inner: RecvStream) -> Self {
Self {
prefix,
prefix_offset: 0,
inner,
}
}
}
impl AsyncRead for PrefixedRecvStream {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
if self.prefix_offset < self.prefix.len() && buf.remaining() > 0 {
let take = (self.prefix.len() - self.prefix_offset).min(buf.remaining());
let start = self.prefix_offset;
buf.put_slice(&self.prefix[start..start + take]);
self.prefix_offset += take;
return Poll::Ready(Ok(()));
}
Pin::new(&mut self.inner).poll_read(cx, buf)
}
}
pub enum IncomingUniClassification {
Control { type_id: u8, payload: Vec<u8> },
Application(PrefixedRecvStream),
MalformedControl,
}
#[derive(Debug, Clone, Default)]
pub struct IrohProbeRegistry {
pending: Arc<RwLock<HashMap<ProbeKey, oneshot::Sender<()>>>>,
next_sequence: Arc<AtomicU64>,
}
impl IrohProbeRegistry {
pub fn new() -> Self {
Self::default()
}
pub async fn probe(
&self,
endpoint_id: &str,
transport_stable_id: u64,
connection: &Connection,
timeout: Duration,
) -> bool {
let sequence = self
.next_sequence
.fetch_add(1, Ordering::Relaxed)
.wrapping_add(1);
let key = (endpoint_id.to_string(), transport_stable_id, sequence);
let (sender, receiver) = oneshot::channel();
self.pending.write().await.insert(key.clone(), sender);
let frame = codec::build_framed_ping(sequence, codec::now_unix_ms());
if !send_control_frame(connection, &frame, timeout).await {
self.pending.write().await.remove(&key);
return false;
}
#[cfg(not(target_arch = "wasm32"))]
let received = tokio::time::timeout(timeout, receiver)
.await
.is_ok_and(|result| result.is_ok());
#[cfg(target_arch = "wasm32")]
let received = {
use futures::{future::Either, FutureExt};
let response = receiver.fuse();
let deadline = gloo_timers::future::sleep(timeout).fuse();
futures::pin_mut!(response, deadline);
matches!(
futures::future::select(response, deadline).await,
Either::Left((Ok(()), _))
)
};
self.pending.write().await.remove(&key);
received
}
pub async fn deliver_pong(&self, endpoint_id: &str, transport_stable_id: u64, pong: &Pong) {
let key = (endpoint_id.to_string(), transport_stable_id, pong.seq);
if let Some(sender) = self.pending.write().await.remove(&key) {
let _ = sender.send(());
}
}
}
pub async fn respond_to_ping(connection: &Connection, payload: &[u8]) -> bool {
let Some(ping) = codec::decode_ping(payload) else {
return false;
};
let frame = codec::build_framed_pong(&ping, codec::now_unix_ms());
send_control_frame(connection, &frame, Duration::from_secs(2)).await
}
pub async fn classify_incoming_uni(mut recv: RecvStream) -> IncomingUniClassification {
let mut header = [0u8; codec::IROH_CONTROL_HEADER_LEN];
let mut read = 0usize;
while read < header.len() {
match AsyncReadExt::read(&mut recv, &mut header[read..]).await {
Ok(0) | Err(_) => {
return IncomingUniClassification::Application(PrefixedRecvStream::new(
header[..read].to_vec(),
recv,
));
}
Ok(count) => read += count,
}
}
if &header[..codec::IROH_CONTROL_MAGIC.len()] != codec::IROH_CONTROL_MAGIC {
return IncomingUniClassification::Application(PrefixedRecvStream::new(
header.to_vec(),
recv,
));
}
let frame_offset = codec::IROH_CONTROL_MAGIC.len();
let len = u32::from_be_bytes(
header[frame_offset..frame_offset + 4]
.try_into()
.expect("fixed control frame header"),
) as usize;
let type_id = header[frame_offset + 4];
let expected_payload_len = match type_id {
codec::TYPE_PING => codec::FRAMED_PING_LEN,
codec::TYPE_PONG => codec::FRAMED_PONG_LEN,
codec::TYPE_MANUAL_DISCONNECT => 0,
_ => return IncomingUniClassification::MalformedControl,
};
if len != expected_payload_len + 1 {
return IncomingUniClassification::MalformedControl;
}
let payload_len = len - 1;
let mut payload = vec![0u8; payload_len];
if AsyncReadExt::read_exact(&mut recv, &mut payload)
.await
.is_err()
{
return IncomingUniClassification::MalformedControl;
}
IncomingUniClassification::Control { type_id, payload }
}
async fn send_control_frame(connection: &Connection, frame: &[u8], timeout: Duration) -> bool {
#[cfg(not(target_arch = "wasm32"))]
{
let open = tokio::time::timeout(timeout, connection.open_uni()).await;
let Ok(Ok(mut send)) = open else {
return false;
};
return tokio::time::timeout(timeout, async {
tokio::io::AsyncWriteExt::write_all(&mut send, frame).await?;
let _ = send.finish();
Ok::<(), std::io::Error>(())
})
.await
.is_ok_and(|result| result.is_ok());
}
#[cfg(target_arch = "wasm32")]
{
let _ = timeout;
let Ok(mut stream) = connection.open_uni().await else {
return false;
};
if tokio::io::AsyncWriteExt::write_all(&mut stream, frame)
.await
.is_err()
{
return false;
}
let _ = stream.finish();
true
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn pong_is_scoped_to_endpoint_transport_and_sequence() {
let registry = IrohProbeRegistry::new();
let (sender, receiver) = oneshot::channel();
registry
.pending
.write()
.await
.insert(("peer-a".to_string(), 7, 11), sender);
let wrong_generation = Pong {
seq: 11,
sender_unix_ms: 0,
recv_unix_ms: 0,
};
registry.deliver_pong("peer-a", 8, &wrong_generation).await;
assert_eq!(registry.pending.read().await.len(), 1);
registry.deliver_pong("peer-a", 7, &wrong_generation).await;
assert!(receiver.await.is_ok());
assert!(registry.pending.read().await.is_empty());
}
}