fix recv_from for UdpConn
@@ -187,7 +187,7 @@ impl ListenConfig {
}
}
- let udp_conn = Arc::new(UdpConn::new(Arc::clone(pconn), raddr));
+ let udp_conn = Arc::new(UdpConn::new(Arc::clone(pconn), Arc::clone(conns), raddr));
{
let accept_ch = accept_ch_tx.lock().await;
if let Some(tx) = &*accept_ch {
@@ -199,6 +199,11 @@ impl ListenConfig {
}
}
+ {
+ let mut m = conns.lock().await;
+ m.insert(raddr.to_string(), Arc::clone(&udp_conn));
+ }
+
Ok(Some(udp_conn))
}
}
@@ -206,14 +211,20 @@ impl ListenConfig {
/// UdpConn augments a connection-oriented connection over a UdpSocket
pub struct UdpConn {
pconn: Arc<dyn Conn + Send + Sync>,
+ conns: Arc<Mutex<HashMap<String, Arc<UdpConn>>>>,
raddr: SocketAddr,
buffer: Buffer,
}
impl UdpConn {
- fn new(pconn: Arc<dyn Conn + Send + Sync>, raddr: SocketAddr) -> Self {
+ fn new(
+ pconn: Arc<dyn Conn + Send + Sync>,
+ conns: Arc<Mutex<HashMap<String, Arc<UdpConn>>>>,
+ raddr: SocketAddr,
+ ) -> Self {
UdpConn {
pconn,
+ conns,
raddr,
buffer: Buffer::new(0, 0),
}
@@ -231,7 +242,8 @@ impl Conn for UdpConn {
}
async fn recv_from(&self, buf: &mut [u8]) -> Result<(usize, SocketAddr)> {
- Ok(self.pconn.recv_from(buf).await?)
+ let n = self.buffer.read(buf, None).await?;
+ Ok((n, self.raddr))
}
async fn send(&self, buf: &[u8]) -> Result<usize> {
@@ -247,7 +259,8 @@ impl Conn for UdpConn {
}
async fn close(&self) -> Result<()> {
- //TODO:
+ let mut conns = self.conns.lock().await;
+ conns.remove(self.raddr.to_string().as_str());
Ok(())
}
}
