udpflow 0.1.0

Stream API for tokio-udp.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

use std::io::Result;
use std::net::SocketAddr;
use std::sync::Arc;

use tokio::net::UdpSocket;
use tokio::sync::mpsc;

use crate::UdpStreamLocal;

use crate::sockmap::{SockMap, Packet};

/// Udp packet listener.
pub struct UdpListener {
    socket: Arc<UdpSocket>,
    sockmap: SockMap,
}

impl UdpListener {
    /// Create from a **bound** udp socket.
    pub fn new(socket: UdpSocket) -> Self {
        Self {
            socket: Arc::new(socket),
            sockmap: SockMap::new(),
        }
    }

    /// Accept a new stream.
    ///
    /// A listener must be continuously polled to recv packets or accept new streams.
    ///
    /// When receiving a packet from a known peer, this function does not return,
    /// and the packet will be copied then sent to the associated
    /// [`UdpStreamLocal`](super::UdpStreamLocal).  
    pub async fn accept(&self, buf: &mut [u8]) -> Result<(UdpStreamLocal, SocketAddr)> {
        loop {
            let (n, addr) = self.socket.recv_from(buf).await?;
            debug_assert!(n != 0);

            // existed session
            if let Some(tx) = self.sockmap.get(&addr) {
                let _ = tx.send(Vec::from(&buf[..n])).await;
                continue;
            }

            // new session
            let (tx, rx) = mpsc::channel::<Packet>(32);
            let _ = tx.send(Vec::from(&buf[..n])).await;
            self.sockmap.insert(addr, tx);

            let stream = UdpStreamLocal::new(rx, self.socket.clone(), self.sockmap.clone(), addr);
            return Ok((stream, addr));
        }
    }
}