feat: accept promise of identity

Author: ilbertt

src/ic-websocket.test.ts

@@ -120,7 +120,7 @@ describe("IcWebsocket class", () => {
     const openMessageBytes = await mockWsServer.nextMessage as ArrayBuffer;
 
     // reconstruct the message that the client should send
-    const clientKey = icWs["_clientKey"];
+    const clientKey = icWs["_clientKey"]!;
     const { envelope: { content: openMessageContent } }: WsAgentRequestMessage<CallRequest> = Cbor.decode(openMessageBytes);
 
     expect(canisterId.compareTo(
@@ -172,7 +172,7 @@ describe("IcWebsocket class", () => {
     icWs.onerror = onError;
     await mockWsServer.connected;
 
-    const originalClientKey = { ...icWs["_clientKey"] };
+    const originalClientKey = { ...icWs["_clientKey"]! };
     // workaround to simulate the client identity
     icWs["_clientKey"] = client1Key;
     // send the open confirmation message from the canister
@@ -205,7 +205,7 @@ describe("IcWebsocket class", () => {
     icWs.onclose = onClose;
     await mockWsServer.connected;
 
-    const originalClientKey = { ...icWs["_clientKey"] };
+    const originalClientKey = { ...icWs["_clientKey"]! };
     // workaround to simulate the client identity
     icWs["_clientKey"] = client1Key;
     // send the open confirmation message from the canister
@@ -240,7 +240,7 @@ describe("IcWebsocket class", () => {
     icWs.onclose = onClose;
     await mockWsServer.connected;
 
-    const originalClientKey = { ...icWs["_clientKey"] };
+    const originalClientKey = { ...icWs["_clientKey"]! };
     // workaround to simulate the client identity
     icWs["_clientKey"] = client1Key;
     // send the open confirmation message from the canister
@@ -280,7 +280,7 @@ describe("IcWebsocket class", () => {
     // wait for the open message from the client
     await mockWsServer.nextMessage;
 
-    const originalClientKey = { ...icWs["_clientKey"] };
+    const originalClientKey = { ...icWs["_clientKey"]! };
     // workaround to simulate the client identity
     icWs["_clientKey"] = client1Key;
     // send the open confirmation message from the canister
@@ -347,7 +347,7 @@ describe("Messages acknowledgement", () => {
     // wait for the open message from the client
     await mockWsServer.nextMessage;
 
-    const originalClientKey = { ...icWs["_clientKey"] };
+    const originalClientKey = { ...icWs["_clientKey"]! };
     // workaround to simulate the client identity
     icWs["_clientKey"] = client1Key;
     // send the open confirmation message from the canister
@@ -389,7 +389,7 @@ describe("Messages acknowledgement", () => {
     // wait for the open message from the client
     await mockWsServer.nextMessage;
 
-    const originalClientKey = { ...icWs["_clientKey"] };
+    const originalClientKey = { ...icWs["_clientKey"]! };
     // workaround to simulate the client identity
     icWs["_clientKey"] = client1Key;
     // send the open confirmation message from the canister
@@ -435,7 +435,7 @@ describe("Messages acknowledgement", () => {
     // wait for the open message from the client
     await mockWsServer.nextMessage;
 
-    const originalClientKey = { ...icWs["_clientKey"] };
+    const originalClientKey = { ...icWs["_clientKey"]! };
     // workaround to simulate the client identity
     icWs["_clientKey"] = client1Key;
     // send the open confirmation message from the canister

src/ic-websocket.ts

@@ -54,7 +54,7 @@ export interface IcWebSocketConfig<S extends _WS_CANISTER_SERVICE> {
   /**
    * The identity to use for signing messages. If empty, a new random temporary identity will be generated.
    */
-  identity: SignIdentity,
+  identity: SignIdentity | Promise<SignIdentity>,
   /**
    * The IC network url to use for the underlying agent. It can be a local replica URL (e.g. http://localhost:4943) or the IC mainnet URL (https://icp0.io).
    */
@@ -91,14 +91,14 @@ export default class IcWebSocket<
   private readonly _httpAgent: HttpAgent;
   private _wsAgent: WsAgent | null = null;
   private readonly _wsInstance: WebSocket;
-  private readonly _identity: SignIdentity;
+  private readonly _identity: Promise<SignIdentity>;
   private _incomingSequenceNum = BigInt(1);
   private _outgoingSequenceNum = BigInt(0);
   private _isConnectionEstablished = false;
   private _incomingMessagesQueue: BaseQueue<ArrayBuffer>;
   private _outgoingMessagesQueue: BaseQueue<Uint8Array>;
   private _ackMessagesQueue: AckMessagesQueue;
-  private _clientKey: ClientKey;
+  private _clientKey: ClientKey | null = null;
   private _maxCertificateAgeInMinutes = 5;
 
   onclose: ((this: IcWebSocket<S, ApplicationMessageType>, ev: CloseEvent) => any) | null = null;
@@ -140,12 +140,7 @@ export default class IcWebSocket<
     if (!(config.identity instanceof SignIdentity)) {
       throw new Error("Identity must be a SignIdentity");
     }
-    this._identity = config.identity;
-
-    this._clientKey = {
-      client_principal: this.getPrincipal(),
-      client_nonce: randomBigInt(),
-    }
+    this._identity = Promise.resolve(config.identity);
 
     if (!config.networkUrl) {
       throw new Error("Network url is required");
@@ -193,8 +188,8 @@ export default class IcWebSocket<
     this._outgoingMessagesQueue.addAndProcess(new Uint8Array(data));
   }
 
-  public getPrincipal(): Principal {
-    return this._identity.getPrincipal();
+  public async getPrincipal(): Promise<Principal> {
+    return (await this._identity).getPrincipal();
   }
 
   public close() {
@@ -213,6 +208,11 @@ export default class IcWebSocket<
   }
 
   private async _onWsOpen() {
+    this._clientKey = {
+      client_principal: await this.getPrincipal(),
+      client_nonce: randomBigInt(),
+    }
+
     this._wsAgent = new WsAgent({
       identity: this._identity,
       httpAgent: this._httpAgent,
@@ -294,7 +294,7 @@ export default class IcWebSocket<
       const serviceMessage = decodeWebsocketServiceMessageContent(content as Uint8Array);
       if ("OpenMessage" in serviceMessage) {
         logger.debug("[onWsMessage] Received open message from canister");
-        if (!isClientKeyEq(serviceMessage.OpenMessage.client_key, this._clientKey)) {
+        if (!isClientKeyEq(serviceMessage.OpenMessage.client_key, this._clientKey!)) {
           throw new Error("Client key does not match");
         }
 
@@ -454,7 +454,7 @@ export default class IcWebSocket<
     this._outgoingSequenceNum++;
 
     const outgoingMessage: WebsocketMessage = {
-      client_key: this._clientKey,
+      client_key: this._clientKey!,
       sequence_num: this._outgoingSequenceNum,
       timestamp: BigInt(Date.now()) * BigInt(10 ** 6),
       content,