From 9a38f53f18fb8cef40c530d25639b4088ded3609 Mon Sep 17 00:00:00 2001
From: Mahad <mahad@xconn.io>
Date: Sat, 11 May 2024 18:07:13 +0500
Subject: [PATCH] detach client on abrupt client disconnect

---
 lib/src/acceptor.dart | 4 +---
 lib/src/server.dart   | 3 +++
 lib/src/types.dart    | 7 ++++++-
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/lib/src/acceptor.dart b/lib/src/acceptor.dart
index dcc7efb..ea5a5eb 100644
--- a/lib/src/acceptor.dart
+++ b/lib/src/acceptor.dart
@@ -15,6 +15,7 @@ class WAMPSessionAcceptor {
 
   IServerAuthenticator? _authenticator;
   late Serializer _serializer;
+  late StreamSubscription wsStreamSubscription;
 
   Future<BaseSession> accept(WebSocket ws) async {
     _serializer = getSerializer(ws.protocol);
@@ -22,9 +23,6 @@ class WAMPSessionAcceptor {
 
     Completer<BaseSession> completer = Completer<BaseSession>();
 
-    // ignore: cancel_subscriptions
-    late StreamSubscription<dynamic> wsStreamSubscription;
-
     wsStreamSubscription = ws.listen((message) {
       MapEntry<Object, bool> received = acceptor.receive(message);
       ws.add(received.key);
diff --git a/lib/src/server.dart b/lib/src/server.dart
index ded2237..3341f7f 100644
--- a/lib/src/server.dart
+++ b/lib/src/server.dart
@@ -45,6 +45,9 @@ class Server {
       _router.attachClient(baseSession);
 
       _handleWebSocket(baseSession, webSocket);
+      acceptor.wsStreamSubscription.onDone(() {
+        _router.detachClient(baseSession);
+      });
     }
   }
 
diff --git a/lib/src/types.dart b/lib/src/types.dart
index be30ff9..3d86ae3 100644
--- a/lib/src/types.dart
+++ b/lib/src/types.dart
@@ -51,7 +51,12 @@ abstract class IBaseSession {
 }
 
 class BaseSession extends IBaseSession {
-  BaseSession(this._ws, this._wsStreamSubscription, this.sessionDetails, this._serializer);
+  BaseSession(this._ws, this._wsStreamSubscription, this.sessionDetails, this._serializer) {
+    // close cleanly on abrupt client disconnect
+    _wsStreamSubscription.onDone(() async {
+      await close();
+    });
+  }
 
   final WebSocket _ws;
   final StreamSubscription<dynamic> _wsStreamSubscription;