1 #[cfg(test)] 2 mod mux_test; 3 4 pub mod endpoint; 5 pub mod mux_func; 6 7 use crate::error::Result; 8 use crate::mux::endpoint::Endpoint; 9 use crate::mux::mux_func::MatchFunc; 10 11 use std::collections::HashMap; 12 use std::sync::atomic::{AtomicUsize, Ordering}; 13 use std::sync::Arc; 14 use tokio::sync::{mpsc, Mutex}; 15 use util::{Buffer, Conn}; 16 17 /// mux multiplexes packets on a single socket (RFC7983) 18 19 /// The maximum amount of data that can be buffered before returning errors. 20 const MAX_BUFFER_SIZE: usize = 1000 * 1000; // 1MB 21 22 /// Config collects the arguments to mux.Mux construction into 23 /// a single structure 24 pub struct Config { 25 pub conn: Arc<dyn Conn + Send + Sync>, 26 pub buffer_size: usize, 27 } 28 29 /// Mux allows multiplexing 30 #[derive(Clone)] 31 pub struct Mux { 32 id: Arc<AtomicUsize>, 33 next_conn: Arc<dyn Conn + Send + Sync>, 34 endpoints: Arc<Mutex<HashMap<usize, Arc<Endpoint>>>>, 35 buffer_size: usize, 36 closed_ch_tx: Option<mpsc::Sender<()>>, 37 } 38 39 impl Mux { new(config: Config) -> Self40 pub fn new(config: Config) -> Self { 41 let (closed_ch_tx, closed_ch_rx) = mpsc::channel(1); 42 let m = Mux { 43 id: Arc::new(AtomicUsize::new(0)), 44 next_conn: Arc::clone(&config.conn), 45 endpoints: Arc::new(Mutex::new(HashMap::new())), 46 buffer_size: config.buffer_size, 47 closed_ch_tx: Some(closed_ch_tx), 48 }; 49 50 let buffer_size = m.buffer_size; 51 let next_conn = Arc::clone(&m.next_conn); 52 let endpoints = Arc::clone(&m.endpoints); 53 tokio::spawn(async move { 54 Mux::read_loop(buffer_size, next_conn, closed_ch_rx, endpoints).await; 55 }); 56 57 m 58 } 59 60 /// creates a new Endpoint new_endpoint(&self, f: MatchFunc) -> Arc<Endpoint>61 pub async fn new_endpoint(&self, f: MatchFunc) -> Arc<Endpoint> { 62 let mut endpoints = self.endpoints.lock().await; 63 64 let id = self.id.fetch_add(1, Ordering::SeqCst); 65 // Set a maximum size of the buffer in bytes. 66 // NOTE: We actually won't get anywhere close to this limit. 67 // SRTP will constantly read from the endpoint and drop packets if it's full. 68 let e = Arc::new(Endpoint { 69 id, 70 buffer: Buffer::new(0, MAX_BUFFER_SIZE), 71 match_fn: f, 72 next_conn: Arc::clone(&self.next_conn), 73 endpoints: Arc::clone(&self.endpoints), 74 }); 75 76 endpoints.insert(e.id, Arc::clone(&e)); 77 78 e 79 } 80 81 /// remove_endpoint removes an endpoint from the Mux remove_endpoint(&mut self, e: &Endpoint)82 pub async fn remove_endpoint(&mut self, e: &Endpoint) { 83 let mut endpoints = self.endpoints.lock().await; 84 endpoints.remove(&e.id); 85 } 86 87 /// Close closes the Mux and all associated Endpoints. close(&mut self)88 pub async fn close(&mut self) { 89 self.closed_ch_tx.take(); 90 91 let mut endpoints = self.endpoints.lock().await; 92 endpoints.clear(); 93 } 94 read_loop( buffer_size: usize, next_conn: Arc<dyn Conn + Send + Sync>, mut closed_ch_rx: mpsc::Receiver<()>, endpoints: Arc<Mutex<HashMap<usize, Arc<Endpoint>>>>, )95 async fn read_loop( 96 buffer_size: usize, 97 next_conn: Arc<dyn Conn + Send + Sync>, 98 mut closed_ch_rx: mpsc::Receiver<()>, 99 endpoints: Arc<Mutex<HashMap<usize, Arc<Endpoint>>>>, 100 ) { 101 let mut buf = vec![0u8; buffer_size]; 102 let mut n = 0usize; 103 loop { 104 tokio::select! { 105 _ = closed_ch_rx.recv() => break, 106 result = next_conn.recv(&mut buf) => { 107 if let Ok(m) = result{ 108 n = m; 109 } 110 } 111 }; 112 113 if let Err(err) = Mux::dispatch(&buf[..n], &endpoints).await { 114 log::error!("mux: ending readLoop dispatch error {:?}", err); 115 break; 116 } 117 } 118 } 119 dispatch( buf: &[u8], endpoints: &Arc<Mutex<HashMap<usize, Arc<Endpoint>>>>, ) -> Result<()>120 async fn dispatch( 121 buf: &[u8], 122 endpoints: &Arc<Mutex<HashMap<usize, Arc<Endpoint>>>>, 123 ) -> Result<()> { 124 let mut endpoint = None; 125 126 { 127 let eps = endpoints.lock().await; 128 for ep in eps.values() { 129 if (ep.match_fn)(buf) { 130 endpoint = Some(Arc::clone(ep)); 131 break; 132 } 133 } 134 } 135 136 if let Some(ep) = endpoint { 137 ep.buffer.write(buf).await?; 138 } else if !buf.is_empty() { 139 log::warn!( 140 "Warning: mux: no endpoint for packet starting with {}", 141 buf[0] 142 ); 143 } else { 144 log::warn!("Warning: mux: no endpoint for zero length packet"); 145 } 146 147 Ok(()) 148 } 149 } 150