feat: close conn. when receiving close message

wip: needs tests
omnia-network · Dec 8, 2023 · b456157 · b456157
1 parent 81e971f
commit b456157
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git a/src/ic-websocket.ts b/src/ic-websocket.ts
@@ -9,6 +9,7 @@ import { IDL } from "@dfinity/candid";
 import { Principal } from "@dfinity/principal";
 import {
   CanisterAckMessageContent,
+  CanisterCloseMessageContent,
   CanisterWsMessageArguments,
   ClientKeepAliveMessageContent,
   ClientKey,
@@ -371,6 +372,10 @@ export default class IcWebSocket<
         this._outgoingMessagesQueue.enableAndProcess();
       } else if ("AckMessage" in serviceMessage) {
         await this._handleAckMessageFromCanister(serviceMessage.AckMessage);
+      } else if ("CloseMessage" in serviceMessage) {
+        await this._handleCloseMessageFromCanister(serviceMessage.CloseMessage);
+        // we don't have to process any further message (there shouldn't be any anyway)
+        return false;
       } else {
         throw new Error("Invalid service message from canister");
       }
@@ -399,6 +404,17 @@ export default class IcWebSocket<
     await this._sendKeepAliveMessage();
   }
 
+  private async _handleCloseMessageFromCanister(content: CanisterCloseMessageContent): Promise<void> {
+    if ("ClosedByApplication" in content.reason) {
+      logger.debug("[onWsMessage] Received close message from canister. Reason: ClosedByApplication");
+      this._wsInstance.close(4001, "ClosedByApplication");
+    } else {
+      logger.error("[onWsMessage] Received close message from canister. Reason:", content.reason);
+      this._callOnErrorCallback(new Error(`Received close message from canister. Reason: ${content.reason}`));
+      this._wsInstance.close(4000, "Received close message from canister");
+    }
+  }
+
   private async _sendKeepAliveMessage(): Promise<void> {
     const keepAliveMessageContent: ClientKeepAliveMessageContent = {
       last_incoming_sequence_num: this._incomingSequenceNum - BigInt(1),

diff --git a/src/idl.ts b/src/idl.ts
@@ -78,12 +78,26 @@ export type CanisterAckMessageContent = {
 export type ClientKeepAliveMessageContent = {
   'last_incoming_sequence_num': bigint,
 };
+export type CloseMessageReason = {
+  WrongSequenceNumber: null,
+} | {
+  InvalidServiceMessage: null,
+} | {
+  KeepAliveTimeout: null,
+} | {
+  ClosedByApplication: null
+};
+export type CanisterCloseMessageContent = {
+  reason: CloseMessageReason,
+};
 export type WebsocketServiceMessageContent = {
   OpenMessage: CanisterOpenMessageContent,
 } | {
   AckMessage: CanisterAckMessageContent,
 } | {
   KeepAliveMessage: ClientKeepAliveMessageContent,
+} | {
+  CloseMessage: CanisterCloseMessageContent,
 };
 
 const CanisterOpenMessageContentIdl = IDL.Record({
@@ -95,10 +109,20 @@ const CanisterAckMessageContentIdl = IDL.Record({
 const ClientKeepAliveMessageContentIdl = IDL.Record({
   'last_incoming_sequence_num': IDL.Nat64,
 });
+const CloseMessageReasonIdl = IDL.Variant({
+  'WrongSequenceNumber': IDL.Null,
+  'InvalidServiceMessage': IDL.Null,
+  'KeepAliveTimeout': IDL.Null,
+  'ClosedByApplication': IDL.Null,
+});
+const CanisterCloseMessageContentIdl = IDL.Record({
+  'reason': CloseMessageReasonIdl,
+})
 const WebsocketServiceMessageContentIdl = IDL.Variant({
   'OpenMessage': CanisterOpenMessageContentIdl,
   'AckMessage': CanisterAckMessageContentIdl,
   'KeepAliveMessage': ClientKeepAliveMessageContentIdl,
+  'CloseMessage': CanisterCloseMessageContentIdl,
 });
 
 export const decodeWebsocketServiceMessageContent = (bytes: Uint8Array): WebsocketServiceMessageContent => {