connlib: disconnect on token expiration (#2890)

Previously, we just expected the portal to disconnects us and 401 on the
retry, right now we harden that behaviour by also just disconnecting
when token expiration.

This seems to work, there's another part to this which is not only
handling the replies but also handling the message generated by the
portal, I'll implement that when I can easily test expirying tokens, for
now this makes the client much more stable.

rust/connlib/clients/shared/src/control.rs

@@ -263,13 +263,14 @@ impl<CB: Callbacks + 'static> ControlPlane<CB> {
         Ok(())
     }
 
+    // Errors here means we need to disconnect
     #[tracing::instrument(level = "trace", skip(self))]
     pub async fn handle_error(
         &mut self,
         reply_error: ErrorReply,
         reference: Option<Reference>,
         topic: String,
-    ) {
+    ) -> Result<()> {
         match (reply_error.error, reference) {
             (ErrorInfo::Offline, Some(reference)) => {
                 let Ok(resource_id) = reference.parse::<ResourceId>() else {
@@ -280,7 +281,7 @@ impl<CB: Callbacks + 'static> ControlPlane<CB> {
                         .tunnel
                         .callbacks()
                         .on_error(&Error::ControlProtocolError);
-                    return;
+                    return Ok(());
                 };
                 // TODO: Rate limit the number of attempts of getting the relays before just trying a local network connection
                 self.tunnel.cleanup_connection(resource_id);
@@ -290,8 +291,12 @@ impl<CB: Callbacks + 'static> ControlPlane<CB> {
                     tracing::debug!(err = ?e, "couldn't join topic: {e:#?}");
                 }
             }
+            (ErrorInfo::Reason(Reason::Known(KnownError::TokenExpired)), _) => {
+                return Err(Error::TokenExpired);
+            }
             _ => {}
         }
+        Ok(())
     }
 
     pub async fn stats_event(&mut self) {

rust/connlib/clients/shared/src/lib.rs

@@ -163,7 +163,10 @@ where
                 fallback_resolver: parking_lot::Mutex::new(None),
             };
 
-            tokio::spawn(async move {
+            tokio::spawn({
+                let runtime_stopper = runtime_stopper.clone();
+                let callbacks = callbacks.clone();
+                async move {
                 let mut log_stats_interval = tokio::time::interval(Duration::from_secs(10));
                 let mut upload_logs_interval = upload_interval();
                 loop {
@@ -171,7 +174,12 @@ where
                         Some((msg, reference, topic)) = control_plane_receiver.recv() => {
                             match msg {
                                 Ok(msg) => control_plane.handle_message(msg, reference).await?,
-                                Err(err) => control_plane.handle_error(err, reference, topic).await,
+                                Err(err) => {
+                                    if let Err(e) = control_plane.handle_error(err, reference, topic).await {
+                                        Self::disconnect_inner(runtime_stopper, &callbacks, Some(e));
+                                        break;
+                                    }
+                                },
                             }
                         },
                         event = control_plane.tunnel.next_event() => control_plane.handle_tunnel_event(event).await,
@@ -182,7 +190,7 @@ where
                 }
 
                 Result::Ok(())
-            });
+            }});
 
             tokio::spawn(async move {
                 let mut exponential_backoff = ExponentialBackoffBuilder::default().build();

rust/connlib/shared/src/control.rs

@@ -385,6 +385,8 @@ pub struct UnknownError(pub String);
 pub enum KnownError {
     #[serde(rename = "unmatched topic")]
     UnmatchedTopic,
+    #[serde(rename = "token_expired")]
+    TokenExpired,
 }
 
 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]

rust/connlib/shared/src/error.rs

@@ -141,6 +141,8 @@ pub enum ConnlibError {
     /// Connection is still being stablished, retry later
     #[error("Pending connection")]
     PendingConnection,
+    #[error("Token has expired")]
+    TokenExpired,
 }
 
 impl ConnlibError {