From 883c38cd3cb04f31466043720b4ba861a4e1fc8f Mon Sep 17 00:00:00 2001 From: Thomas Eizinger Date: Tue, 18 Mar 2025 15:35:57 +1100 Subject: [PATCH] fix(connlib): remove explicit `Session::disconnect` (#8474) Within the event-loop, we already react to the channel being closed which happens when the `Sender` within the `Session` gets dropped. As such, there is no need to send an explicit `Stop` command, dropping the `Session` is equivalent. As it turns out, `swift-bridge` already calls `Drop` for us when the last pointer is set to `nil`: https://github.com/firezone/firezone/blob/280a9dd999e4754ac57c768b1d423b8a358cfe40/swift/apple/FirezoneNetworkExtension/Connlib/Generated/connlib-client-apple/connlib-client-apple.swift#L24-L28 Thus, we can also remove the explicit `disconnect` call to `WrappedSession` entirely. --- rust/connlib/clients/android/src/lib.rs | 3 +-- rust/connlib/clients/apple/src/lib.rs | 6 +++--- rust/connlib/clients/shared/src/eventloop.rs | 3 +-- rust/connlib/clients/shared/src/lib.rs | 12 ++++++------ rust/headless-client/src/ipc_service.rs | 9 ++------- rust/headless-client/src/main.rs | 2 +- swift/apple/FirezoneNetworkExtension/Adapter.swift | 13 +++---------- .../connlib-client-apple/connlib-client-apple.h | 1 - .../connlib-client-apple/connlib-client-apple.swift | 4 ---- 9 files changed, 17 insertions(+), 36 deletions(-) diff --git a/rust/connlib/clients/android/src/lib.rs b/rust/connlib/clients/android/src/lib.rs index 1dcd036e0..48eafb619 100644 --- a/rust/connlib/clients/android/src/lib.rs +++ b/rust/connlib/clients/android/src/lib.rs @@ -460,10 +460,9 @@ pub unsafe extern "system" fn Java_dev_firezone_android_tunnel_ConnlibSession_di ) { let session = session_ptr as *mut SessionWrapper; catch_and_throw(&mut env, |_| { - let mut session = Box::from_raw(session); + let mut session = Box::from_raw(session); // Creating an owned `Box` from this will properly drop this at the end of the scope. session.runtime.block_on(session.telemetry.stop()); - session.inner.disconnect(); }); } diff --git a/rust/connlib/clients/apple/src/lib.rs b/rust/connlib/clients/apple/src/lib.rs index 76918430b..9c39ed30f 100644 --- a/rust/connlib/clients/apple/src/lib.rs +++ b/rust/connlib/clients/apple/src/lib.rs @@ -77,7 +77,6 @@ mod ffi { #[swift_bridge(swift_name = "setDisabledResources", return_with = err_to_string)] fn set_disabled_resources(&mut self, disabled_resources: String) -> Result<(), String>; - fn disconnect(self); } extern "Swift" { @@ -320,10 +319,11 @@ impl WrappedSession { Ok(()) } +} - fn disconnect(mut self) { +impl Drop for WrappedSession { + fn drop(&mut self) { self.runtime.block_on(self.telemetry.stop()); - self.inner.disconnect(); } } diff --git a/rust/connlib/clients/shared/src/eventloop.rs b/rust/connlib/clients/shared/src/eventloop.rs index ce283fa36..eaa1de184 100644 --- a/rust/connlib/clients/shared/src/eventloop.rs +++ b/rust/connlib/clients/shared/src/eventloop.rs @@ -28,7 +28,6 @@ pub struct Eventloop { /// Commands that can be sent to the [`Eventloop`]. pub enum Command { - Stop, Reset, SetDns(Vec), SetTun(Box), @@ -60,7 +59,7 @@ where pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll> { loop { match self.rx.poll_recv(cx) { - Poll::Ready(Some(Command::Stop)) | Poll::Ready(None) => return Poll::Ready(Ok(())), + Poll::Ready(None) => return Poll::Ready(Ok(())), Poll::Ready(Some(Command::SetDns(dns))) => { self.tunnel.state_mut().update_system_resolvers(dns); diff --git a/rust/connlib/clients/shared/src/lib.rs b/rust/connlib/clients/shared/src/lib.rs index 1d5fbe017..27718741f 100644 --- a/rust/connlib/clients/shared/src/lib.rs +++ b/rust/connlib/clients/shared/src/lib.rs @@ -26,7 +26,8 @@ const PHOENIX_TOPIC: &str = "client"; /// A session is the entry-point for connlib, maintains the runtime and the tunnel. /// -/// A session is created using [Session::connect], then to stop a session we use [Session::disconnect]. +/// A session is created using [`Session::connect`]. +/// To stop the session, simply drop this struct. #[derive(Clone)] pub struct Session { channel: tokio::sync::mpsc::UnboundedSender, @@ -103,12 +104,11 @@ impl Session { pub fn set_tun(&self, new_tun: Box) { let _ = self.channel.send(Command::SetTun(new_tun)); } +} - /// Disconnect a [`Session`]. - /// - /// This consumes [`Session`] which cleans up all state associated with it. - pub fn disconnect(self) { - let _ = self.channel.send(Command::Stop); +impl Drop for Session { + fn drop(&mut self) { + tracing::debug!("`Session` dropped") } } diff --git a/rust/headless-client/src/ipc_service.rs b/rust/headless-client/src/ipc_service.rs index ce7eeba2d..5bb4be860 100644 --- a/rust/headless-client/src/ipc_service.rs +++ b/rust/headless-client/src/ipc_service.rs @@ -434,10 +434,7 @@ impl<'a> Handler<'a> { error_msg, is_authentication_error, } => { - if let Some(session) = self.session.take() { - // Identical to dropping, but looks nicer - session.connlib.disconnect(); - } + let _ = self.session.take(); self.dns_controller.deactivate()?; self.send_ipc(ServerMsg::OnDisconnect { error_msg, @@ -491,9 +488,7 @@ impl<'a> Handler<'a> { self.send_ipc(ServerMsg::ConnectResult(result)).await? } ClientMsg::Disconnect => { - if let Some(session) = self.session.take() { - // Identical to dropping it, but looks nicer. - session.connlib.disconnect(); + if self.session.take().is_some() { self.dns_controller.deactivate()?; } // Always send `DisconnectedGracefully` even if we weren't connected, diff --git a/rust/headless-client/src/main.rs b/rust/headless-client/src/main.rs index 2f82022ac..4dcafcacd 100644 --- a/rust/headless-client/src/main.rs +++ b/rust/headless-client/src/main.rs @@ -323,7 +323,7 @@ fn main() -> Result<()> { telemetry.stop().await; // Stop telemetry before dropping session. `connlib` needs to be active for this, otherwise we won't be able to resolve the DNS name for sentry. - session.disconnect(); + drop(session); result }) diff --git a/swift/apple/FirezoneNetworkExtension/Adapter.swift b/swift/apple/FirezoneNetworkExtension/Adapter.swift index 85034df46..1a0e8715f 100644 --- a/swift/apple/FirezoneNetworkExtension/Adapter.swift +++ b/swift/apple/FirezoneNetworkExtension/Adapter.swift @@ -204,11 +204,6 @@ class Adapter { // Cancel network monitor networkMonitor?.cancel() - - // Shutdown the tunnel - session?.disconnect() - - session = nil } /// Start the tunnel. @@ -263,7 +258,7 @@ class Adapter { public func stop() { Log.log("Adapter.stop") - session?.disconnect() + // Assigning `nil` will invoke `Drop` on the Rust side session = nil networkMonitor?.cancel() @@ -416,10 +411,8 @@ extension Adapter: CallbackHandlerDelegate { } public func onDisconnect(error: String) { - // We must call disconnect() for connlib to free its session state - session?.disconnect() - - // Immediately invalidate our session pointer to prevent workQueue items from trying to use it + // Immediately invalidate our session pointer to prevent workQueue items from trying to use it. + // Assigning to `nil` will invoke `Drop` on the Rust side. session = nil // Since connlib has already shutdown by this point, we queue this callback diff --git a/swift/apple/FirezoneNetworkExtension/Connlib/Generated/connlib-client-apple/connlib-client-apple.h b/swift/apple/FirezoneNetworkExtension/Connlib/Generated/connlib-client-apple/connlib-client-apple.h index 7f8d035f8..48809a030 100644 --- a/swift/apple/FirezoneNetworkExtension/Connlib/Generated/connlib-client-apple/connlib-client-apple.h +++ b/swift/apple/FirezoneNetworkExtension/Connlib/Generated/connlib-client-apple/connlib-client-apple.h @@ -15,6 +15,5 @@ struct __private__ResultPtrAndPtr __swift_bridge__$WrappedSession$connect(void* void __swift_bridge__$WrappedSession$reset(void* self); void* __swift_bridge__$WrappedSession$set_dns(void* self, void* dns_servers); void* __swift_bridge__$WrappedSession$set_disabled_resources(void* self, void* disabled_resources); -void __swift_bridge__$WrappedSession$disconnect(void* self); diff --git a/swift/apple/FirezoneNetworkExtension/Connlib/Generated/connlib-client-apple/connlib-client-apple.swift b/swift/apple/FirezoneNetworkExtension/Connlib/Generated/connlib-client-apple/connlib-client-apple.swift index 797a13206..f64d2e804 100644 --- a/swift/apple/FirezoneNetworkExtension/Connlib/Generated/connlib-client-apple/connlib-client-apple.swift +++ b/swift/apple/FirezoneNetworkExtension/Connlib/Generated/connlib-client-apple/connlib-client-apple.swift @@ -31,10 +31,6 @@ extension WrappedSession { class public func connect(_ api_url: GenericIntoRustString, _ token: GenericIntoRustString, _ device_id: GenericIntoRustString, _ account_slug: GenericIntoRustString, _ device_name_override: Optional, _ os_version_override: Optional, _ log_dir: GenericIntoRustString, _ log_filter: GenericIntoRustString, _ callback_handler: CallbackHandler, _ device_info: GenericIntoRustString) throws -> WrappedSession { try { let val = __swift_bridge__$WrappedSession$connect({ let rustString = api_url.intoRustString(); rustString.isOwned = false; return rustString.ptr }(), { let rustString = token.intoRustString(); rustString.isOwned = false; return rustString.ptr }(), { let rustString = device_id.intoRustString(); rustString.isOwned = false; return rustString.ptr }(), { let rustString = account_slug.intoRustString(); rustString.isOwned = false; return rustString.ptr }(), { if let rustString = optionalStringIntoRustString(device_name_override) { rustString.isOwned = false; return rustString.ptr } else { return nil } }(), { if let rustString = optionalStringIntoRustString(os_version_override) { rustString.isOwned = false; return rustString.ptr } else { return nil } }(), { let rustString = log_dir.intoRustString(); rustString.isOwned = false; return rustString.ptr }(), { let rustString = log_filter.intoRustString(); rustString.isOwned = false; return rustString.ptr }(), Unmanaged.passRetained(callback_handler).toOpaque(), { let rustString = device_info.intoRustString(); rustString.isOwned = false; return rustString.ptr }()); if val.is_ok { return WrappedSession(ptr: val.ok_or_err!) } else { throw RustString(ptr: val.ok_or_err!) } }() } - - public func disconnect() { - __swift_bridge__$WrappedSession$disconnect({isOwned = false; return ptr;}()) - } } public class WrappedSessionRefMut: WrappedSessionRef { public override init(ptr: UnsafeMutableRawPointer) {