fix: fix mux udp relaying

src/protocol/direct/mod.rs

@@ -62,8 +62,8 @@ pub struct DirectUdpStream {
 #[async_trait]
 impl UdpRead for DirectUdpStream {
     async fn read_from(&mut self, buf: &mut [u8]) -> io::Result<(usize, Address)> {
-        let (size, addr) = self.inner.recv_from(buf).await?;
-        Ok((size, Address::SocketAddress(addr)))
+        let (len, addr) = self.inner.recv_from(buf).await?;
+        Ok((len, Address::SocketAddress(addr)))
     }
 }
 

src/protocol/dokodemo/acceptor.rs

@@ -26,8 +26,8 @@ pub struct DokodemoUdpStream {
 #[async_trait]
 impl UdpRead for DokodemoUdpStream {
     async fn read_from(&mut self, buf: &mut [u8]) -> Result<(usize, Address)> {
-        let (size, _) = self.inner.recv_from(buf).await?;
-        Ok((size, self.addr.clone()))
+        let (len, _) = self.inner.recv_from(buf).await?;
+        Ok((len, self.addr.clone()))
     }
 }
 

src/protocol/mux/acceptor.rs

@@ -35,7 +35,7 @@ impl ProxyAcceptor for MuxAcceptor {
 
     async fn accept(&self) -> io::Result<AcceptResult<Self::TS, Self::US>> {
         if let Some(result) = self.accept_stream_rx.lock().await.recv().await {
-            return Ok(result);
+            Ok(result)
         } else {
             Err(io::ErrorKind::ConnectionReset.into())
         }

src/protocol/mux/connector.rs

@@ -114,7 +114,7 @@ impl<T: ProxyConnector> ProxyConnector for MuxConnector<T> {
     async fn connect_udp(&self) -> io::Result<Self::US> {
         let mut stream = self.spawn_mux_stream().await?;
         RequestHeader::new(
-            Command::TcpConnect,
+            Command::UdpAssociate,
             &Address::DomainNameAddress("UDP_CONN".to_string(), 0),
         )
         .write_to(&mut stream)

src/protocol/mux/mod.rs

@@ -141,10 +141,9 @@ impl UdpHeader {
         R: AsyncRead + Unpin,
     {
         let address = Address::read_from_stream(stream).await?;
-        log::debug!("udp addr read: {}", address);
-        let mut len = [0u8; 2];
-        stream.read_exact(&mut len).await?;
-        let len = ((len[0] as u16) << 8) | (len[1] as u16);
+        let mut len_buf = [0u8; 2];
+        stream.read_exact(&mut len_buf).await?;
+        let len = ((len_buf[0] as u16) << 8) | (len_buf[1] as u16);
         Ok(Self {
             address,
             payload_len: len,
@@ -156,7 +155,7 @@ impl UdpHeader {
         W: AsyncWrite + Unpin,
     {
         self.address.write_to_stream(w).await?;
-        self.payload_len.to_be_bytes();
+
         w.write(&self.payload_len.to_be_bytes()).await?;
         Ok(())
     }
@@ -445,6 +444,7 @@ impl UdpWrite for MuxStreamWriteHalf {
         let len = min(buf.len(), MAX_DATA_LEN);
         let udp_header = UdpHeader::new(addr, len);
         udp_header.write_to(self).await?;
+        self.write(buf).await?;
         Ok(())
     }
 }

src/protocol/trojan/mod.rs

@@ -106,9 +106,9 @@ impl RequestHeader {
         R: AsyncRead + Unpin,
     {
         let mut hash_buf = [0u8; 56];
-        let size = stream.read(&mut hash_buf).await?;
-        if size != 56 {
-            first_packet.extend_from_slice(&hash_buf[..size]);
+        let len = stream.read(&mut hash_buf).await?;
+        if len != 56 {
+            first_packet.extend_from_slice(&hash_buf[..len]);
             return Err(new_error("first packet too short"));
         }
 
@@ -200,7 +200,6 @@ impl TrojanUdpHeader {
         W: AsyncWrite + Unpin,
     {
         self.address.write_to_stream(w).await?;
-        self.payload_len.to_be_bytes();
         w.write(&self.payload_len.to_be_bytes()).await?;
         let crlf = b"\r\n";
         w.write(crlf).await?;

src/proxy/mod.rs

@@ -41,11 +41,12 @@ const RELAY_BUFFER_SIZE: usize = 0x4000;
 async fn copy_udp<R: UdpRead, W: UdpWrite>(r: &mut R, w: &mut W) -> io::Result<()> {
     let mut buf = [0u8; RELAY_BUFFER_SIZE];
     loop {
-        let (size, addr) = r.read_from(&mut buf).await?;
-        if size == 0 {
+        let (len, addr) = r.read_from(&mut buf).await?;
+        log::debug!("udp packet from {} len {}", addr, len);
+        if len == 0 {
             break;
         }
-        w.write_to(&buf[..size], &addr).await?;
+        w.write_to(&buf[..len], &addr).await?;
     }
     Ok(())
 }
@@ -56,11 +57,11 @@ async fn copy_tcp<R: AsyncRead + Unpin, W: AsyncWrite + Unpin>(
 ) -> io::Result<()> {
     let mut buf = [0u8; RELAY_BUFFER_SIZE];
     loop {
-        let size = r.read(&mut buf).await?;
-        if size == 0 {
+        let len = r.read(&mut buf).await?;
+        if len == 0 {
             break;
         }
-        w.write(&buf[..size]).await?;
+        w.write(&buf[..len]).await?;
     }
     Ok(())
 }
@@ -77,8 +78,8 @@ pub async fn relay_udp<T: ProxyUdpStream, U: ProxyUdpStream>(a: T, b: U) {
     if let Err(e) = e {
         log::debug!("udp session ends: {}", e)
     }
-    let _ = T::reunite(a_rx, a_tx).close();
-    let _ = U::reunite(b_rx, b_tx).close();
+    let _ = T::reunite(a_rx, a_tx).close().await;
+    let _ = U::reunite(b_rx, b_tx).close().await;
 }
 
 pub async fn relay_tcp<T: ProxyTcpStream, U: ProxyTcpStream>(a: T, b: U) {