+++ /dev/null
-use super::*;
-use async_recursion::async_recursion;
-use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
-use std::{
- borrow::Cow,
- collections::{HashMap, VecDeque},
- io,
- pin::Pin,
- sync::Arc,
- time::{Duration, Instant},
-};
-use tokio::{
- sync::watch,
- time::{interval, sleep, Interval, Sleep},
-};
-
-fn to_seqnum(seqnum: u16) -> usize {
- (seqnum as usize) & (REL_BUFFER - 1)
-}
-
-type Result<T> = std::result::Result<T, Error>;
-
-#[derive(Debug)]
-struct Split {
- timestamp: Option<Instant>,
- chunks: Vec<Option<Vec<u8>>>,
- got: usize,
-}
-
-struct RecvChan {
- packets: Vec<Option<Vec<u8>>>, // char ** 😛
- splits: HashMap<u16, Split>,
- seqnum: u16,
-}
-
-pub struct RudpReceiver<P: UdpPeer> {
- pub(crate) share: Arc<RudpShare<P>>,
- chans: [RecvChan; NUM_CHANS],
- udp: P::Receiver,
- close: watch::Receiver<bool>,
- closed: bool,
- resend: Interval,
- ping: Interval,
- cleanup: Interval,
- timeout: Pin<Box<Sleep>>,
- queue: VecDeque<Result<Pkt<'static>>>,
-}
-
-impl<P: UdpPeer> RudpReceiver<P> {
- pub(crate) fn new(
- udp: P::Receiver,
- share: Arc<RudpShare<P>>,
- close: watch::Receiver<bool>,
- ) -> Self {
- Self {
- udp,
- share,
- close,
- closed: false,
- resend: interval(Duration::from_millis(500)),
- ping: interval(Duration::from_secs(PING_TIMEOUT)),
- cleanup: interval(Duration::from_secs(TIMEOUT)),
- timeout: Box::pin(sleep(Duration::from_secs(TIMEOUT))),
- chans: std::array::from_fn(|_| RecvChan {
- packets: (0..REL_BUFFER).map(|_| None).collect(),
- seqnum: INIT_SEQNUM,
- splits: HashMap::new(),
- }),
- queue: VecDeque::new(),
- }
- }
-
- fn handle_err(&mut self, res: Result<()>) -> Result<()> {
- use Error::*;
-
- match res {
- Err(RemoteDisco(_)) | Err(LocalDisco) => {
- self.closed = true;
- res
- }
- Ok(_) => res,
- Err(e) => {
- self.queue.push_back(Err(e));
- Ok(())
- }
- }
- }
-
- pub async fn recv(&mut self) -> Option<Result<Pkt<'static>>> {
- use Error::*;
-
- if self.closed {
- return None;
- }
-
- loop {
- if let Some(x) = self.queue.pop_front() {
- return Some(x);
- }
-
- tokio::select! {
- _ = self.close.changed() => {
- self.closed = true;
- return Some(Err(LocalDisco));
- },
- _ = self.cleanup.tick() => {
- let timeout = Duration::from_secs(TIMEOUT);
-
- for chan in self.chans.iter_mut() {
- chan.splits = chan
- .splits
- .drain_filter(
- |_k, v| !matches!(v.timestamp, Some(t) if t.elapsed() < timeout),
- )
- .collect();
- }
- },
- _ = self.resend.tick() => {
- for chan in self.share.chans.iter() {
- for (_, ack) in chan.lock().await.acks.iter() {
- self.share.send_raw(&ack.data).await.ok(); // TODO: handle error (?)
- }
- }
- },
- _ = self.ping.tick() => {
- self.share
- .send(
- PktType::Ctl,
- Pkt {
- chan: 0,
- unrel: false,
- data: Cow::Borrowed(&[CtlType::Ping as u8]),
- },
- )
- .await
- .ok();
- }
- _ = &mut self.timeout => {
- self.closed = true;
- return Some(Err(RemoteDisco(true)));
- },
- pkt = self.udp.recv() => {
- if let Err(e) = self.handle_pkt(pkt).await {
- return Some(Err(e));
- }
- }
- }
- }
- }
-
- async fn handle_pkt(&mut self, pkt: io::Result<Vec<u8>>) -> Result<()> {
- use Error::*;
-
- let mut cursor = io::Cursor::new(pkt?);
-
- self.timeout
- .as_mut()
- .reset(tokio::time::Instant::now() + Duration::from_secs(TIMEOUT));
-
- let proto_id = cursor.read_u32::<BigEndian>()?;
- if proto_id != PROTO_ID {
- return Err(InvalidProtoId(proto_id));
- }
-
- let _peer_id = cursor.read_u16::<BigEndian>()?;
-
- let chan = cursor.read_u8()?;
- if chan >= NUM_CHANS as u8 {
- return Err(InvalidChannel(chan));
- }
-
- let res = self.process_pkt(cursor, true, chan).await;
- self.handle_err(res)?;
-
- Ok(())
- }
-
- #[async_recursion]
- async fn process_pkt(
- &mut self,
- mut cursor: io::Cursor<Vec<u8>>,
- unrel: bool,
- chan: u8,
- ) -> Result<()> {
- use Error::*;
-
- let ch = chan as usize;
- match cursor.read_u8()?.try_into()? {
- PktType::Ctl => match cursor.read_u8()?.try_into()? {
- CtlType::Ack => {
- // println!("Ack");
-
- let seqnum = cursor.read_u16::<BigEndian>()?;
- if let Some(ack) = self.share.chans[ch].lock().await.acks.remove(&seqnum) {
- ack.tx.send(true).ok();
- }
- }
- CtlType::SetPeerID => {
- // println!("SetPeerID");
-
- let mut id = self.share.remote_id.write().await;
-
- if *id != PeerID::Nil as u16 {
- return Err(PeerIDAlreadySet);
- }
-
- *id = cursor.read_u16::<BigEndian>()?;
- }
- CtlType::Ping => {
- // println!("Ping");
- }
- CtlType::Disco => {
- // println!("Disco");
- return Err(RemoteDisco(false));
- }
- },
- PktType::Orig => {
- // println!("Orig");
-
- self.queue.push_back(Ok(Pkt {
- chan,
- unrel,
- data: Cow::Owned(cursor.remaining_slice().into()),
- }));
- }
- PktType::Split => {
- // println!("Split");
-
- let seqnum = cursor.read_u16::<BigEndian>()?;
- let chunk_count = cursor.read_u16::<BigEndian>()? as usize;
- let chunk_index = cursor.read_u16::<BigEndian>()? as usize;
-
- let mut split = self.chans[ch]
- .splits
- .entry(seqnum)
- .or_insert_with(|| Split {
- got: 0,
- chunks: (0..chunk_count).map(|_| None).collect(),
- timestamp: None,
- });
-
- if split.chunks.len() != chunk_count {
- return Err(InvalidChunkCount(split.chunks.len(), chunk_count));
- }
-
- if split
- .chunks
- .get_mut(chunk_index)
- .ok_or(InvalidChunkIndex(chunk_index, chunk_count))?
- .replace(cursor.remaining_slice().into())
- .is_none()
- {
- split.got += 1;
- }
-
- split.timestamp = if unrel { Some(Instant::now()) } else { None };
-
- if split.got == chunk_count {
- let split = self.chans[ch].splits.remove(&seqnum).unwrap();
-
- self.queue.push_back(Ok(Pkt {
- chan,
- unrel,
- data: split
- .chunks
- .into_iter()
- .map(|x| x.unwrap())
- .reduce(|mut a, mut b| {
- a.append(&mut b);
- a
- })
- .unwrap_or_default()
- .into(),
- }));
- }
- }
- PktType::Rel => {
- // println!("Rel");
-
- let seqnum = cursor.read_u16::<BigEndian>()?;
- self.chans[ch].packets[to_seqnum(seqnum)].replace(cursor.remaining_slice().into());
-
- let mut ack_data = Vec::with_capacity(3);
- ack_data.write_u8(CtlType::Ack as u8)?;
- ack_data.write_u16::<BigEndian>(seqnum)?;
-
- self.share
- .send(
- PktType::Ctl,
- Pkt {
- chan,
- unrel: true,
- data: ack_data.into(),
- },
- )
- .await?;
-
- let next_pkt = |chan: &mut RecvChan| chan.packets[to_seqnum(chan.seqnum)].take();
- while let Some(pkt) = next_pkt(&mut self.chans[ch]) {
- let res = self.process_pkt(io::Cursor::new(pkt), false, chan).await;
- self.handle_err(res)?;
- self.chans[ch].seqnum = self.chans[ch].seqnum.overflowing_add(1).0;
- }
- }
- }
-
- Ok(())
- }
-}
projects / mt_rudp.git / blobdiff