Get allocations info (#288)

Author: logist322

.github/workflows/cargo.yml

@@ -18,7 +18,7 @@ jobs:
     name: Test
     strategy:
       matrix:
-        os: ['ubuntu-latest', 'macos-latest']
+        os: ["ubuntu-latest", "macos-latest"]
         toolchain:
           - 1.57.0 # min supported version (https://github.com/webrtc-rs/webrtc/#toolchain)
           - stable
@@ -28,10 +28,10 @@ jobs:
       - name: Install Rust ${{ matrix.toolchain }}
         uses: actions-rs/toolchain@v1
         with:
-            toolchain: ${{ matrix.toolchain }}
-            override: true      
+          toolchain: ${{ matrix.toolchain }}
+          override: true
       - name: Install Rust
-        run: rustup update stable            
+        run: rustup update stable
       - name: 📦 Cache cargo registry
         uses: actions/cache@v3
         with:
@@ -47,7 +47,7 @@ jobs:
       - name: 🏭 Cache dependencies
         uses: Swatinem/rust-cache@v2
       - name: Test
-        run: cargo test
+        run: cargo test --features metrics
 
   test_windows:
     name: Test (windows)
@@ -62,10 +62,10 @@ jobs:
       - name: Install Rust ${{ matrix.toolchain }}
         uses: actions-rs/toolchain@v1
         with:
-            toolchain: ${{ matrix.toolchain }}
-            override: true      
+          toolchain: ${{ matrix.toolchain }}
+          override: true
       - name: Install Rust
-        run: rustup update stable            
+        run: rustup update stable
       - name: 📦 Cache cargo registry
         uses: actions/cache@v3
         with:
@@ -84,7 +84,7 @@ jobs:
       #   uses: Swatinem/rust-cache@v2
       - name: Test
         working-directory: "C:\\a\\webrtc\\webrtc"
-        run: cargo test
+        run: cargo test --features metrics
 
   rustfmt_and_clippy:
     name: Check formatting style and run clippy

turn/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 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).
+* Added `get_allocations_info` method on `Server`. This method provides possibility to get information about allocations [#288](https://github.com/webrtc-rs/webrtc/pull/288) by [@logist322](https://github.com/logist322).
 
 
 ## v0.6.0

turn/Cargo.toml

@@ -34,6 +34,9 @@ hex = "0.4.3"
 clap = "3.2.6"
 criterion = "0.3.5"
 
+[features]
+metrics = []
+
 [[bench]]
 name = "bench"
 harness = false

turn/src/allocation/allocation_manager.rs

@@ -41,6 +41,33 @@ impl Manager {
         Ok(())
     }
 
+    // Returns the information about the all [`Allocation`]s associated with
+    // the specified [`FiveTuple`]s.
+    pub async fn get_allocations_info(
+        &self,
+        five_tuples: Option<Vec<FiveTuple>>,
+    ) -> HashMap<FiveTuple, AllocationInfo> {
+        let mut infos = HashMap::new();
+
+        let guarded = self.allocations.lock().await;
+
+        guarded.iter().for_each(|(five_tuple, alloc)| {
+            if five_tuples.is_none() || five_tuples.as_ref().unwrap().contains(five_tuple) {
+                infos.insert(
+                    *five_tuple,
+                    AllocationInfo::new(
+                        *five_tuple,
+                        alloc.username.text.clone(),
+                        #[cfg(feature = "metrics")]
+                        alloc.relayed_bytes.load(Ordering::Acquire),
+                    ),
+                );
+            }
+        });
+
+        infos
+    }
+
     // get_allocation fetches the allocation matching the passed FiveTuple
     pub async fn get_allocation(&self, five_tuple: &FiveTuple) -> Option<Arc<Allocation>> {
         let allocations = self.allocations.lock().await;
@@ -68,13 +95,7 @@ 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(),
-            username,
-        );
+        let mut a = Allocation::new(turn_socket, relay_socket, relay_addr, five_tuple, username);
         a.allocations = Some(Arc::clone(&self.allocations));
 
         log::debug!("listening on relay addr: {:?}", a.relay_addr);

turn/src/allocation/allocation_manager/allocation_manager_test.rs

@@ -1,8 +1,21 @@
 use super::*;
 
-use crate::{error::Result, proto::lifetime::DEFAULT_LIFETIME, relay::relay_none::*};
-
-use std::{net::Ipv4Addr, str::FromStr};
+use crate::{
+    auth::{generate_auth_key, AuthHandler},
+    client::{Client, ClientConfig},
+    error::Result,
+    proto::lifetime::DEFAULT_LIFETIME,
+    relay::{relay_none::*, relay_static::RelayAddressGeneratorStatic},
+    server::{
+        config::{ConnConfig, ServerConfig},
+        Server,
+    },
+};
+
+use std::{
+    net::{IpAddr, Ipv4Addr},
+    str::FromStr,
+};
 use stun::{attributes::ATTR_USERNAME, textattrs::TextAttribute};
 use tokio::net::UdpSocket;
 use util::vnet::net::*;
@@ -357,7 +370,7 @@ async fn test_delete_allocation_by_username() -> Result<()> {
             Arc::clone(&turn_socket),
             0,
             DEFAULT_LIFETIME,
-            TextAttribute::new(ATTR_USERNAME, String::from("user2")),
+            TextAttribute::new(ATTR_USERNAME, "user2".into()),
         )
         .await?;
 
@@ -375,3 +388,174 @@ async fn test_delete_allocation_by_username() -> Result<()> {
 
     Ok(())
 }
+
+struct TestAuthHandler;
+impl AuthHandler for TestAuthHandler {
+    fn auth_handle(&self, username: &str, realm: &str, _src_addr: SocketAddr) -> Result<Vec<u8>> {
+        Ok(generate_auth_key(username, realm, "pass"))
+    }
+}
+
+async fn create_server() -> Result<(Server, u16)> {
+    let conn = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
+    let server_port = conn.local_addr()?.port();
+
+    let server = Server::new(ServerConfig {
+        conn_configs: vec![ConnConfig {
+            conn,
+            relay_addr_generator: Box::new(RelayAddressGeneratorStatic {
+                relay_address: IpAddr::from_str("127.0.0.1")?,
+                address: "0.0.0.0".to_owned(),
+                net: Arc::new(Net::new(None)),
+            }),
+        }],
+        realm: "webrtc.rs".to_owned(),
+        auth_handler: Arc::new(TestAuthHandler {}),
+        channel_bind_timeout: Duration::from_secs(0),
+    })
+    .await?;
+
+    Ok((server, server_port))
+}
+
+async fn create_client(username: String, server_port: u16) -> Result<Client> {
+    let conn = Arc::new(UdpSocket::bind("0.0.0.0:0").await?);
+
+    Ok(Client::new(ClientConfig {
+        stun_serv_addr: format!("127.0.0.1:{}", server_port),
+        turn_serv_addr: format!("127.0.0.1:{}", server_port),
+        username,
+        password: "pass".to_owned(),
+        realm: String::new(),
+        software: String::new(),
+        rto_in_ms: 0,
+        conn,
+        vnet: None,
+    })
+    .await?)
+}
+
+#[cfg(feature = "metrics")]
+#[tokio::test]
+async fn test_get_allocations_info() -> Result<()> {
+    let (server, server_port) = create_server().await?;
+
+    let client1 = create_client("user1".to_owned(), server_port).await?;
+    client1.listen().await?;
+
+    let client2 = create_client("user2".to_owned(), server_port).await?;
+    client2.listen().await?;
+
+    let client3 = create_client("user3".to_owned(), server_port).await?;
+    client3.listen().await?;
+
+    assert!(server.get_allocations_info(None).await?.is_empty());
+
+    let user1 = client1.allocate().await?;
+    let user2 = client2.allocate().await?;
+    let user3 = client3.allocate().await?;
+
+    assert_eq!(server.get_allocations_info(None).await?.len(), 3);
+
+    let addr1 = client1
+        .send_binding_request_to(format!("127.0.0.1:{}", server_port).as_str())
+        .await?;
+    let addr2 = client2
+        .send_binding_request_to(format!("127.0.0.1:{}", server_port).as_str())
+        .await?;
+    let addr3 = client3
+        .send_binding_request_to(format!("127.0.0.1:{}", server_port).as_str())
+        .await?;
+
+    user1.send_to(b"1", addr1).await?;
+    user2.send_to(b"12", addr2).await?;
+    user3.send_to(b"123", addr3).await?;
+
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    server
+        .get_allocations_info(None)
+        .await?
+        .iter()
+        .for_each(|(_, ai)| match ai.username.as_str() {
+            "user1" => assert_eq!(ai.relayed_bytes, 1),
+            "user2" => assert_eq!(ai.relayed_bytes, 2),
+            "user3" => assert_eq!(ai.relayed_bytes, 3),
+            _ => unreachable!(),
+        });
+
+    Ok(())
+}
+
+#[cfg(feature = "metrics")]
+#[tokio::test]
+async fn test_get_allocations_info_bytes_count() -> Result<()> {
+    let (server, server_port) = create_server().await?;
+
+    let client = create_client("foo".to_owned(), server_port).await?;
+
+    client.listen().await?;
+
+    assert!(server.get_allocations_info(None).await?.is_empty());
+
+    let conn = client.allocate().await?;
+    let addr = client
+        .send_binding_request_to(format!("127.0.0.1:{}", server_port).as_str())
+        .await?;
+
+    assert!(!server.get_allocations_info(None).await?.is_empty());
+
+    assert_eq!(
+        server
+            .get_allocations_info(None)
+            .await?
+            .values()
+            .last()
+            .unwrap()
+            .relayed_bytes,
+        0
+    );
+
+    for _ in 0..10 {
+        conn.send_to(b"Hello", addr).await?;
+
+        tokio::time::sleep(Duration::from_millis(100)).await;
+    }
+
+    tokio::time::sleep(Duration::from_millis(1000)).await;
+
+    assert_eq!(
+        server
+            .get_allocations_info(None)
+            .await?
+            .values()
+            .last()
+            .unwrap()
+            .relayed_bytes,
+        50
+    );
+
+    for _ in 0..10 {
+        conn.send_to(b"Hello", addr).await?;
+
+        tokio::time::sleep(Duration::from_millis(100)).await;
+    }
+
+    tokio::time::sleep(Duration::from_millis(1000)).await;
+
+    assert_eq!(
+        server
+            .get_allocations_info(None)
+            .await?
+            .values()
+            .last()
+            .unwrap()
+            .relayed_bytes,
+        100
+    );
+
+    client.close().await?;
+    server.close().await?;
+
+    Ok(())
+}

turn/src/allocation/five_tuple.rs

@@ -12,7 +12,7 @@ use std::net::{Ipv4Addr, SocketAddr};
 // server.  The 5-tuple uniquely identifies this communication
 // stream.  The 5-tuple also uniquely identifies the Allocation on
 // the server.
-#[derive(PartialEq, Eq, Clone, Hash)]
+#[derive(PartialEq, Eq, Clone, Copy, Hash)]
 pub struct FiveTuple {
     pub protocol: Protocol,
     pub src_addr: SocketAddr,
@@ -34,3 +34,13 @@ impl fmt::Display for FiveTuple {
         write!(f, "{}_{}_{}", self.protocol, self.src_addr, self.dst_addr)
     }
 }
+
+impl fmt::Debug for FiveTuple {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("FiveTuple")
+            .field("protocol", &self.protocol)
+            .field("src_addr", &self.src_addr)
+            .field("dst_addr", &self.dst_addr)
+            .finish()
+    }
+}