From f05f8315cbd2d3e41f3b41d2bb57b56ebdbd6d66 Mon Sep 17 00:00:00 2001
From: Artemis Tosini <me@artem.ist>
Date: Thu, 5 Sep 2024 04:18:35 +0000
Subject: [PATCH] lib: Make CMSession users poll the context

---
 lib/examples/login_qr.rs         |  4 +++-
 lib/examples/login_steamguard.rs |  4 +++-
 lib/src/connection.rs            | 14 +++++++++-----
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/lib/examples/login_qr.rs b/lib/examples/login_qr.rs
index 1832eff..510c2c3 100644
--- a/lib/examples/login_qr.rs
+++ b/lib/examples/login_qr.rs
@@ -22,7 +22,9 @@ pub async fn main() -> eyre::Result<()> {
     let servers = vapore::selection::bootstrap_find_servers().await?;
     log::debug!("Found servers: {:?}", servers);
 
-    let session = vapore::connection::CMSession::connect(&servers[0]).await?;
+    let (session, context) = vapore::connection::CMSession::connect(&servers[0]).await?;
+
+    tokio::spawn(context);
 
     session.send_notification(
         EMsg::k_EMsgClientHello,
diff --git a/lib/examples/login_steamguard.rs b/lib/examples/login_steamguard.rs
index 722bd6a..d5b4454 100644
--- a/lib/examples/login_steamguard.rs
+++ b/lib/examples/login_steamguard.rs
@@ -30,7 +30,9 @@ pub async fn main() -> eyre::Result<()> {
     let servers = vapore::selection::bootstrap_find_servers().await?;
     log::debug!("Found servers: {:?}", servers);
 
-    let session = CMSession::connect(&servers[0]).await?;
+    let (session, context) = CMSession::connect(&servers[0]).await?;
+
+    tokio::spawn(context);
 
     session.send_notification(
         EMsg::k_EMsgClientHello,
diff --git a/lib/src/connection.rs b/lib/src/connection.rs
index b3f4004..70bad47 100644
--- a/lib/src/connection.rs
+++ b/lib/src/connection.rs
@@ -71,7 +71,12 @@ impl CMSessionInner {
     }
 }
 
-struct CMContext {
+/// Task to manage a single Connection Manager socket.
+/// Functions on the matching [CMSession] objects will use this Context
+/// to send and recieve messages.
+/// Users _must_ poll this to make any progress, e.g. with `tokio::spawn(context)`
+#[must_use = "Messages will not be sent or received unless context is polled"]
+pub struct CMContext {
     socket: WebSocketStream<ConnectStream>,
     session: Arc<Mutex<CMSessionInner>>,
 }
@@ -204,7 +209,8 @@ pub struct CMSession {
 }
 
 impl CMSession {
-    pub async fn connect(server: &str) -> Result<Self, ClientError> {
+    /// Connect to a given Steam Connection Manager server
+    pub async fn connect(server: &str) -> Result<(Self, CMContext), ClientError> {
         let (socket, _) = connect_async(server)
             .await
             .with_context(|_| WebSocketConnectSnafu {
@@ -229,13 +235,11 @@ impl CMSession {
             session: inner_wrapped.clone(),
         };
 
-        tokio::spawn(context);
-
         let session = Self {
             inner: inner_wrapped,
         };
 
-        Ok(session)
+        Ok((session, context))
     }
 
     pub fn begin_heartbeat(&self, interval: u32) {