Merge pull request #4 from andrewvy/connection-stats

Expose a function for retrieving stats about a connection

3 files changed, 25 insertions, 9 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 0639a6d..bfefaee 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -522,6 +522,7 @@ dependencies = [
  "futures",
  "futures-util",
  "quinn",
+ "quinn-proto",
  "rand",
  "rcgen",
  "ring",
diff --git a/Cargo.toml b/Cargo.toml
index 13935e6..67e2e24 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,6 +20,7 @@ tokio = { version = "1.21.2", features = ["sync", "rt-multi-thread", "macros"] }
 tokio-util = { version = "0.7.4", features = ["codec"] }
 rcgen = "0.10.0"
 quinn = "0.9.1"
+quinn-proto = "0.9.1"
 futures-util = "0.3.24"
 futures = "0.3.24"
 bincode = "1.3.3"
diff --git a/src/client.rs b/src/client.rs
index 6f97e55..98eeb86 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -12,7 +12,8 @@ use bevy::prelude::*;
 use bytes::Bytes;
 use futures::sink::SinkExt;
 use futures_util::StreamExt;
-use quinn::{ClientConfig, Endpoint};
+use quinn::{ClientConfig, Connection as QuinnConnection, Endpoint};
+use quinn_proto::ConnectionStats;
 use serde::Deserialize;
 use tokio::{
     runtime::{self},
@@ -94,16 +95,18 @@ impl ConnectionConfiguration {
     }
 }
 
-/// Current state of the client driver
-#[derive(Debug, PartialEq, Eq)]
+type InternalConnectionRef = QuinnConnection;
+
+/// Current state of a client connection
+#[derive(Debug)]
 enum ConnectionState {
     Disconnected,
-    Connected,
+    Connected(InternalConnectionRef),
 }
 
 #[derive(Debug)]
 pub(crate) enum InternalAsyncMessage {
-    Connected,
+    Connected(InternalConnectionRef),
     LostConnection,
     CertificateInteractionRequest {
         status: CertVerificationStatus,
@@ -228,7 +231,18 @@ impl Connection {
     }
 
     pub fn is_connected(&self) -> bool {
-        return self.state == ConnectionState::Connected;
+        match self.state {
+            ConnectionState::Disconnected => false,
+            ConnectionState::Connected(_) => true,
+        }
+    }
+
+    /// Returns statistics about the current connection if connected.
+    pub fn stats(&self) -> Option<ConnectionStats> {
+        match &self.state {
+            ConnectionState::Disconnected => None,
+            ConnectionState::Connected(connection) => Some(connection.stats()),
+        }
     }
 }
 
@@ -445,7 +459,7 @@ async fn connection_task(mut spawn_config: ConnectionSpawnConfig) {
 
             spawn_config
                 .to_sync_client
-                .send(InternalAsyncMessage::Connected)
+                .send(InternalAsyncMessage::Connected(connection.clone()))
                 .await
                 .expect("Failed to signal connection to sync client");
 
@@ -522,8 +536,8 @@ fn update_sync_client(
     for (connection_id, mut connection) in &mut client.connections {
         while let Ok(message) = connection.internal_receiver.try_recv() {
             match message {
-                InternalAsyncMessage::Connected => {
-                    connection.state = ConnectionState::Connected;
+                InternalAsyncMessage::Connected(internal_connection) => {
+                    connection.state = ConnectionState::Connected(internal_connection);
                     connection_events.send(ConnectionEvent(*connection_id));
                 }
                 InternalAsyncMessage::LostConnection => {