Add method to delete allocation by username (#263)

Author: logist322

turn/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## Unreleased
 
+* Added `delete_allocations_by_username` method on `Server`. This method provides possibility to manually delete allocation [#263](https://github.com/webrtc-rs/webrtc/pull/263) by [@logist322](https://github.com/logist322).
+
+
 ## v0.6.0
 
 * [#15 update deps + loosen some requirements](https://github.com/webrtc-rs/turn/pull/15) by [@melekes](https://github.com/melekes).

turn/Cargo.toml

@@ -12,6 +12,7 @@ repository = "https://github.com/webrtc-rs/turn"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+futures = "0.3.21"
 util = { version = "0.6.0", path = "../util", package = "webrtc-util", default-features = false, features = ["conn", "vnet"] }
 stun = { version = "0.4.3", path = "../stun" }
 

turn/examples/turn_client_udp.rs

@@ -7,7 +7,7 @@ use tokio::net::UdpSocket;
 use tokio::time::Duration;
 use util::Conn;
 
-// RUST_LOG=trace cargo run --color=always --package webrtc-turn --example turn_client_udp -- --host 0.0.0.0 --user user=pass --ping
+// RUST_LOG=trace cargo run --color=always --package turn --example turn_client_udp -- --host 0.0.0.0 --user user=pass --ping
 
 #[tokio::main]
 async fn main() -> Result<(), Error> {

turn/examples/turn_server_udp.rs

@@ -39,7 +39,7 @@ impl AuthHandler for MyAuthHandler {
     }
 }
 
-// RUST_LOG=trace cargo run --color=always --package webrtc-turn --example turn_server_udp -- --public-ip 0.0.0.0 --users user=pass
+// RUST_LOG=trace cargo run --color=always --package turn --example turn_server_udp -- --public-ip 0.0.0.0 --users user=pass
 
 #[tokio::main]
 async fn main() -> Result<(), Error> {

turn/src/allocation/allocation_manager.rs

@@ -5,7 +5,9 @@ use super::*;
 use crate::error::*;
 use crate::relay::*;
 
+use futures::future;
 use std::collections::HashMap;
+use stun::textattrs::Username;
 use util::Conn;
 
 // ManagerConfig a bag of config params for Manager.
@@ -34,16 +36,15 @@ impl Manager {
     pub async fn close(&self) -> Result<()> {
         let allocations = self.allocations.lock().await;
         for a in allocations.values() {
-            let mut a = a.lock().await;
             a.close().await?;
         }
         Ok(())
     }
 
     // get_allocation fetches the allocation matching the passed FiveTuple
-    pub async fn get_allocation(&self, five_tuple: &FiveTuple) -> Option<Arc<Mutex<Allocation>>> {
+    pub async fn get_allocation(&self, five_tuple: &FiveTuple) -> Option<Arc<Allocation>> {
         let allocations = self.allocations.lock().await;
-        allocations.get(&five_tuple.fingerprint()).map(Arc::clone)
+        allocations.get(five_tuple).map(Arc::clone)
     }
 
     // create_allocation creates a new allocation and starts relaying
@@ -53,7 +54,8 @@ impl Manager {
         turn_socket: Arc<dyn Conn + Send + Sync>,
         requested_port: u16,
         lifetime: Duration,
-    ) -> Result<Arc<Mutex<Allocation>>> {
+        username: Username,
+    ) -> Result<Arc<Allocation>> {
         if lifetime == Duration::from_secs(0) {
             return Err(Error::ErrLifetimeZero);
         }
@@ -66,36 +68,68 @@ impl Manager {
             .relay_addr_generator
             .allocate_conn(true, requested_port)
             .await?;
-        let mut a = Allocation::new(turn_socket, relay_socket, relay_addr, five_tuple.clone());
+        let mut a = Allocation::new(
+            turn_socket,
+            relay_socket,
+            relay_addr,
+            five_tuple.clone(),
+            username,
+        );
         a.allocations = Some(Arc::clone(&self.allocations));
 
         log::debug!("listening on relay addr: {:?}", a.relay_addr);
         a.start(lifetime).await;
         a.packet_handler().await;
 
-        let a = Arc::new(Mutex::new(a));
+        let a = Arc::new(a);
         {
             let mut allocations = self.allocations.lock().await;
-            allocations.insert(five_tuple.fingerprint(), Arc::clone(&a));
+            allocations.insert(five_tuple, Arc::clone(&a));
         }
 
         Ok(a)
     }
 
     // delete_allocation removes an allocation
     pub async fn delete_allocation(&self, five_tuple: &FiveTuple) {
-        let fingerprint = five_tuple.fingerprint();
+        let allocation = self.allocations.lock().await.remove(five_tuple);
 
-        let mut allocations = self.allocations.lock().await;
-        let allocation = allocations.remove(&fingerprint);
         if let Some(a) = allocation {
-            let mut a = a.lock().await;
             if let Err(err) = a.close().await {
                 log::error!("Failed to close allocation: {}", err);
             }
         }
     }
 
+    /// Deletes the [`Allocation`]s according to the specified `username`.
+    pub async fn delete_allocations_by_username(&self, name: &str) {
+        let to_delete = {
+            let mut allocations = self.allocations.lock().await;
+
+            let mut to_delete = Vec::new();
+
+            // TODO(logist322): Use `.drain_filter()` once stabilized.
+            allocations.retain(|_, allocation| {
+                let match_name = allocation.username.text == name;
+
+                if match_name {
+                    to_delete.push(Arc::clone(allocation));
+                }
+
+                !match_name
+            });
+
+            to_delete
+        };
+
+        future::join_all(to_delete.iter().map(|a| async move {
+            if let Err(err) = a.close().await {
+                log::error!("Failed to close allocation: {}", err);
+            }
+        }))
+        .await;
+    }
+
     // create_reservation stores the reservation for the token+port
     pub async fn create_reservation(&self, reservation_token: String, port: u16) {
         let reservations = Arc::clone(&self.reservations);

turn/src/allocation/allocation_manager/allocation_manager_test.rs

@@ -1,10 +1,9 @@
 use super::*;
-use crate::error::Result;
-use crate::relay::relay_none::*;
 
-use crate::proto::lifetime::DEFAULT_LIFETIME;
-use std::net::Ipv4Addr;
-use std::str::FromStr;
+use crate::{error::Result, proto::lifetime::DEFAULT_LIFETIME, relay::relay_none::*};
+
+use std::{net::Ipv4Addr, str::FromStr};
+use stun::{attributes::ATTR_USERNAME, textattrs::TextAttribute};
 use tokio::net::UdpSocket;
 use util::vnet::net::*;
 
@@ -62,6 +61,7 @@ async fn test_packet_handler() -> Result<()> {
             Arc::new(turn_socket),
             0,
             DEFAULT_LIFETIME,
+            TextAttribute::new(ATTR_USERNAME, "user".into()),
         )
         .await?;
 
@@ -74,8 +74,6 @@ async fn test_packet_handler() -> Result<()> {
     );
 
     let port = {
-        let a = a.lock().await;
-
         // add permission with peer1 address
         a.add_permission(Permission::new(peer_listener1.local_addr()?))
             .await;
@@ -168,11 +166,18 @@ async fn test_create_allocation_duplicate_five_tuple() -> Result<()> {
             Arc::clone(&turn_socket),
             0,
             DEFAULT_LIFETIME,
+            TextAttribute::new(ATTR_USERNAME, "user".into()),
         )
         .await?;
 
     let result = m
-        .create_allocation(five_tuple, Arc::clone(&turn_socket), 0, DEFAULT_LIFETIME)
+        .create_allocation(
+            five_tuple,
+            Arc::clone(&turn_socket),
+            0,
+            DEFAULT_LIFETIME,
+            TextAttribute::new(ATTR_USERNAME, "user".into()),
+        )
         .await;
     assert!(result.is_err(), "expected error, but got ok");
 
@@ -196,6 +201,7 @@ async fn test_delete_allocation() -> Result<()> {
             Arc::clone(&turn_socket),
             0,
             DEFAULT_LIFETIME,
+            TextAttribute::new(ATTR_USERNAME, "user".into()),
         )
         .await?;
 
@@ -231,7 +237,13 @@ async fn test_allocation_timeout() -> Result<()> {
         let five_tuple = random_five_tuple();
 
         let a = m
-            .create_allocation(five_tuple, Arc::clone(&turn_socket), 0, lifetime)
+            .create_allocation(
+                five_tuple,
+                Arc::clone(&turn_socket),
+                0,
+                lifetime,
+                TextAttribute::new(ATTR_USERNAME, "user".into()),
+            )
             .await?;
 
         allocations.push(a);
@@ -250,8 +262,7 @@ async fn test_allocation_timeout() -> Result<()> {
 
         let any_outstanding = false;
 
-        for allocation in &allocations {
-            let mut a = allocation.lock().await;
+        for a in &allocations {
             if a.close().await.is_ok() {
                 continue 'outer;
             }
@@ -280,6 +291,7 @@ async fn test_manager_close() -> Result<()> {
             Arc::clone(&turn_socket),
             0,
             Duration::from_millis(100),
+            TextAttribute::new(ATTR_USERNAME, "user".into()),
         )
         .await?;
     allocations.push(a1);
@@ -290,6 +302,7 @@ async fn test_manager_close() -> Result<()> {
             Arc::clone(&turn_socket),
             0,
             Duration::from_millis(200),
+            TextAttribute::new(ATTR_USERNAME, "user".into()),
         )
         .await?;
     allocations.push(a2);
@@ -300,8 +313,7 @@ async fn test_manager_close() -> Result<()> {
 
     m.close().await?;
 
-    for allocation in allocations {
-        let mut a = allocation.lock().await;
+    for a in allocations {
         assert!(
             a.close().await.is_err(),
             "Allocation should be closed if lifetime timeout"
@@ -310,3 +322,56 @@ async fn test_manager_close() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_delete_allocation_by_username() -> Result<()> {
+    let turn_socket: Arc<dyn Conn + Send + Sync> = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
+
+    let m = new_test_manager();
+
+    let five_tuple1 = random_five_tuple();
+    let five_tuple2 = random_five_tuple();
+    let five_tuple3 = random_five_tuple();
+
+    let _ = m
+        .create_allocation(
+            five_tuple1.clone(),
+            Arc::clone(&turn_socket),
+            0,
+            DEFAULT_LIFETIME,
+            TextAttribute::new(ATTR_USERNAME, "user".into()),
+        )
+        .await?;
+    let _ = m
+        .create_allocation(
+            five_tuple2.clone(),
+            Arc::clone(&turn_socket),
+            0,
+            DEFAULT_LIFETIME,
+            TextAttribute::new(ATTR_USERNAME, "user".into()),
+        )
+        .await?;
+    let _ = m
+        .create_allocation(
+            five_tuple3.clone(),
+            Arc::clone(&turn_socket),
+            0,
+            DEFAULT_LIFETIME,
+            TextAttribute::new(ATTR_USERNAME, String::from("user2")),
+        )
+        .await?;
+
+    assert_eq!(m.allocations.lock().await.len(), 3);
+
+    m.delete_allocations_by_username("user").await;
+
+    assert_eq!(m.allocations.lock().await.len(), 1);
+
+    assert!(
+        m.get_allocation(&five_tuple1).await.is_none()
+            && m.get_allocation(&five_tuple2).await.is_none()
+            && m.get_allocation(&five_tuple3).await.is_some()
+    );
+
+    Ok(())
+}