add onDisconnect callback

Author: Mahad-10

lib/session.ts

@@ -42,6 +42,7 @@ export class Session {
     private _baseSession: IBaseSession;
     private _wampSession: WAMPSession;
     private _idGen: SessionScopeIDGenerator = new SessionScopeIDGenerator();
+    private _onDisconnect?: () => Promise<void>;
 
     private _callRequests: Map<number, {
         resolve: (value: Result) => void,
@@ -75,6 +76,7 @@ export class Session {
     constructor(baseSession: IBaseSession) {
         this._baseSession = baseSession;
         this._wampSession = new WAMPSession(baseSession.serializer());
+        this._baseSession.setOnDisconnect(async () => await this.markDisconnected());
 
         (async () => {
             for (; ;) {
@@ -84,6 +86,10 @@ export class Session {
         })();
     }
 
+    setDisconnectCallback(callback: () => Promise<void>): void {
+        this._onDisconnect = callback;
+    }
+
     private get _nextID(): number {
         return this._idGen.next();
     }
@@ -219,13 +225,17 @@ export class Session {
                     throw new ProtocolError(wampErrorString(message));
             }
         } else if (message instanceof Goodbye) {
-            this.markDisconnected()
+            await this.markDisconnected()
         } else {
             throw new ProtocolError(`Unexpected message type ${typeof message}`);
         }
     }
 
-    private markDisconnected() {
+    private async markDisconnected() {
+        if (this._onDisconnect) {
+            await this._onDisconnect();
+        }
+
         if (!this.isConnected()) {
             return
         }

lib/types.ts

@@ -48,13 +48,18 @@ export abstract class IBaseSession {
     isConnected(): boolean {
         throw new Error("UnimplementedError");
     }
+
+    setOnDisconnect(callback: () => Promise<void>): void {
+        throw new Error("UnimplementedError");
+    }
 }
 
 export class BaseSession extends IBaseSession {
     private readonly _ws: WebSocket;
     private readonly _wsMessageHandler: any;
     private readonly sessionDetails: SessionDetails;
     private readonly _serializer: Serializer;
+    private _onDisconnect?: () => Promise<void>;
 
     constructor(
         ws: WebSocket,
@@ -70,6 +75,9 @@ export class BaseSession extends IBaseSession {
 
         // close cleanly on abrupt client disconnect
         this._ws.addEventListener("close", async () => {
+            if (this._onDisconnect) {
+                await this._onDisconnect();
+            }
             await this.close();
         });
     }
@@ -134,6 +142,10 @@ export class BaseSession extends IBaseSession {
     isConnected(): boolean {
         return this._ws.readyState === WebSocket.OPEN;
     }
+
+    setOnDisconnect(callback: () => Promise<void>): void {
+        this._onDisconnect = callback;
+    }
 }
 
 export class Result {