move connect_udp to network.rs

Co-authored-by: Eskil Queseth <eskilq@kth.se>

2 files changed, 31 insertions, 19 deletions

diff --git a/mumd/src/main.rs b/mumd/src/main.rs
index c8cf68b..ffd5547 100644
--- a/mumd/src/main.rs
+++ b/mumd/src/main.rs
@@ -21,13 +21,9 @@ use mumble_protocol::voice::VoicePacketPayload;
 use std::convert::Into;
 use std::convert::TryInto;
 
-use std::net::Ipv6Addr;
 use std::net::SocketAddr;
 use std::net::ToSocketAddrs;
 
-use tokio::net::UdpSocket;
-use tokio_util::udp::UdpFramed;
-
 use tokio::time::{self, Duration};
 
 use std::sync::Arc;
@@ -130,19 +126,6 @@ async fn handle_udp(
     server_addr: SocketAddr,
     crypt_state: oneshot::Receiver<ClientCryptState>,
 ) {
-    // Bind UDP socket
-    let udp_socket = UdpSocket::bind((Ipv6Addr::from(0u128), 0u16))
-        .await
-        .expect("Failed to bind UDP socket");
-
-    // Wait for initial CryptState
-    let crypt_state = match crypt_state.await {
-        Ok(crypt_state) => crypt_state,
-        // disconnected before we received the CryptSetup packet, oh well
-        Err(_) => return,
-    };
-    println!("UDP ready!");
-
     let audio = Audio::new();
     audio.output_stream.play().unwrap();
 
@@ -156,8 +139,7 @@ async fn handle_udp(
         }
     ).unwrap();
 
-    // Wrap the raw UDP packets in Mumble's crypto and voice codec (CryptState does both)
-    let (mut sink, mut source) = UdpFramed::new(udp_socket, crypt_state).split();
+    let (mut sink, mut source) = network::connect_udp(server_addr, crypt_state).await;
 
     // Note: A normal application would also send periodic Ping packets, and its own audio
     //       via UDP. We instead trick the server into accepting us by sending it one
diff --git a/mumd/src/network.rs b/mumd/src/network.rs
index 6aec343..34b5d01 100644
--- a/mumd/src/network.rs
+++ b/mumd/src/network.rs
@@ -1,15 +1,21 @@
+use futures::channel::oneshot;
 use futures::StreamExt;
 use futures_util::stream::{SplitSink, SplitStream};
 use mumble_protocol::{Serverbound, Clientbound};
 use mumble_protocol::control::ClientControlCodec;
 use mumble_protocol::control::ControlCodec;
 use mumble_protocol::control::ControlPacket;
+use mumble_protocol::crypt::ClientCryptState;
+use mumble_protocol::voice::VoicePacket;
+use std::net::Ipv6Addr;
 use std::net::SocketAddr;
 use tokio::net::TcpStream;
+use tokio::net::UdpSocket;
 use tokio_tls::TlsConnector;
 use tokio_tls::TlsStream;
 use tokio_util::codec::Decoder;
 use tokio_util::codec::Framed;
+use tokio_util::udp::UdpFramed;
 
 pub async fn connect_tcp(
     server_addr: SocketAddr,
@@ -38,3 +44,27 @@ pub async fn connect_tcp(
     // Wrap the TLS stream with Mumble's client-side control-channel codec
     ClientControlCodec::new().framed(tls_stream).split()
 }
+
+pub async fn connect_udp(
+    server_addr: SocketAddr,
+    crypt_state: oneshot::Receiver<ClientCryptState>,
+) -> (
+    SplitSink<UdpFramed<ClientCryptState>, (VoicePacket<Serverbound>, SocketAddr)>,
+    SplitStream<UdpFramed<ClientCryptState>>
+) {
+    // Bind UDP socket
+    let udp_socket = UdpSocket::bind((Ipv6Addr::from(0u128), 0u16))
+        .await
+        .expect("Failed to bind UDP socket");
+
+    // Wait for initial CryptState
+    let crypt_state = match crypt_state.await {
+        Ok(crypt_state) => crypt_state,
+        // disconnected before we received the CryptSetup packet, oh well
+        Err(_) => panic!("disconnect before crypt packet received"), //TODO exit gracefully
+    };
+    println!("UDP ready!");
+
+    // Wrap the raw UDP packets in Mumble's crypto and voice codec (CryptState does both)
+    UdpFramed::new(udp_socket, crypt_state).split()
+}