From 5be8d03c440c6c85825f7aefc10628b04b56c067 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Mon, 25 Sep 2023 08:53:18 +0200 Subject: [PATCH 1/4] wip: integration tests --- Cargo.lock | 28 +- tests/integration/canister.test.ts | 1312 ++++++++++----------------- tests/integration/utils/actors.ts | 2 + tests/integration/utils/api.ts | 157 ++-- tests/integration/utils/crypto.ts | 22 - tests/integration/utils/identity.ts | 5 + tests/integration/utils/idl.ts | 67 ++ tests/integration/utils/messages.ts | 31 + tests/integration/utils/random.ts | 20 + tests/package.json | 2 +- tests/src/lib.rs | 17 +- tests/test_canister.did | 3 +- 12 files changed, 711 insertions(+), 955 deletions(-) delete mode 100644 tests/integration/utils/crypto.ts create mode 100644 tests/integration/utils/idl.ts create mode 100644 tests/integration/utils/messages.ts create mode 100644 tests/integration/utils/random.ts diff --git a/Cargo.lock b/Cargo.lock index 5791941..61394ca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -181,9 +181,9 @@ checksum = "a2bd12c1caf447e69cd4528f47f94d203fd2582878ecb9e9465484c4148a8223" [[package]] name = "candid" -version = "0.9.6" +version = "0.9.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88f6eec0ae850e006ef0fe306f362884d370624094ec55a6a26de18b251774be" +checksum = "f391a0d11d997af68e1a06b5e2ab354079cecb82b6eefb26addb38adf66d351d" dependencies = [ "anyhow", "binread", @@ -628,9 +628,9 @@ checksum = "2c6201b9ff9fd90a5a3bac2e56a830d0caa509576f0e503818ee82c181b3437a" [[package]] name = "hermit-abi" -version = "0.3.2" +version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "443144c8cdadd93ebf52ddb4056d257f5b52c04d3c804e657d19eb73fc33668b" +checksum = "d77f7ec81a6d05a3abb01ab6eb7590f6083d08449fe5a1c8b1e620283546ccb7" [[package]] name = "hex" @@ -1352,9 +1352,9 @@ checksum = "d626bb9dae77e28219937af045c257c28bfd3f69333c512553507f5f9798cb76" [[package]] name = "rustix" -version = "0.38.13" +version = "0.38.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7db8590df6dfcd144d22afd1b83b36c21a18d7cbc1dc4bb5295a8712e9eb662" +checksum = "747c788e9ce8e92b12cd485c49ddf90723550b654b32508f979b71a7b1ecda4f" dependencies = [ "bitflags 2.4.0", "errno", @@ -1398,9 +1398,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.101.5" +version = "0.101.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45a27e3b59326c16e23d30aeb7a36a24cc0d29e71d68ff611cdfb4a01d013bed" +checksum = "3c7d5dece342910d9ba34d259310cae3e0154b873b35408b787b59bce53d34fe" dependencies = [ "ring", "untrusted", @@ -1804,9 +1804,9 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.8" +version = "0.7.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "806fe8c2c87eccc8b3267cbae29ed3ab2d0bd37fca70ab622e46aaa9375ddb7d" +checksum = "1d68074620f57a0b21594d9735eb2e98ab38b17f80d3fcb189fca266771ca60d" dependencies = [ "bytes", "futures-core", @@ -1912,9 +1912,9 @@ checksum = "1dd624098567895118886609431a7c3b8f516e41d30e0643f03d94592a147e36" [[package]] name = "unicode-width" -version = "0.1.10" +version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c0edd1e5b14653f783770bce4a4dabb4a5108a5370a5f5d8cfe8710c361f6c8b" +checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85" [[package]] name = "untrusted" @@ -2086,9 +2086,9 @@ checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" [[package]] name = "winapi-util" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" +checksum = "f29e6f9198ba0d26b4c9f07dbe6f9ed633e1f3d5b8b414090084349e46a52596" dependencies = [ "winapi", ] diff --git a/tests/integration/canister.test.ts b/tests/integration/canister.test.ts index 5ab291c..ecf5eaa 100644 --- a/tests/integration/canister.test.ts +++ b/tests/integration/canister.test.ts @@ -2,22 +2,24 @@ import { IDL } from "@dfinity/candid"; import { Principal } from "@dfinity/principal"; import { Cbor } from "@dfinity/agent"; import { + anonymousClient, canisterId, client1, + client1Data, client2, + client2Data, commonAgent, gateway1, gateway2, } from "./utils/actors"; -import { getKeyPair, getMessageSignature } from "./utils/crypto"; import { - getWebsocketMessage, isMessageBodyValid, isValidCertificate, + reinitialize, wsClose, + wsGetMessages, wsMessage, wsOpen, - wsRegister, wsSend, wsWipe, } from "./utils/api"; @@ -27,583 +29,263 @@ import type { CanisterWsGetMessagesResult, CanisterWsMessageResult, CanisterWsOpenResult, - CanisterWsRegisterResult, CanisterWsSendResult, + ClientKey, + WebsocketMessage, } from "../src/declarations/test_canister/test_canister.did"; -import type { WebsocketMessage } from "./utils/api"; +import { generateClientKey, getRandomClientNonce } from "./utils/random"; +import { CanisterOpenMessageContent, WebsocketServiceMessageContent, encodeWebsocketServiceMessageContent, getServiceMessageFromCanisterMessage, isClientKeyEq } from "./utils/idl"; +import { createWebsocketMessage, decodeWebsocketMessage, filterServiceMessagesFromCanisterMessages } from "./utils/messages"; const MAX_NUMBER_OF_RETURNED_MESSAGES = 10; // set in the CDK const SEND_MESSAGES_COUNT = MAX_NUMBER_OF_RETURNED_MESSAGES + 2; // test with more messages to check the indexes and limits const MAX_GATEWAY_KEEP_ALIVE_TIME_MS = 15_000; // set in the CDK +const DEFAULT_TEST_SEND_ACK_INTERVAL_MS = 300_000; // 5 minutes to make sure the canister doesn't reset the client +const DEFAULT_TEST_KEEP_ALIVE_DELAY_MS = 300_000; // 5 minutes to make sure the canister doesn't reset the client -let client1KeyPair: { publicKey: Uint8Array; secretKey: Uint8Array | string; }; -let client2KeyPair: { publicKey: Uint8Array; secretKey: Uint8Array | string; }; +let client1Key: ClientKey; +let client2Key: ClientKey; -// the status index used by the gateway to send a keep-alive message -let gatewayStatusIndex = 0; - -const sendGatewayStatusMessage = async (index?: number) => { - const statusIndex = index !== undefined ? index : gatewayStatusIndex; - - await wsMessage({ - message: { - IcWebSocketGatewayStatus: { - status_index: BigInt(statusIndex), - } - }, - actor: gateway1, - }, true); - - gatewayStatusIndex += 1; -}; - -const assignKeyPairsToClients = async () => { - if (!client1KeyPair) { - client1KeyPair = await getKeyPair(); +const assignKeysToClients = async () => { + if (!client1Key) { + client1Key = generateClientKey((await client1Data.identity).getPrincipal()); } - if (!client2KeyPair) { - client2KeyPair = await getKeyPair(); + if (!client2Key) { + client2Key = generateClientKey((await client2Data.identity).getPrincipal()); } }; // testing again canister takes quite a while jest.setTimeout(60_000); -describe("Canister - ws_register", () => { - beforeAll(async () => { - await assignKeyPairsToClients(); - }); - - afterAll(async () => { - await wsWipe(gateway1); - }); - - it("should register a client", async () => { - const res = await wsRegister({ - clientActor: client1, - clientKey: client1KeyPair.publicKey, - }); - - expect(res).toMatchObject({ - Ok: null, - }); - }); -}); - describe("Canister - ws_open", () => { beforeAll(async () => { - await assignKeyPairsToClients(); - - await wsRegister({ - clientActor: client1, - clientKey: client1KeyPair.publicKey, - }, true); + await assignKeysToClients(); }); afterAll(async () => { - await wsWipe(gateway1); + await wsWipe(); }); - beforeEach(async () => { - await sendGatewayStatusMessage(); - }); - - it("fails for a gateway which is not registered", async () => { + it("fails for an anonymous client", async () => { const res = await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client1KeyPair.secretKey, canisterId, - gatewayActor: gateway2, - }); + clientActor: anonymousClient, + clientNonce: getRandomClientNonce(), + }) expect(res).toMatchObject({ - Err: "caller is not the gateway that has been registered during CDK initialization", + Err: "anonymous principal cannot open a connection", }); }); - it("fails if a registered gateway relays a wrong first message", async () => { - // empty message - let content = Cbor.encode({}) - let res = await gateway1.ws_open({ - content: new Uint8Array(content), - sig: await getMessageSignature(content, client1KeyPair.secretKey), - }); - expect(res).toMatchObject({ - Err: "missing field `client_key`", - }); - - // with client_key - content = Cbor.encode({ - client_key: client1KeyPair.publicKey, - }); - res = await gateway1.ws_open({ - content: new Uint8Array(content), - sig: await getMessageSignature(content, client1KeyPair.secretKey), - }); - expect(res).toMatchObject({ - Err: "missing field `canister_id`", - }); - }); - - it("fails for a client which is not registered", async () => { + it("fails for the registered gateway", async () => { const res = await wsOpen({ - clientPublicKey: client2KeyPair.publicKey, - clientSecretKey: client2KeyPair.secretKey, canisterId, - gatewayActor: gateway1, + clientActor: gateway1, + clientNonce: getRandomClientNonce(), }); expect(res).toMatchObject({ - Err: "client's public key has not been previously registered by client", + Err: "caller is the registered gateway which can't open a connection for itself", }); }); - it("fails for an invalid signature", async () => { - // sign message with client2 secret key but send client1 public key + it("should open a connection", async () => { const res = await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client2KeyPair.secretKey, canisterId, - gatewayActor: gateway1, + clientActor: client1, + clientNonce: client1Key.client_nonce, }); expect(res).toMatchObject({ - Err: "Signature doesn't verify", + Ok: null, }); - }); - - it("fails for a client which is not registered after the gateway has been reset", async () => { - await sendGatewayStatusMessage(0); - const res = await wsOpen({ - clientPublicKey: client2KeyPair.publicKey, - clientSecretKey: client2KeyPair.secretKey, - canisterId, + const msgs = await wsGetMessages({ + fromNonce: 0, gatewayActor: gateway1, }); - expect(res).toMatchObject({ - Err: "client's public key has not been previously registered by client", + const serviceMessages = filterServiceMessagesFromCanisterMessages(msgs.messages); + + expect(isClientKeyEq(serviceMessages[0].client_key, client1Key)).toBe(true); + const openMessage = getServiceMessageFromCanisterMessage(serviceMessages[0]); + expect(openMessage).toMatchObject({ + OpenMessage: expect.any(Object), }); + const openMessageContent = (openMessage as { OpenMessage: CanisterOpenMessageContent }).OpenMessage; + expect(isClientKeyEq(openMessageContent.client_key, client1Key)).toBe(true); }); - it("fails for a client which is registered, but after the gateway increased the status index by two and then been reset", async () => { - // reset the canister state from the previous test - await wsWipe(gateway1); - // register the client again - await wsRegister({ - clientActor: client1, - clientKey: client1KeyPair.publicKey, - }, true); - - // send two status messages to make the client key shift out of the tmp ones - await sendGatewayStatusMessage(); - await sendGatewayStatusMessage(); - - // reset the gateway on the canister - await sendGatewayStatusMessage(0); - + it("fails for a client with the same nonce", async () => { const res = await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client1KeyPair.secretKey, canisterId, - gatewayActor: gateway1, + clientActor: client1, + clientNonce: client1Key.client_nonce, }); expect(res).toMatchObject({ - Err: "client's public key has not been previously registered by client", + Err: `client with key ${client1Key.client_principal.toText()}_${client1Key.client_nonce} already has an open connection`, }); }); - it("should open the websocket for a registered client after gateway has been reset", async () => { - // reset the canister state from the previous test - await wsWipe(gateway1); - // register the client again - await wsRegister({ - clientActor: client1, - clientKey: client1KeyPair.publicKey, - }, true); - - // reset the gateway on the canister - await sendGatewayStatusMessage(0); - + it("should open a connection for the same client with a different nonce", async () => { + const clientKey = { + ...client1Key, + client_nonce: getRandomClientNonce(), + } const res = await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client1KeyPair.secretKey, canisterId, - gatewayActor: gateway1, + clientActor: client1, + clientNonce: clientKey.client_nonce, }); expect(res).toMatchObject({ - Ok: { - client_key: client1KeyPair.publicKey, - canister_id: Principal.fromText(canisterId), - nonce: BigInt(0), - }, + Ok: null, }); - }); - it("should open the websocket for a registered client", async () => { - // reset the canister state from the previous test - await wsWipe(gateway1); - // setup the canister state again - await sendGatewayStatusMessage(); - await wsRegister({ - clientActor: client1, - clientKey: client1KeyPair.publicKey, - }, true); - - // open the websocket - const res = await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client1KeyPair.secretKey, - canisterId, + const msgs = await wsGetMessages({ + fromNonce: 0, gatewayActor: gateway1, }); - expect(res).toMatchObject({ - Ok: { - client_key: client1KeyPair.publicKey, - canister_id: Principal.fromText(canisterId), - nonce: BigInt(0), - }, + const serviceMessages = filterServiceMessagesFromCanisterMessages(msgs.messages); + const serviceMessagesForClient = serviceMessages.filter((msg) => isClientKeyEq(msg.client_key, clientKey)); + + const openMessage = getServiceMessageFromCanisterMessage(serviceMessagesForClient[0]); + expect(openMessage).toMatchObject({ + OpenMessage: expect.any(Object), }); + const openMessageContent = (openMessage as { OpenMessage: CanisterOpenMessageContent }).OpenMessage; + expect(isClientKeyEq(openMessageContent.client_key, clientKey)).toBe(true); }); }); describe("Canister - ws_message", () => { beforeAll(async () => { - await assignKeyPairsToClients(); - - await wsRegister({ - clientActor: client1, - clientKey: client1KeyPair.publicKey, - }, true); + await assignKeysToClients(); await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client1KeyPair.secretKey, + clientNonce: client1Key.client_nonce, canisterId, - gatewayActor: gateway1, + clientActor: client1, }, true); }); afterAll(async () => { - await wsWipe(gateway1); - }); - - beforeEach(async () => { - await sendGatewayStatusMessage(); + await wsWipe(); }); - it("fails if a non registered gateway sends an IcWebSocketEstablished message", async () => { - const res = await wsMessage({ - message: { - IcWebSocketEstablished: client1KeyPair.publicKey, - }, - actor: gateway2, - }); - - expect(res).toMatchObject({ - Err: "caller is not the gateway that has been registered during CDK initialization", - }); - }); - - it("fails if a non registered gateway sends a RelayedByGateway message", async () => { - const content = getWebsocketMessage(client1KeyPair.publicKey, 0); - const res = await wsMessage({ - message: { - RelayedByGateway: { - content, - sig: await getMessageSignature(content, client1KeyPair.secretKey), - } - }, - actor: gateway2, - }); - - expect(res).toMatchObject({ - Err: "caller is not the gateway that has been registered during CDK initialization", - }); - }); - - it("fails if a non registered client sends a DirectlyFromClient message", async () => { - const message = IDL.encode([IDL.Record({ 'text': IDL.Text })], [{ text: "pong" }]); - const res = await wsMessage({ - message: { - DirectlyFromClient: { - client_key: client2KeyPair.publicKey, - message: new Uint8Array(message), - } - }, - actor: client2, - }); - - expect(res).toMatchObject({ - Err: "client is not registered, call ws_register first", - }); - }); - - it("fails if a non registered client sends a DirectlyFromClient message using a registered client key", async () => { - const message = IDL.encode([IDL.Record({ 'text': IDL.Text })], [{ text: "pong" }]); + it("fails if client is not registered", async () => { const res = await wsMessage({ - message: { - DirectlyFromClient: { - client_key: client1KeyPair.publicKey, - message: new Uint8Array(message), - } - }, + message: createWebsocketMessage(client2Key, 0), actor: client2, }); expect(res).toMatchObject({ - Err: "caller is not the same that registered the public key", + Err: `client with principal ${client2Key.client_principal.toText()} doesn't have an open connection`, }); }); - it("fails if a registered gateway sends an IcWebSocketEstablished message for a non registered client", async () => { + it("fails if client sends a message with a different client key", async () => { + // first, send a message with a different principal const res = await wsMessage({ - message: { - IcWebSocketEstablished: client2KeyPair.publicKey, - }, - actor: gateway1, - }); - - expect(res).toMatchObject({ - Err: "client's public key has not been previously registered by client", - }); - }); - - it("fails if a registered gateway sends a wrong RelayedByGateway message", async () => { - // empty message - let content = Cbor.encode({}); - let res = await wsMessage({ - message: { - RelayedByGateway: { - content: new Uint8Array(content), - sig: await getMessageSignature(content, client2KeyPair.secretKey), - }, - }, - actor: gateway1, - }); - expect(res).toMatchObject({ - Err: "missing field `client_key`", + message: createWebsocketMessage({ ...client1Key, client_principal: client2Key.client_principal }, 0), + actor: client1, }); - // with client_key - content = Cbor.encode({ - client_key: client1KeyPair.publicKey, - }); - res = await wsMessage({ - message: { - RelayedByGateway: { - content: new Uint8Array(content), - sig: await getMessageSignature(content, client2KeyPair.secretKey), - }, - }, - actor: gateway1, - }); expect(res).toMatchObject({ - Err: "missing field `sequence_num`", + Err: `client with principal ${client1Key.client_principal.toText()} has a different key than the one used in the message`, }); - // with client_key, sequence_num - content = Cbor.encode({ - client_key: client1KeyPair.publicKey, - sequence_num: 0, - }); - res = await wsMessage({ - message: { - RelayedByGateway: { - content: new Uint8Array(content), - sig: await getMessageSignature(content, client2KeyPair.secretKey), - }, - }, - actor: gateway1, - }); - expect(res).toMatchObject({ - Err: "missing field `timestamp`", + // then, send a message with a different nonce + const res2 = await wsMessage({ + message: createWebsocketMessage({ ...client1Key, client_nonce: getRandomClientNonce() }, 0), + actor: client1, }); - // with client_key, sequence_num, timestamp - content = Cbor.encode({ - client_key: client1KeyPair.publicKey, - sequence_num: 0, - timestamp: Date.now(), - }); - res = await wsMessage({ - message: { - RelayedByGateway: { - content: new Uint8Array(content), - sig: await getMessageSignature(content, client2KeyPair.secretKey), - }, - }, - actor: gateway1, - }); - expect(res).toMatchObject({ - Err: "missing field `message`", + expect(res2).toMatchObject({ + Err: `client with principal ${client1Key.client_principal.toText()} has a different key than the one used in the message`, }); }); - it("fails if a registered gateway sends a RelayedByGateway message with an invalid signature", async () => { - const content = getWebsocketMessage(client1KeyPair.publicKey, 0); + it("should send a message from a registered client", async () => { const res = await wsMessage({ - message: { - RelayedByGateway: { - content, - sig: await getMessageSignature(content, client2KeyPair.secretKey), - } - }, - actor: gateway1, + message: createWebsocketMessage(client1Key, 1), + actor: client1, }); expect(res).toMatchObject({ - Err: "Signature doesn't verify", + Ok: null, }); }); - it("fails if registered gateway sends a RelayedByGateway message with a wrong sequence number", async () => { - const appMessage = IDL.Record({ 'text': IDL.Text }).encodeValue({ text: "pong" }); - let content = getWebsocketMessage(client1KeyPair.publicKey, 1, appMessage); - let res = await wsMessage({ - message: { - RelayedByGateway: { - content, - sig: await getMessageSignature(content, client1KeyPair.secretKey), - } - }, - actor: gateway1, - }); - expect(res).toMatchObject({ - Err: "incoming client's message relayed from WS Gateway does not have the expected sequence number", + it("fails if client sends a message with a wrong sequence number", async () => { + const actualSequenceNumber = 1; + const expectedSequenceNumber = 2; // first valid message with sequence number 1 was sent in the previous test + const res = await wsMessage({ + message: createWebsocketMessage(client1Key, actualSequenceNumber), + actor: client1, }); - // send a correct message to increase the sequence number - content = getWebsocketMessage(client1KeyPair.publicKey, 0, appMessage); - res = await wsMessage({ - message: { - RelayedByGateway: { - content, - sig: await getMessageSignature(content, client1KeyPair.secretKey), - } - }, - actor: gateway1, - }); expect(res).toMatchObject({ - Ok: null, + Err: `incoming client's message does not have the expected sequence number. Expected: ${expectedSequenceNumber}, actual: ${actualSequenceNumber}. Client removed.`, }); - // send a message with the old sequence number - content = getWebsocketMessage(client1KeyPair.publicKey, 0, appMessage); - res = await wsMessage({ - message: { - RelayedByGateway: { - content, - sig: await getMessageSignature(content, client1KeyPair.secretKey), - } - }, - actor: gateway1, - }); - expect(res).toMatchObject({ - Err: "incoming client's message relayed from WS Gateway does not have the expected sequence number", + // check if client has been removed + const res2 = await wsMessage({ + message: createWebsocketMessage(client1Key, 0), // here the sequence number doesn't matter + actor: client1, }); - // send a message with a sequence number that is too high - content = getWebsocketMessage(client1KeyPair.publicKey, 2, appMessage); - res = await wsMessage({ - message: { - RelayedByGateway: { - content, - sig: await getMessageSignature(content, client1KeyPair.secretKey), - } - }, - actor: gateway1, - }); - expect(res).toMatchObject({ - Err: "incoming client's message relayed from WS Gateway does not have the expected sequence number", + expect(res2).toMatchObject({ + Err: `client with principal ${client1Key.client_principal.toText()} doesn't have an open connection`, }); }); - it("fails if a registered gateway sends a RelayedByGateway for a registered client that doesn't have open connection", async () => { - // register another client, but don't call ws_open for it - await wsRegister({ - clientActor: client2, - clientKey: client2KeyPair.publicKey, + it("fails if a client sends a wrong service message", async () => { + // open the connection again + await wsOpen({ + clientNonce: client1Key.client_nonce, + canisterId, + clientActor: client1, }, true); - const content = getWebsocketMessage(client2KeyPair.publicKey, 0); + // wring content encoding const res = await wsMessage({ - message: { - RelayedByGateway: { - content, - sig: await getMessageSignature(content, client2KeyPair.secretKey), - } - }, - actor: gateway1, + message: createWebsocketMessage(client1Key, 1, true, new Uint8Array([1, 2, 3])), + actor: client1, }); expect(res).toMatchObject({ - Err: "expected incoming message num not initialized for client", + Err: expect.stringContaining("Error decoding service message from client:"), }); - }); - it("fails if registered gateway sends a DirectlyFromClient message", async () => { - const message = IDL.encode([IDL.Record({ 'text': IDL.Text })], [{ text: "pong" }]); - const res = await wsMessage({ - message: { - DirectlyFromClient: { - client_key: client1KeyPair.publicKey, - message: new Uint8Array(message), - } - }, - actor: gateway1, + const wrongServiceMessage: WebsocketServiceMessageContent = { + // the client can only send KeepAliveMessage variant + AckMessage: { + last_incoming_sequence_num: BigInt(0), + } + }; + const res2 = await wsMessage({ + message: createWebsocketMessage(client1Key, 2, true, encodeWebsocketServiceMessageContent(wrongServiceMessage)), + actor: client1, }); - expect(res).toMatchObject({ - Err: "caller is not the same that registered the public key", + expect(res2).toMatchObject({ + Err: "invalid keep alive message content", }); }); - it("a registered gateway should send a message (IcWebSocketEstablished) for a registered client", async () => { - const res = await wsMessage({ - message: { - IcWebSocketEstablished: client1KeyPair.publicKey, + it("should send a service message from a registered client", async () => { + const clientServiceMessage: WebsocketServiceMessageContent = { + KeepAliveMessage: { + last_incoming_sequence_num: BigInt(0), }, - actor: gateway1, - }); - - expect(res).toMatchObject({ - Ok: null, - }); - }); - - it("a registered gateway should send a message (RelayedByGateway) for a registered client", async () => { - const appMessage = IDL.encode([IDL.Record({ 'text': IDL.Text })], [{ text: "pong" }]); - // the message with sequence number 0 has been sent in a previous test, so we send a message with sequence number 1 - const content = getWebsocketMessage(client1KeyPair.publicKey, 1, appMessage); + }; const res = await wsMessage({ - message: { - RelayedByGateway: { - content, - sig: await getMessageSignature(content, client1KeyPair.secretKey), - } - }, - actor: gateway1, - }); - - expect(res).toMatchObject({ - Ok: null, - }); - }); - - it("a registered client should send a message (DirectlyFromClient)", async () => { - const message = IDL.encode([IDL.Record({ 'text': IDL.Text })], [{ text: "pong" }]); - const res = await wsMessage({ - message: { - DirectlyFromClient: { - client_key: client1KeyPair.publicKey, - message: new Uint8Array(message), - } - }, + message: createWebsocketMessage(client1Key, 3, true, encodeWebsocketServiceMessageContent(clientServiceMessage)), actor: client1, }); @@ -614,32 +296,6 @@ describe("Canister - ws_message", () => { }); describe("Canister - ws_get_messages (failures,empty)", () => { - beforeAll(async () => { - await assignKeyPairsToClients(); - - await wsRegister({ - clientActor: client1, - clientKey: client1KeyPair.publicKey, - }, true); - - await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client1KeyPair.secretKey, - canisterId, - gatewayActor: gateway1, - }, true); - - await commonAgent.fetchRootKey(); - }); - - afterAll(async () => { - await wsWipe(gateway1); - }); - - beforeEach(async () => { - await sendGatewayStatusMessage(); - }); - it("fails if a non registered gateway tries to get messages", async () => { const res = await gateway2.ws_get_messages({ nonce: BigInt(0), @@ -677,245 +333,245 @@ describe("Canister - ws_get_messages (failures,empty)", () => { }); }); -describe("Canister - ws_message (gateway status)", () => { +// describe("Canister - ws_message (gateway status)", () => { +// beforeAll(async () => { +// await assignKeysToClients(); + +// await wsRegister({ +// clientActor: client1, +// clientKey: client1Key.publicKey, +// }, true); + +// await wsOpen({ +// clientPublicKey: client1Key.publicKey, +// clientSecretKey: client1Key.secretKey, +// canisterId, +// clientActor: gateway1, +// }, true); + +// await wsSend({ +// clientKey: client1Key.publicKey, +// actor: client1, +// message: { text: "test" }, +// }, true); +// }); + +// afterAll(async () => { +// await wsWipe(gateway1); +// }); + +// it("fails if a non registered gateway sends an IcWebSocketGatewayStatus message", async () => { +// const res = await wsMessage({ +// message: { +// IcWebSocketGatewayStatus: { +// status_index: BigInt(1), +// }, +// }, +// actor: gateway2, +// }); + +// expect(res).toMatchObject({ +// Err: "caller is not the gateway that has been registered during CDK initialization", +// }); +// }); + +// it("registered gateway should update the status index", async () => { +// const res = await wsMessage({ +// message: { +// IcWebSocketGatewayStatus: { +// status_index: BigInt(2), // set it high to test behavior for indexes behind the current one +// }, +// }, +// actor: gateway1, +// }); + +// expect(res).toMatchObject({ +// Ok: null, +// }); +// }); + +// it("fails if a registered gateway sends an IcWebSocketGatewayStatus with a wrong status index (equal to current)", async () => { +// const res = await wsMessage({ +// message: { +// IcWebSocketGatewayStatus: { +// status_index: BigInt(2), +// }, +// }, +// actor: gateway1, +// }); + +// expect(res).toMatchObject({ +// Err: "Gateway status index is equal to or behind the current one", +// }); +// }); + +// it("fails if a registered gateway sends an IcWebSocketGatewayStatus with a wrong status index (behind the current)", async () => { +// const res = await wsMessage({ +// message: { +// IcWebSocketGatewayStatus: { +// status_index: BigInt(1), +// }, +// }, +// actor: gateway1, +// }); + +// expect(res).toMatchObject({ +// Err: "Gateway status index is equal to or behind the current one", +// }); +// }); + +// it("registered gateway should disconnect after maximum time", async () => { +// let res = await gateway1.ws_get_messages({ +// nonce: BigInt(0), +// }); + +// expect(res).toMatchObject({ +// Ok: { +// messages: expect.any(Array), +// cert: expect.any(Uint8Array), +// tree: expect.any(Uint8Array), +// }, +// }); +// expect((res as { Ok: CanisterOutputCertifiedMessages }).Ok.messages.length).toEqual(1); + +// // wait for the maximum time the gateway can send a status message, +// // so that the internal canister state is reset +// // double the time to make sure the canister state is reset +// await new Promise((resolve) => setTimeout(resolve, 2 * MAX_GATEWAY_KEEP_ALIVE_TIME_MS)); + +// // check if messages have been deleted +// res = await gateway1.ws_get_messages({ +// nonce: BigInt(0), +// }); +// expect(res).toMatchObject({ +// Ok: { +// messages: [], +// cert: expect.any(Uint8Array), +// tree: expect.any(Uint8Array), +// }, +// }); + +// // check if registered client has been deleted +// const sendRes = await wsSend({ +// clientKey: client1Key.publicKey, +// actor: client1, +// message: { text: "test" }, +// }); +// expect(sendRes).toMatchObject({ +// Err: "client's public key has not been previously registered by client", +// }); +// }); + +// it("registered gateway should reconnect by resetting the status index", async () => { +// let res = await wsMessage({ +// message: { +// IcWebSocketGatewayStatus: { +// status_index: BigInt(0), +// }, +// }, +// actor: gateway1, +// }); + +// expect(res).toMatchObject({ +// Ok: null, +// }); + +// res = await wsMessage({ +// message: { +// IcWebSocketGatewayStatus: { +// status_index: BigInt(1), +// }, +// }, +// actor: gateway1, +// }); + +// expect(res).toMatchObject({ +// Ok: null, +// }); +// }); + +// it("registered gateway should reconnect before maximum time", async () => { +// // reconnect the client +// await wsRegister({ +// clientActor: client1, +// clientKey: client1Key.publicKey, +// }, true); + +// await wsOpen({ +// clientPublicKey: client1Key.publicKey, +// clientSecretKey: client1Key.secretKey, +// canisterId, +// clientActor: gateway1, +// }, true); + +// // send a test message from the canister to check if the internal state is reset +// await wsSend({ +// clientKey: client1Key.publicKey, +// actor: client1, +// message: { text: "test" }, +// }, true); + +// // check if the canister has the message in the queue +// let messagesRes = await gateway1.ws_get_messages({ +// nonce: BigInt(0), +// }); +// expect(messagesRes).toMatchObject({ +// Ok: { +// messages: expect.any(Array), +// cert: expect.any(Uint8Array), +// tree: expect.any(Uint8Array), +// }, +// }); +// expect((messagesRes as { Ok: CanisterOutputCertifiedMessages }).Ok.messages.length).toEqual(1); + +// // simulate a reconnection +// const res = await wsMessage({ +// message: { +// IcWebSocketGatewayStatus: { +// status_index: BigInt(0), +// }, +// }, +// actor: gateway1, +// }); +// expect(res).toMatchObject({ +// Ok: null, +// }); + +// // check if the canister reset the internal state +// messagesRes = await gateway1.ws_get_messages({ +// nonce: BigInt(0), +// }); +// expect(messagesRes).toMatchObject({ +// Ok: { +// messages: [], +// cert: expect.any(Uint8Array), +// tree: expect.any(Uint8Array), +// }, +// }); +// }); +// }); + +describe.only("Canister - ws_get_messages (receive)", () => { beforeAll(async () => { - await assignKeyPairsToClients(); - - await wsRegister({ - clientActor: client1, - clientKey: client1KeyPair.publicKey, - }, true); + await assignKeysToClients(); - await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client1KeyPair.secretKey, - canisterId, - gatewayActor: gateway1, - }, true); - - await wsSend({ - clientPublicKey: client1KeyPair.publicKey, - actor: client1, - message: { text: "test" }, - }, true); - }); - - afterAll(async () => { - await wsWipe(gateway1); - }); - - it("fails if a non registered gateway sends an IcWebSocketGatewayStatus message", async () => { - const res = await wsMessage({ - message: { - IcWebSocketGatewayStatus: { - status_index: BigInt(1), - }, - }, - actor: gateway2, - }); - - expect(res).toMatchObject({ - Err: "caller is not the gateway that has been registered during CDK initialization", - }); - }); - - it("registered gateway should update the status index", async () => { - const res = await wsMessage({ - message: { - IcWebSocketGatewayStatus: { - status_index: BigInt(2), // set it high to test behavior for indexes behind the current one - }, - }, - actor: gateway1, - }); - - expect(res).toMatchObject({ - Ok: null, - }); - }); - - it("fails if a registered gateway sends an IcWebSocketGatewayStatus with a wrong status index (equal to current)", async () => { - const res = await wsMessage({ - message: { - IcWebSocketGatewayStatus: { - status_index: BigInt(2), - }, - }, - actor: gateway1, - }); - - expect(res).toMatchObject({ - Err: "Gateway status index is equal to or behind the current one", - }); - }); - - it("fails if a registered gateway sends an IcWebSocketGatewayStatus with a wrong status index (behind the current)", async () => { - const res = await wsMessage({ - message: { - IcWebSocketGatewayStatus: { - status_index: BigInt(1), - }, - }, - actor: gateway1, - }); - - expect(res).toMatchObject({ - Err: "Gateway status index is equal to or behind the current one", - }); - }); - - it("registered gateway should disconnect after maximum time", async () => { - let res = await gateway1.ws_get_messages({ - nonce: BigInt(0), - }); - - expect(res).toMatchObject({ - Ok: { - messages: expect.any(Array), - cert: expect.any(Uint8Array), - tree: expect.any(Uint8Array), - }, - }); - expect((res as { Ok: CanisterOutputCertifiedMessages }).Ok.messages.length).toEqual(1); - - // wait for the maximum time the gateway can send a status message, - // so that the internal canister state is reset - // double the time to make sure the canister state is reset - await new Promise((resolve) => setTimeout(resolve, 2 * MAX_GATEWAY_KEEP_ALIVE_TIME_MS)); - - // check if messages have been deleted - res = await gateway1.ws_get_messages({ - nonce: BigInt(0), - }); - expect(res).toMatchObject({ - Ok: { - messages: [], - cert: expect.any(Uint8Array), - tree: expect.any(Uint8Array), - }, + // reset the internal timers + await reinitialize({ + sendAckIntervalMs: DEFAULT_TEST_SEND_ACK_INTERVAL_MS, + keepAliveDelayMs: DEFAULT_TEST_KEEP_ALIVE_DELAY_MS, }); - // check if registered client has been deleted - const sendRes = await wsSend({ - clientPublicKey: client1KeyPair.publicKey, - actor: client1, - message: { text: "test" }, - }); - expect(sendRes).toMatchObject({ - Err: "client's public key has not been previously registered by client", - }); - }); - - it("registered gateway should reconnect by resetting the status index", async () => { - let res = await wsMessage({ - message: { - IcWebSocketGatewayStatus: { - status_index: BigInt(0), - }, - }, - actor: gateway1, - }); - - expect(res).toMatchObject({ - Ok: null, - }); - - res = await wsMessage({ - message: { - IcWebSocketGatewayStatus: { - status_index: BigInt(1), - }, - }, - actor: gateway1, - }); - - expect(res).toMatchObject({ - Ok: null, - }); - }); - - it("registered gateway should reconnect before maximum time", async () => { - // reconnect the client - await wsRegister({ - clientActor: client1, - clientKey: client1KeyPair.publicKey, - }, true); - await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client1KeyPair.secretKey, + clientNonce: client1Key.client_nonce, canisterId, - gatewayActor: gateway1, - }, true); - - // send a test message from the canister to check if the internal state is reset - await wsSend({ - clientPublicKey: client1KeyPair.publicKey, - actor: client1, - message: { text: "test" }, - }, true); - - // check if the canister has the message in the queue - let messagesRes = await gateway1.ws_get_messages({ - nonce: BigInt(0), - }); - expect(messagesRes).toMatchObject({ - Ok: { - messages: expect.any(Array), - cert: expect.any(Uint8Array), - tree: expect.any(Uint8Array), - }, - }); - expect((messagesRes as { Ok: CanisterOutputCertifiedMessages }).Ok.messages.length).toEqual(1); - - // simulate a reconnection - const res = await wsMessage({ - message: { - IcWebSocketGatewayStatus: { - status_index: BigInt(0), - }, - }, - actor: gateway1, - }); - expect(res).toMatchObject({ - Ok: null, - }); - - // check if the canister reset the internal state - messagesRes = await gateway1.ws_get_messages({ - nonce: BigInt(0), - }); - expect(messagesRes).toMatchObject({ - Ok: { - messages: [], - cert: expect.any(Uint8Array), - tree: expect.any(Uint8Array), - }, - }); - }); -}); - -describe("Canister - ws_get_messages (receive)", () => { - beforeAll(async () => { - await assignKeyPairsToClients(); - - await wsRegister({ clientActor: client1, - clientKey: client1KeyPair.publicKey, - }, true); - - await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client1KeyPair.secretKey, - canisterId, - gatewayActor: gateway1, }, true); // prepare the messages for (let i = 0; i < SEND_MESSAGES_COUNT; i++) { const appMessage = { text: `test${i}` }; await wsSend({ - clientPublicKey: client1KeyPair.publicKey, + clientPrincipal: client1Key.client_principal, actor: client1, message: appMessage, }, true); @@ -925,15 +581,13 @@ describe("Canister - ws_get_messages (receive)", () => { }); afterAll(async () => { - await wsWipe(gateway1); - }); - - beforeEach(async () => { - await sendGatewayStatusMessage(); + await wsWipe(); }); it("registered gateway can receive correct amount of messages", async () => { - for (let i = 0; i < SEND_MESSAGES_COUNT; i++) { + // on open, the canister puts a service message in the queue + const messagesCount = SEND_MESSAGES_COUNT + 1; // +1 for the service message + for (let i = 0; i < messagesCount; i++) { const res = await gateway1.ws_get_messages({ nonce: BigInt(i), }); @@ -948,15 +602,15 @@ describe("Canister - ws_get_messages (receive)", () => { const messagesResult = (res as { Ok: CanisterOutputCertifiedMessages }).Ok; expect(messagesResult.messages.length).toBe( - SEND_MESSAGES_COUNT - i > MAX_NUMBER_OF_RETURNED_MESSAGES + messagesCount - i > MAX_NUMBER_OF_RETURNED_MESSAGES ? MAX_NUMBER_OF_RETURNED_MESSAGES - : SEND_MESSAGES_COUNT - i + : messagesCount - i ); } // try to get more messages than available const res = await gateway1.ws_get_messages({ - nonce: BigInt(SEND_MESSAGES_COUNT), + nonce: BigInt(messagesCount), }); expect(res).toMatchObject({ @@ -971,21 +625,25 @@ describe("Canister - ws_get_messages (receive)", () => { it("registered gateway can receive certified messages", async () => { // first batch of messages const firstBatchRes = await gateway1.ws_get_messages({ - nonce: BigInt(0), + nonce: BigInt(1), }); const firstBatchMessagesResult = (firstBatchRes as { Ok: CanisterOutputCertifiedMessages }).Ok; + console.log(firstBatchMessagesResult.messages.map((msg) => msg.key)); for (let i = 0; i < firstBatchMessagesResult.messages.length; i++) { const message = firstBatchMessagesResult.messages[i]; - expect(message.client_key).toEqual(client1KeyPair.publicKey); - const decodedContent = Cbor.decode(new Uint8Array(message.content)); - expect(decodedContent).toMatchObject({ - client_key: client1KeyPair.publicKey, - message: expect.any(Uint8Array), - sequence_num: i + 1, - timestamp: expect.any(Object), // weird timestamp deserialization + expect(isClientKeyEq(message.client_key, client1Key)).toEqual(true); + const websocketMessage = decodeWebsocketMessage(new Uint8Array(message.content)); + console.log(websocketMessage); + expect(websocketMessage).toMatchObject({ + client_key: expect.any(Object), + content: expect.any(Uint8Array), + sequence_num: BigInt(i + 1), + timestamp: expect.any(Object), // weird cbor bigint deserialization + is_service_message: false, }); - expect(IDL.decode([IDL.Record({ 'text': IDL.Text })], decodedContent.message as Uint8Array)).toEqual([{ text: `test${i}` }]); + expect(isClientKeyEq(websocketMessage.client_key, client1Key)).toEqual(true); + expect(IDL.decode([IDL.Record({ 'text': IDL.Text })], websocketMessage.content as Uint8Array)).toEqual([{ text: `test${i}` }]); // check the certification await expect( @@ -1013,15 +671,17 @@ describe("Canister - ws_get_messages (receive)", () => { const secondBatchMessagesResult = (secondBatchRes as { Ok: CanisterOutputCertifiedMessages }).Ok; for (let i = 0; i < secondBatchMessagesResult.messages.length; i++) { const message = secondBatchMessagesResult.messages[i]; - expect(message.client_key).toEqual(client1KeyPair.publicKey); - const decodedContent = Cbor.decode(new Uint8Array(message.content)); - expect(decodedContent).toMatchObject({ - client_key: client1KeyPair.publicKey, - message: expect.any(Uint8Array), - sequence_num: i + MAX_NUMBER_OF_RETURNED_MESSAGES + 1, - timestamp: expect.any(Object), // weird timestamp deserialization + expect(isClientKeyEq(message.client_key, client1Key)).toEqual(true); + const websocketMessage = decodeWebsocketMessage(new Uint8Array(message.content)); + expect(websocketMessage).toMatchObject({ + client_key: expect.any(Object), + content: expect.any(Uint8Array), + sequence_num: BigInt(i + MAX_NUMBER_OF_RETURNED_MESSAGES + 1), + timestamp: expect.any(Object), // weird cbor bigint deserialization + is_service_message: false, }); - expect(IDL.decode([IDL.Record({ 'text': IDL.Text })], decodedContent.message as Uint8Array)).toEqual([{ text: `test${i + MAX_NUMBER_OF_RETURNED_MESSAGES}` }]); + expect(isClientKeyEq(websocketMessage.client_key, client1Key)).toEqual(true); + expect(IDL.decode([IDL.Record({ 'text': IDL.Text })], websocketMessage.content as Uint8Array)).toEqual([{ text: `test${i + MAX_NUMBER_OF_RETURNED_MESSAGES}` }]); // check the certification await expect( @@ -1043,107 +703,103 @@ describe("Canister - ws_get_messages (receive)", () => { }); }); -describe("Canister - ws_close", () => { - beforeAll(async () => { - await assignKeyPairsToClients(); - - await wsRegister({ - clientActor: client1, - clientKey: client1KeyPair.publicKey, - }, true); - - await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client1KeyPair.secretKey, - canisterId, - gatewayActor: gateway1, - }, true); - }); - - afterAll(async () => { - await wsWipe(gateway1); - }); - - beforeEach(async () => { - await sendGatewayStatusMessage(); - }); - - it("fails if gateway is not registered", async () => { - const res = await wsClose({ - clientPublicKey: client1KeyPair.publicKey, - gatewayActor: gateway2, - }); - - expect(res).toMatchObject({ - Err: "caller is not the gateway that has been registered during CDK initialization", - }); - }); - - it("fails if client is not registered", async () => { - const res = await wsClose({ - clientPublicKey: client2KeyPair.publicKey, - gatewayActor: gateway1, - }); - - expect(res).toMatchObject({ - Err: "client's public key has not been previously registered by client", - }); - }); - - it("should close the websocket for a registered client", async () => { - const res = await wsClose({ - clientPublicKey: client1KeyPair.publicKey, - gatewayActor: gateway1, - }); - - expect(res).toMatchObject({ - Ok: null, - }); - }); -}); - -describe("Canister - ws_send", () => { - beforeAll(async () => { - await assignKeyPairsToClients(); - - await wsRegister({ - clientActor: client1, - clientKey: client1KeyPair.publicKey, - }, true); - - await wsOpen({ - clientPublicKey: client1KeyPair.publicKey, - clientSecretKey: client1KeyPair.secretKey, - canisterId, - gatewayActor: gateway1, - }, true); - }); - - afterAll(async () => { - await wsWipe(gateway1); - }); - - it("fails if sending a message to a non registered client", async () => { - const res = await wsSend({ - clientPublicKey: client2KeyPair.publicKey, - actor: client1, - message: { text: "test" }, - }); - - expect(res).toMatchObject({ - Err: "client's public key has not been previously registered by client", - }); - }); - - it("should send a message to a registered client", async () => { - const res = await wsSend({ - clientPublicKey: client1KeyPair.publicKey, - actor: client1, - message: { text: "test" }, - }); - - expect(res).toMatchObject({ - Ok: null, - }); - }); -}); +// describe("Canister - ws_close", () => { +// beforeAll(async () => { +// await assignKeysToClients(); + +// await wsRegister({ +// clientActor: client1, +// clientKey: client1Key.publicKey, +// }, true); + +// await wsOpen({ +// clientPublicKey: client1Key.publicKey, +// clientSecretKey: client1Key.secretKey, +// canisterId, +// clientActor: gateway1, +// }, true); +// }); + +// afterAll(async () => { +// await wsWipe(gateway1); +// }); + +// it("fails if gateway is not registered", async () => { +// const res = await wsClose({ +// clientPublicKey: client1Key.publicKey, +// gatewayActor: gateway2, +// }); + +// expect(res).toMatchObject({ +// Err: "caller is not the gateway that has been registered during CDK initialization", +// }); +// }); + +// it("fails if client is not registered", async () => { +// const res = await wsClose({ +// clientPublicKey: client2Key.publicKey, +// gatewayActor: gateway1, +// }); + +// expect(res).toMatchObject({ +// Err: "client's public key has not been previously registered by client", +// }); +// }); + +// it("should close the websocket for a registered client", async () => { +// const res = await wsClose({ +// clientPublicKey: client1Key.publicKey, +// gatewayActor: gateway1, +// }); + +// expect(res).toMatchObject({ +// Ok: null, +// }); +// }); +// }); + +// describe("Canister - ws_send", () => { +// beforeAll(async () => { +// await assignKeysToClients(); + +// await wsRegister({ +// clientActor: client1, +// clientKey: client1Key.publicKey, +// }, true); + +// await wsOpen({ +// clientPublicKey: client1Key.publicKey, +// clientSecretKey: client1Key.secretKey, +// canisterId, +// clientActor: gateway1, +// }, true); +// }); + +// afterAll(async () => { +// await wsWipe(gateway1); +// }); + +// it("fails if sending a message to a non registered client", async () => { +// const res = await wsSend({ +// clientKey: client2Key.publicKey, +// actor: client1, +// message: { text: "test" }, +// }); + +// expect(res).toMatchObject({ +// Err: "client's public key has not been previously registered by client", +// }); +// }); + +// it("should send a message to a registered client", async () => { +// const res = await wsSend({ +// clientKey: client1Key.publicKey, +// actor: client1, +// message: { text: "test" }, +// }); + +// expect(res).toMatchObject({ +// Ok: null, +// }); +// }); +// }); diff --git a/tests/integration/utils/actors.ts b/tests/integration/utils/actors.ts index 46635a5..525be15 100644 --- a/tests/integration/utils/actors.ts +++ b/tests/integration/utils/actors.ts @@ -84,3 +84,5 @@ export const client2 = createActor(canisterId, { identity: client2Data.identity, }, }); + +export const anonymousClient = createActor(canisterId); diff --git a/tests/integration/utils/api.ts b/tests/integration/utils/api.ts index ef33d37..a380a87 100644 --- a/tests/integration/utils/api.ts +++ b/tests/integration/utils/api.ts @@ -4,124 +4,120 @@ import { ActorSubclass, Cbor, Certificate, HashTree, HttpAgent, compare, lookup_ import { Secp256k1KeyIdentity } from "@dfinity/identity-secp256k1"; import { Principal } from "@dfinity/principal"; import { IDL } from "@dfinity/candid"; -import { getMessageSignature } from "./crypto"; -import type { CanisterIncomingMessage, ClientPublicKey, _SERVICE } from "../../src/declarations/test_canister/test_canister.did"; +import { anonymousClient, gateway1Data } from "./actors"; +import type { CanisterOutputCertifiedMessages, ClientKey, ClientPrincipal, WebsocketMessage, _SERVICE } from "../../src/declarations/test_canister/test_canister.did"; -type WsRegisterArgs = { - clientActor: ActorSubclass<_SERVICE>, - clientKey: Uint8Array, +type GenericResult = { + Ok: T, +} | { + Err: string, }; -export const wsRegister = async (args: WsRegisterArgs, throwIfError = false) => { - const res = await args.clientActor.ws_register({ - client_key: args.clientKey, - }); - - if (throwIfError) { - if ('Err' in res) { - throw new Error(res.Err); - } +const resolveResult = (result: GenericResult, throwIfError: boolean) => { + if (throwIfError && 'Err' in result) { + throw new Error(result.Err); } - return res; + return result; }; type WsOpenArgs = { - clientPublicKey: Uint8Array, - clientSecretKey: Uint8Array | string, + clientNonce: bigint, canisterId: string, - gatewayActor: ActorSubclass<_SERVICE>, -}; - -export type CanisterOpenMessageContent = { - client_key: ClientPublicKey, - canister_id: Principal, + clientActor: ActorSubclass<_SERVICE>, }; +/** + * Sends an update call to the canister to the **ws_open** method, using the provided actor. + * @param args {@link WsOpenArgs} + * @param throwIfError whether to throw if the result is an error (defaults to `false`) + * @returns the result of the **ws_open** method + */ export const wsOpen = async (args: WsOpenArgs, throwIfError = false) => { - const firstMessage: CanisterOpenMessageContent = { - client_key: args.clientPublicKey, - canister_id: Principal.fromText(args.canisterId), - }; - const contentBuf = new Uint8Array(Cbor.encode(firstMessage)); - const sig = await getMessageSignature(contentBuf, args.clientSecretKey); - - const res = await args.gatewayActor.ws_open({ - content: contentBuf, - sig, + const res = await args.clientActor.ws_open({ + client_nonce: args.clientNonce, }); - if (throwIfError) { - if ('Err' in res) { - throw new Error(res.Err); - } - } - - return res; + return resolveResult(res, throwIfError); }; type WsMessageArgs = { - message: CanisterIncomingMessage, + message: WebsocketMessage, actor: ActorSubclass<_SERVICE>, }; +/** + * Sends an update call to the canister to the **ws_message** method, using the provided actor. + * @param args {@link WsMessageArgs} + * @param throwIfError whether to throw if the result is an error (defaults to `false`) + * @returns the result of the **ws_message** method + */ export const wsMessage = async (args: WsMessageArgs, throwIfError = false) => { const res = await args.actor.ws_message({ msg: args.message, }); - if (throwIfError) { - if ('Err' in res) { - throw new Error(res.Err); - } - } - - return res; + return resolveResult(res, throwIfError); }; -export type WebsocketMessage = { - client_key: ClientPublicKey, - sequence_num: number, - timestamp: number, - message: ArrayBuffer | Uint8Array, +type WsCloseArgs = { + clientKey: ClientKey, + gatewayActor: ActorSubclass<_SERVICE>, }; -export const getWebsocketMessage = (clientPublicKey: ClientPublicKey, sequenceNumber: number, content?: ArrayBuffer | Uint8Array): Uint8Array => { - const websocketMessage: WebsocketMessage = { - client_key: clientPublicKey, - sequence_num: sequenceNumber, - timestamp: Date.now(), - message: new Uint8Array(content || []), - }; +/** + * Sends an update call to the canister to the **ws_close** method, using the provided gateway actor. + * @param args {@link WsCloseArgs} + * @param throwIfError whether to throw if the result is an error (defaults to `false`) + * @returns the result of the **ws_close** method + */ +export const wsClose = async (args: WsCloseArgs, throwIfError = false) => { + const res = await args.gatewayActor.ws_close({ + client_key: args.clientKey, + }); - return new Uint8Array(Cbor.encode(websocketMessage)); + return resolveResult(res, throwIfError); }; -type WsCloseArgs = { - clientPublicKey: Uint8Array, +type WsGetMessagesArgs = { + fromNonce: number, gatewayActor: ActorSubclass<_SERVICE>, }; -export const wsClose = async (args: WsCloseArgs, throwIfError = false) => { - const res = await args.gatewayActor.ws_close({ - client_key: args.clientPublicKey, +/** + * Sends a query call to the canister to the **ws_get_messages** method, using the provided gateway actor. + * @param args {@link WsGetMessagesArgs} + */ +export const wsGetMessages = async (args: WsGetMessagesArgs): Promise => { + const res = await args.gatewayActor.ws_get_messages({ + nonce: BigInt(args.fromNonce), }); - if (throwIfError) { - if ('Err' in res) { - throw new Error(res.Err); - } - } + const messages = resolveResult(res, true); - return res; + return (messages as { Ok: CanisterOutputCertifiedMessages }).Ok; }; -export const wsWipe = async (gatewayActor: ActorSubclass<_SERVICE>) => { - await gatewayActor.ws_wipe(); +export const wsWipe = async () => { + await anonymousClient.ws_wipe(); +}; + +type ReinitializeArgs = { + sendAckIntervalMs: number, + keepAliveDelayMs: number, +}; + +/** + * Used to reinitialize the canister with the provided intervals. + * @param args {@link ReinitializeArgs} + */ +export const reinitialize = async (args: ReinitializeArgs) => { + const gatewayPrincipal = (await gateway1Data.identity).getPrincipal().toText(); + await anonymousClient.reinitialize(gatewayPrincipal, BigInt(args.sendAckIntervalMs), BigInt(args.keepAliveDelayMs)); }; type WsSendArgs = { - clientPublicKey: Uint8Array, + clientPrincipal: ClientPrincipal, actor: ActorSubclass<_SERVICE>, message: { text: string, @@ -130,15 +126,9 @@ type WsSendArgs = { export const wsSend = async (args: WsSendArgs, throwIfError = false) => { const msgBytes = IDL.encode([IDL.Record({ 'text': IDL.Text })], [args.message]); - const res = await args.actor.ws_send(args.clientPublicKey, new Uint8Array(msgBytes)); + const res = await args.actor.ws_send(args.clientPrincipal, new Uint8Array(msgBytes)); - if (throwIfError) { - if ('Err' in res) { - throw new Error(res.Err); - } - } - - return res; + return resolveResult(res, throwIfError); }; export const getCertifiedMessageKey = async (gatewayIdentity: Promise, nonce: number) => { @@ -146,7 +136,6 @@ export const getCertifiedMessageKey = async (gatewayIdentity: Promise { const canisterPrincipal = Principal.fromText(canisterId); let cert: Certificate; diff --git a/tests/integration/utils/crypto.ts b/tests/integration/utils/crypto.ts deleted file mode 100644 index a46d653..0000000 --- a/tests/integration/utils/crypto.ts +++ /dev/null @@ -1,22 +0,0 @@ -import * as ed from "@noble/ed25519"; - -export const getKeyPair = async (secretKey?: string | Uint8Array): Promise<{ publicKey: Uint8Array, secretKey: Uint8Array | string }> => { - if (!secretKey) { - secretKey = ed.utils.randomPrivateKey(); - } - - const publicKey = await ed.getPublicKeyAsync(secretKey); - - return { - publicKey, - secretKey, - }; -}; - -export const getMessageSignature = async (buf: ArrayBuffer | Uint8Array, secretKey: Uint8Array | string): Promise => { - // Sign the message so that the gateway can verify canister and client ids match - const toSign = new Uint8Array(buf); - const sig = await ed.signAsync(toSign, secretKey); - - return sig; -} diff --git a/tests/integration/utils/identity.ts b/tests/integration/utils/identity.ts index 21a1d8b..cfac2ee 100644 --- a/tests/integration/utils/identity.ts +++ b/tests/integration/utils/identity.ts @@ -9,3 +9,8 @@ export const identityFromSeed = async (phrase: string) => { return Secp256k1KeyIdentity.generate(addrnode.privateKey); }; + +export const generateRandomIdentity = async () => { + const mnemonic = bip39.generateMnemonic(); + return identityFromSeed(mnemonic); +}; diff --git a/tests/integration/utils/idl.ts b/tests/integration/utils/idl.ts new file mode 100644 index 0000000..317486e --- /dev/null +++ b/tests/integration/utils/idl.ts @@ -0,0 +1,67 @@ +import { IDL } from "@dfinity/candid"; +import { CanisterOutputMessage, ClientKey, WebsocketMessage } from "../../src/declarations/test_canister/test_canister.did"; +import { Cbor } from "@dfinity/agent"; + +export const ClientPrincipalIdl = IDL.Principal; +export const ClientKeyIdl = IDL.Record({ + 'client_principal': ClientPrincipalIdl, + 'client_nonce': IDL.Nat64, +}); + +export type CanisterOpenMessageContent = { + 'client_key': ClientKey, +}; +export type CanisterAckMessageContent = { + 'last_incoming_sequence_num': bigint, +}; +export type ClientKeepAliveMessageContent = { + 'last_incoming_sequence_num': bigint, +}; +export type WebsocketServiceMessageContent = { + OpenMessage: CanisterOpenMessageContent, +} | { + AckMessage: CanisterAckMessageContent, +} | { + KeepAliveMessage: ClientKeepAliveMessageContent, +}; + +export const CanisterOpenMessageContentIdl = IDL.Record({ + 'client_key': ClientKeyIdl, +}); +export const CanisterAckMessageContentIdl = IDL.Record({ + 'last_incoming_sequence_num': IDL.Nat64, +}); +export const ClientKeepAliveMessageContentIdl = IDL.Record({ + 'last_incoming_sequence_num': IDL.Nat64, +}); +export const WebsocketServiceMessageContentIdl = IDL.Variant({ + 'OpenMessage': CanisterOpenMessageContentIdl, + 'AckMessage': CanisterAckMessageContentIdl, + 'KeepAliveMessage': ClientKeepAliveMessageContentIdl, +}); + +export const decodeWebsocketServiceMessageContent = (bytes: Uint8Array): WebsocketServiceMessageContent => { + const decoded = IDL.decode([WebsocketServiceMessageContentIdl], bytes); + if (decoded.length !== 1) { + throw new Error("Invalid CanisterServiceMessage"); + } + return decoded[0] as unknown as WebsocketServiceMessageContent; +}; + +export const encodeWebsocketServiceMessageContent = (msg: WebsocketServiceMessageContent): Uint8Array => { + return new Uint8Array(IDL.encode([WebsocketServiceMessageContentIdl], [msg])); +}; + +export const isClientKeyEq = (a: ClientKey, b: ClientKey): boolean => { + return a.client_principal.compareTo(b.client_principal) === "eq" && a.client_nonce === b.client_nonce; +} + +export const getServiceMessageFromCanisterMessage = (msg: CanisterOutputMessage): WebsocketServiceMessageContent => { + const content = getWebsocketMessageFromCanisterMessage(msg).content; + return decodeWebsocketServiceMessageContent(content as Uint8Array); +} + +export const getWebsocketMessageFromCanisterMessage = (msg: CanisterOutputMessage): WebsocketMessage => { + const websocketMessage: WebsocketMessage = Cbor.decode(msg.content as Uint8Array); + return websocketMessage; +} diff --git a/tests/integration/utils/messages.ts b/tests/integration/utils/messages.ts new file mode 100644 index 0000000..993e36c --- /dev/null +++ b/tests/integration/utils/messages.ts @@ -0,0 +1,31 @@ +import { Cbor } from "@dfinity/agent"; +import { CanisterOutputMessage, ClientKey, WebsocketMessage } from "../../src/declarations/test_canister/test_canister.did"; +import { getWebsocketMessageFromCanisterMessage } from "./idl"; + +export const filterServiceMessagesFromCanisterMessages = (messages: CanisterOutputMessage[]): CanisterOutputMessage[] => { + return messages.filter((msg) => { + const websocketMessage = getWebsocketMessageFromCanisterMessage(msg); + return websocketMessage.is_service_message; + }); +}; + +export const createWebsocketMessage = ( + clientKey: ClientKey, + sequenceNumber: number, + isServiceMessage = false, + content?: ArrayBuffer | Uint8Array +): WebsocketMessage => { + const websocketMessage: WebsocketMessage = { + client_key: clientKey, + sequence_num: BigInt(sequenceNumber), + timestamp: BigInt(Date.now()) * BigInt(1_000_000), // in nanoseconds + content: new Uint8Array(content || []), + is_service_message: isServiceMessage, + }; + + return websocketMessage; +}; + +export const decodeWebsocketMessage = (bytes: Uint8Array): WebsocketMessage => { + return Cbor.decode(bytes); +}; diff --git a/tests/integration/utils/random.ts b/tests/integration/utils/random.ts new file mode 100644 index 0000000..d469381 --- /dev/null +++ b/tests/integration/utils/random.ts @@ -0,0 +1,20 @@ +import { ClientKey, ClientPrincipal } from "../../src/declarations/test_canister/test_canister.did"; +import { generateRandomIdentity } from "./identity"; + +export const getRandomClientNonce = (): bigint => { + const array = new BigUint64Array(1); + globalThis.crypto.getRandomValues(array); + return array[0]; +}; + +export const generateClientKey = (clientPrincipal: ClientPrincipal): ClientKey => { + return { + client_principal: clientPrincipal, + client_nonce: getRandomClientNonce(), + }; +}; + +export const getRandomPrincipal = async (): Promise => { + const identity = await generateRandomIdentity(); + return identity.getPrincipal(); +}; diff --git a/tests/package.json b/tests/package.json index 240d132..7313deb 100644 --- a/tests/package.json +++ b/tests/package.json @@ -10,7 +10,7 @@ ], "scripts": { "generate": "dfx generate test_canister", - "deploy:tests": "dfx deploy test_canister --no-wallet --argument '(\"i3gux-m3hwt-5mh2w-t7wwm-fwx5j-6z6ht-hxguo-t4rfw-qp24z-g5ivt-2qe\")'", + "deploy:tests": "dfx deploy test_canister --no-wallet --argument '(\"i3gux-m3hwt-5mh2w-t7wwm-fwx5j-6z6ht-hxguo-t4rfw-qp24z-g5ivt-2qe\", 300_000 : nat64, 300_000 : nat64)'", "test:integration": "jest integration" }, "devDependencies": { diff --git a/tests/src/lib.rs b/tests/src/lib.rs index 9e09999..26932a3 100644 --- a/tests/src/lib.rs +++ b/tests/src/lib.rs @@ -11,7 +11,7 @@ use ic_websocket_cdk::{ mod canister; #[init] -fn init(gateway_principal: String) { +fn init(gateway_principal: String, send_ack_interval_ms: u64, keep_alive_delay_ms: u64) { let handlers = WsHandlers { on_open: Some(on_open), on_message: Some(on_message), @@ -21,16 +21,16 @@ fn init(gateway_principal: String) { let params = WsInitParams { handlers, gateway_principal, - send_ack_interval_ms: 10_000, - keep_alive_delay_ms: 5_000, + send_ack_interval_ms, + keep_alive_delay_ms, }; ic_websocket_cdk::init(params) } #[post_upgrade] -fn post_upgrade(gateway_principal: String) { - init(gateway_principal); +fn post_upgrade(gateway_principal: String, send_ack_interval_ms: u64, keep_alive_delay_ms: u64) { + init(gateway_principal, send_ack_interval_ms, keep_alive_delay_ms); } // method called by the WS Gateway after receiving FirstMessage from the client @@ -69,3 +69,10 @@ fn ws_wipe() { fn ws_send(client_principal: ClientPrincipal, msg_bytes: Vec) -> CanisterWsSendResult { ic_websocket_cdk::ws_send(client_principal, msg_bytes) } + +// reinitialize the canister +#[update] +fn reinitialize(gateway_principal: String, send_ack_interval_ms: u64, keep_alive_delay_ms: u64) { + ic_websocket_cdk::wipe(); + init(gateway_principal, send_ack_interval_ms, keep_alive_delay_ms); +} diff --git a/tests/test_canister.did b/tests/test_canister.did index a98d79e..da066d3 100644 --- a/tests/test_canister.did +++ b/tests/test_canister.did @@ -5,7 +5,7 @@ type CanisterWsSendResult = variant { Err : text; }; -service : (text) -> { +service : (text, nat64, nat64) -> { "ws_open" : (CanisterWsOpenArguments) -> (CanisterWsOpenResult); "ws_close" : (CanisterWsCloseArguments) -> (CanisterWsCloseResult); "ws_message" : (CanisterWsMessageArguments) -> (CanisterWsMessageResult); @@ -14,4 +14,5 @@ service : (text) -> { // methods used just for debugging/testing "ws_wipe" : () -> (); "ws_send" : (ClientPrincipal, blob) -> (CanisterWsSendResult); + "reinitialize" : (text, nat64, nat64) -> (); }; From 3c2ec90f42b41a2ea34fa89173c6b25f41d5f22f Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Mon, 9 Oct 2023 13:45:13 +0200 Subject: [PATCH 2/4] chore: unit test panic handling --- Cargo.lock | 255 +++++++++++++++++++------------- src/ic-websocket-cdk/src/lib.rs | 43 +++++- 2 files changed, 190 insertions(+), 108 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 61394ca..baab2d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -169,9 +169,9 @@ checksum = "7f30e7476521f6f8af1a1c4c0b8cc94f0bee37d91763d0ca2665f299b6cd8aec" [[package]] name = "byteorder" -version = "1.4.3" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" @@ -181,9 +181,9 @@ checksum = "a2bd12c1caf447e69cd4528f47f94d203fd2582878ecb9e9465484c4148a8223" [[package]] name = "candid" -version = "0.9.7" +version = "0.9.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f391a0d11d997af68e1a06b5e2ab354079cecb82b6eefb26addb38adf66d351d" +checksum = "aa0f00717c71b8e9ee4c090b4880ec2418c8506bb6828a2c72df72d3896e905d" dependencies = [ "anyhow", "binread", @@ -201,7 +201,7 @@ dependencies = [ "pretty", "serde", "serde_bytes", - "sha2 0.10.7", + "sha2 0.10.8", "stacker", "thiserror", ] @@ -215,7 +215,7 @@ dependencies = [ "lazy_static", "proc-macro2", "quote", - "syn 2.0.37", + "syn 2.0.38", ] [[package]] @@ -249,6 +249,22 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28c122c3980598d243d63d9a704629a2d748d101f278052ff068be5a4423ab6f" +[[package]] +name = "core-foundation" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "194a7a9e6de53fa55116934067c844d9d749312f75c6f6d0980e8c252f8c2146" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e496a50fda8aacccc86d7529e2c1e0892dbd0f898a6b5645b5561b89c3210efa" + [[package]] name = "cpufeatures" version = "0.2.9" @@ -355,9 +371,9 @@ checksum = "a26ae43d7bcc3b814de94796a5e736d4029efb0ee900c12e2d54c993ad1a1e07" [[package]] name = "elliptic-curve" -version = "0.13.5" +version = "0.13.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "968405c8fdc9b3bf4df0a6638858cc0b52462836ab6b1c87377785dd09cf1c0b" +checksum = "d97ca172ae9dc9f9b779a6e3a65d308f2af74e5b8c921299075bdb4a0370e914" dependencies = [ "base16ct", "crypto-bigint", @@ -390,30 +406,19 @@ checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" [[package]] name = "errno" -version = "0.3.3" +version = "0.3.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "136526188508e25c6fef639d7927dfb3e0e3084488bf202267829cf7fc23dbdd" +checksum = "ac3e13f66a2f95e32a39eaa81f6b95d42878ca0e1db0c7543723dfe12557e860" dependencies = [ - "errno-dragonfly", "libc", "windows-sys", ] -[[package]] -name = "errno-dragonfly" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa68f1b12764fab894d2755d2518754e71b4fd80ecfb822714a1206c2aab39bf" -dependencies = [ - "cc", - "libc", -] - [[package]] name = "fastrand" -version = "2.0.0" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6999dc1837253364c2ebb0704ba97994bd874e8f195d665c50b7548f6ea92764" +checksum = "25cbce373ec4653f1a01a31e8a5e5ec0c622dc27ff9c4e6606eefef5cbbed4a5" [[package]] name = "ff" @@ -506,7 +511,7 @@ checksum = "89ca545a94061b6365f2c7355b4b32bd20df3ff95f02da9329b34ccc3bd6ee72" dependencies = [ "proc-macro2", "quote", - "syn 2.0.37", + "syn 2.0.38", ] [[package]] @@ -622,9 +627,9 @@ checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" [[package]] name = "hashbrown" -version = "0.14.0" +version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c6201b9ff9fd90a5a3bac2e56a830d0caa509576f0e503818ee82c181b3437a" +checksum = "7dfda62a12f55daeae5015f81b0baea145391cb4520f86c248fc615d72640d12" [[package]] name = "hermit-abi" @@ -739,14 +744,14 @@ dependencies = [ "pkcs8", "rand", "reqwest", - "ring", + "ring 0.16.20", "rustls 0.20.9", "sec1", "serde", "serde_bytes", "serde_cbor", "serde_repr", - "sha2 0.10.7", + "sha2 0.10.8", "simple_asn1", "thiserror", "time", @@ -804,7 +809,7 @@ dependencies = [ "hex", "serde", "serde_bytes", - "sha2 0.10.7", + "sha2 0.10.8", ] [[package]] @@ -815,7 +820,7 @@ checksum = "197524aecec47db0b6c0c9f8821aad47272c2bd762c7a0ffe9715eaca0364061" dependencies = [ "serde", "serde_bytes", - "sha2 0.10.7", + "sha2 0.10.8", ] [[package]] @@ -843,11 +848,11 @@ dependencies = [ "ic-certified-map", "proptest", "rand", - "ring", + "ring 0.16.20", "serde", "serde_bytes", "serde_cbor", - "sha2 0.10.7", + "sha2 0.10.8", ] [[package]] @@ -878,12 +883,12 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.0.0" +version = "2.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5477fe2230a79769d8dc68e0eabf5437907c0457a5614a9e8dddb67f65eb65d" +checksum = "8adf3ddd720272c6ea8bf59463c04e0f93d0bbf7c5439b691bca2987e0270897" dependencies = [ "equivalent", - "hashbrown 0.14.0", + "hashbrown 0.14.1", ] [[package]] @@ -926,7 +931,7 @@ dependencies = [ "ecdsa", "elliptic-curve", "once_cell", - "sha2 0.10.7", + "sha2 0.10.8", "signature", ] @@ -944,21 +949,21 @@ checksum = "884e2677b40cc8c339eaefcb701c32ef1fd2493d71118dc0ca4b6a736c93bd67" [[package]] name = "libc" -version = "0.2.148" +version = "0.2.149" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9cdc71e17332e86d2e1d38c1f99edcb6288ee11b815fb1a4b049eaa2114d369b" +checksum = "a08173bc88b7955d1b3145aa561539096c421ac8debde8cbc3612ec635fee29b" [[package]] name = "libm" -version = "0.2.7" +version = "0.2.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7012b1bbb0719e1097c47611d3898568c546d597c2e74d66f6087edd5233ff4" +checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" [[package]] name = "linux-raw-sys" -version = "0.4.7" +version = "0.4.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a9bad9f94746442c783ca431b22403b519cd7fbeed0533fdd6328b2f2212128" +checksum = "da2479e8c062e40bf0066ffa0bc823de0a9368974af99c9f6df941d2c231e03f" [[package]] name = "log" @@ -968,9 +973,9 @@ checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f" [[package]] name = "memchr" -version = "2.6.3" +version = "2.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f232d6ef707e1956a43342693d2a31e72989554d58299d7a88738cc95b0d35c" +checksum = "f665ee40bc4a3c5590afb1e9677db74a508659dfd71e126420da8274909a0167" [[package]] name = "mime" @@ -1022,9 +1027,9 @@ dependencies = [ [[package]] name = "num-traits" -version = "0.2.16" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f30b0abd723be7e2ffca1272140fac1a2f084c77ec3e123c192b66af1ee9e6c2" +checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c" dependencies = [ "autocfg", "libm", @@ -1058,7 +1063,7 @@ dependencies = [ "proc-macro-crate", "proc-macro2", "quote", - "syn 2.0.37", + "syn 2.0.38", ] [[package]] @@ -1152,13 +1157,13 @@ checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" [[package]] name = "pretty" -version = "0.12.1" +version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "563c9d701c3a31dfffaaf9ce23507ba09cbe0b9125ba176d15e629b0235e9acc" +checksum = "b55c4d17d994b637e2f4daf6e5dc5d660d209d5642377d675d7a1c3ab69fa579" dependencies = [ "arrayvec", "typed-arena", - "unicode-segmentation", + "unicode-width", ] [[package]] @@ -1173,22 +1178,22 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.67" +version = "1.0.69" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d433d9f1a3e8c1263d9456598b16fec66f4acc9a74dacffd35c7bb09b3a1328" +checksum = "134c189feb4956b20f6f547d2cf727d4c0fe06722b20a0eec87ed445a97f92da" dependencies = [ "unicode-ident", ] [[package]] name = "proptest" -version = "1.2.0" +version = "1.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4e35c06b98bf36aba164cc17cb25f7e232f5c4aeea73baa14b8a9f0d92dbfa65" +checksum = "7c003ac8c77cb07bb74f5f198bce836a689bcd5a42574612bf14d17bfd08c20e" dependencies = [ "bit-set", - "bitflags 1.3.2", - "byteorder", + "bit-vec", + "bitflags 2.4.0", "lazy_static", "num-traits", "rand", @@ -1274,15 +1279,15 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.6.29" +version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" +checksum = "dbb5fb1acd8a1a18b3dd5be62d25485eb770e05afb408a9627d14d451bae12da" [[package]] name = "reqwest" -version = "0.11.20" +version = "0.11.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e9ad3fe7488d7e34558a2033d45a0c90b72d97b4f80705666fea71472e2e6a1" +checksum = "046cd98826c46c2ac8ddecae268eb5c2e58628688a5fc7a2643704a73faba95b" dependencies = [ "base64", "bytes", @@ -1306,6 +1311,7 @@ dependencies = [ "serde", "serde_json", "serde_urlencoded", + "system-configuration", "tokio", "tokio-rustls", "tokio-util", @@ -1338,12 +1344,26 @@ dependencies = [ "cc", "libc", "once_cell", - "spin", - "untrusted", + "spin 0.5.2", + "untrusted 0.7.1", "web-sys", "winapi", ] +[[package]] +name = "ring" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "911b295d2d302948838c8ac142da1ee09fa7863163b44e6715bc9357905878b8" +dependencies = [ + "cc", + "getrandom", + "libc", + "spin 0.9.8", + "untrusted 0.9.0", + "windows-sys", +] + [[package]] name = "rustc-demangle" version = "0.1.23" @@ -1352,9 +1372,9 @@ checksum = "d626bb9dae77e28219937af045c257c28bfd3f69333c512553507f5f9798cb76" [[package]] name = "rustix" -version = "0.38.14" +version = "0.38.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "747c788e9ce8e92b12cd485c49ddf90723550b654b32508f979b71a7b1ecda4f" +checksum = "f25469e9ae0f3d0047ca8b93fc56843f38e6774f0914a107ff8b41be8be8e0b7" dependencies = [ "bitflags 2.4.0", "errno", @@ -1370,7 +1390,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b80e3dec595989ea8510028f30c408a4630db12c9cbb8de34203b89d6577e99" dependencies = [ "log", - "ring", + "ring 0.16.20", "sct", "webpki", ] @@ -1382,7 +1402,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cd8d6c9f025a446bc4d18ad9632e69aec8f287aa84499ee335599fabd20c3fd8" dependencies = [ "log", - "ring", + "ring 0.16.20", "rustls-webpki", "sct", ] @@ -1402,8 +1422,8 @@ version = "0.101.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c7d5dece342910d9ba34d259310cae3e0154b873b35408b787b59bce53d34fe" dependencies = [ - "ring", - "untrusted", + "ring 0.16.20", + "untrusted 0.7.1", ] [[package]] @@ -1436,8 +1456,8 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d53dcdb7c9f8158937a7981b48accfd39a43af418591a5d008c7b22b5e1b7ca4" dependencies = [ - "ring", - "untrusted", + "ring 0.16.20", + "untrusted 0.7.1", ] [[package]] @@ -1490,7 +1510,7 @@ checksum = "4eca7ac642d82aa35b60049a6eccb4be6be75e599bd2e9adb5f875a737654af2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.37", + "syn 2.0.38", ] [[package]] @@ -1512,7 +1532,7 @@ checksum = "8725e1dfadb3a50f7e5ce0b1a540466f6ed3fe7a0fca2ac2b8b831d31316bd00" dependencies = [ "proc-macro2", "quote", - "syn 2.0.37", + "syn 2.0.38", ] [[package]] @@ -1553,9 +1573,9 @@ dependencies = [ [[package]] name = "sha2" -version = "0.10.7" +version = "0.10.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "479fb9d862239e610720565ca91403019f2f00410f1864c5aa7479b950a76ed8" +checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8" dependencies = [ "cfg-if", "cpufeatures", @@ -1628,6 +1648,12 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d" +[[package]] +name = "spin" +version = "0.9.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" + [[package]] name = "spki" version = "0.7.2" @@ -1670,15 +1696,36 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.37" +version = "2.0.38" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7303ef2c05cd654186cb250d29049a24840ca25d2747c25c0381c8d9e2f582e8" +checksum = "e96b79aaa137db8f61e26363a0c9b47d8b4ec75da28b7d1d614c2303e232408b" dependencies = [ "proc-macro2", "quote", "unicode-ident", ] +[[package]] +name = "system-configuration" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba3a3adc5c275d719af8cb4272ea1c4a6d668a777f37e115f6d11ddbc1c8e0e7" +dependencies = [ + "bitflags 1.3.2", + "core-foundation", + "system-configuration-sys", +] + +[[package]] +name = "system-configuration-sys" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a75fb188eb626b924683e3b95e3a48e63551fcfb51949de2f06a9d91dbee93c9" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "tempfile" version = "3.8.0" @@ -1715,29 +1762,29 @@ dependencies = [ [[package]] name = "thiserror" -version = "1.0.48" +version = "1.0.49" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d6d7a740b8a666a7e828dd00da9c0dc290dff53154ea77ac109281de90589b7" +checksum = "1177e8c6d7ede7afde3585fd2513e611227efd6481bd78d2e82ba1ce16557ed4" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.48" +version = "1.0.49" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49922ecae66cc8a249b77e68d1d0623c1b2c514f0060c27cdc68bd62a1219d35" +checksum = "10712f02019e9288794769fba95cd6847df9874d49d871d062172f9dd41bc4cc" dependencies = [ "proc-macro2", "quote", - "syn 2.0.37", + "syn 2.0.38", ] [[package]] name = "time" -version = "0.3.28" +version = "0.3.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17f6bb557fd245c28e6411aa56b6403c689ad95061f50e4be16c274e70a17e48" +checksum = "426f806f4089c493dcac0d24c29c01e2c38baf8e30f1b716ee37e83d200b18fe" dependencies = [ "deranged", "itoa", @@ -1748,15 +1795,15 @@ dependencies = [ [[package]] name = "time-core" -version = "0.1.1" +version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7300fbefb4dadc1af235a9cef3737cea692a9d97e1b9cbcd4ebdae6f8868e6fb" +checksum = "ef927ca75afb808a4d64dd374f00a2adf8d0fcff8e7b184af886c3c87ec4a3f3" [[package]] name = "time-macros" -version = "0.2.14" +version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a942f44339478ef67935ab2bbaec2fb0322496cf3cbe84b261e06ac3814c572" +checksum = "4ad70d68dba9e1f8aceda7aa6711965dfec1cac869f311a51bd08b3a2ccbce20" dependencies = [ "time-core", ] @@ -1778,9 +1825,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.32.0" +version = "1.33.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17ed6077ed6cd6c74735e21f37eb16dc3935f96878b1fe961074089cc80893f9" +checksum = "4f38200e3ef7995e5ef13baec2f432a6da0aa9ac495b2c0e8f3b7eec2c92d653" dependencies = [ "backtrace", "bytes", @@ -1828,7 +1875,7 @@ version = "0.19.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b5bb770da30e5cbfde35a2d7b9b8a2c4b8ef89548a7a6aeab5c9a576e3e7421" dependencies = [ - "indexmap 2.0.0", + "indexmap 2.0.2", "toml_datetime", "winnow", ] @@ -1904,12 +1951,6 @@ dependencies = [ "tinyvec", ] -[[package]] -name = "unicode-segmentation" -version = "1.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1dd624098567895118886609431a7c3b8f516e41d30e0643f03d94592a147e36" - [[package]] name = "unicode-width" version = "0.1.11" @@ -1922,6 +1963,12 @@ version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a" +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + [[package]] name = "url" version = "2.4.1" @@ -1984,7 +2031,7 @@ dependencies = [ "once_cell", "proc-macro2", "quote", - "syn 2.0.37", + "syn 2.0.38", "wasm-bindgen-shared", ] @@ -2018,7 +2065,7 @@ checksum = "54681b18a46765f095758388f2d0cf16eb8d4169b639ab575a8f5693af210c7b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.37", + "syn 2.0.38", "wasm-bindgen-backend", "wasm-bindgen-shared", ] @@ -2054,12 +2101,12 @@ dependencies = [ [[package]] name = "webpki" -version = "0.22.1" +version = "0.22.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0e74f82d49d545ad128049b7e88f6576df2da6b02e9ce565c6f533be576957e" +checksum = "ed63aea5ce73d0ff405984102c42de94fc55a6b75765d621c65262469b3c9b53" dependencies = [ - "ring", - "untrusted", + "ring 0.17.2", + "untrusted 0.9.0", ] [[package]] @@ -2167,9 +2214,9 @@ checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" [[package]] name = "winnow" -version = "0.5.15" +version = "0.5.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c2e3184b9c4e92ad5167ca73039d0c42476302ab603e2fec4487511f38ccefc" +checksum = "037711d82167854aff2018dfd193aa0fef5370f456732f0d5a0c59b0f1b4b907" dependencies = [ "memchr", ] diff --git a/src/ic-websocket-cdk/src/lib.rs b/src/ic-websocket-cdk/src/lib.rs index 2816771..648ff2f 100644 --- a/src/ic-websocket-cdk/src/lib.rs +++ b/src/ic-websocket-cdk/src/lib.rs @@ -374,7 +374,6 @@ fn get_messages_for_gateway_range(gateway_principal: Principal, nonce: u64) -> ( MESSAGES_FOR_GATEWAY.with(|m| { let queue_len = m.borrow().len(); - // TODO: test if nonce == 0 && queue_len > 0 { // this is the case in which the poller on the gateway restarted // the range to return is end:last index and start: max(end - MAX_NUMBER_OF_RETURNED_MESSAGES, 0) @@ -595,7 +594,6 @@ pub struct WsHandlers { impl WsHandlers { fn call_on_open(&self, args: OnOpenCallbackArgs) { if let Some(on_open) = self.on_open { - // TODO: test the panic handling let res = panic::catch_unwind(|| { on_open(args); }); @@ -608,7 +606,6 @@ impl WsHandlers { fn call_on_message(&self, args: OnMessageCallbackArgs) { if let Some(on_message) = self.on_message { - // TODO: test the panic handling let res = panic::catch_unwind(|| { on_message(args); }); @@ -621,7 +618,6 @@ impl WsHandlers { fn call_on_close(&self, args: OnCloseCallbackArgs) { if let Some(on_close) = self.on_close { - // TODO: test the panic handling let res = panic::catch_unwind(|| { on_close(args); }); @@ -990,6 +986,45 @@ mod test { assert!(CUSTOM_STATE.with(|h| h.borrow().is_on_close_called)); } + #[test] + fn test_ws_handlers_panic_is_handled() { + let handlers = WsHandlers { + on_open: Some(|_| { + panic!("on_open_panic"); + }), + on_message: Some(|_| { + panic!("on_close_panic"); + }), + on_close: Some(|_| { + panic!("on_close_panic"); + }), + }; + + initialize_handlers(handlers); + + let handlers = HANDLERS.with(|h| h.borrow().clone()); + + let res = panic::catch_unwind(|| { + handlers.call_on_open(OnOpenCallbackArgs { + client_principal: test_utils::generate_random_principal(), + }); + }); + assert!(res.is_ok()); + let res = panic::catch_unwind(|| { + handlers.call_on_message(OnMessageCallbackArgs { + client_principal: test_utils::generate_random_principal(), + message: vec![], + }); + }); + assert!(res.is_ok()); + let res = panic::catch_unwind(|| { + handlers.call_on_close(OnCloseCallbackArgs { + client_principal: test_utils::generate_random_principal(), + }); + }); + assert!(res.is_ok()); + } + #[test] fn test_current_time() { // test From 924e3544d4fde8d74f2a67d7f5f735650abffe39 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Mon, 9 Oct 2023 14:04:22 +0200 Subject: [PATCH 3/4] chore: handle received service message (mock) --- src/ic-websocket-cdk/src/lib.rs | 31 +++++++++++++++++++++++++++---- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/src/ic-websocket-cdk/src/lib.rs b/src/ic-websocket-cdk/src/lib.rs index 648ff2f..ffcfad7 100644 --- a/src/ic-websocket-cdk/src/lib.rs +++ b/src/ic-websocket-cdk/src/lib.rs @@ -1,4 +1,4 @@ -use candid::{encode_one, CandidType, Principal}; +use candid::{decode_one, encode_one, CandidType, Principal}; #[cfg(not(test))] use ic_cdk::api::time; use ic_cdk::api::{caller, data_certificate, set_certified_data}; @@ -489,7 +489,7 @@ struct ClientKeepAliveMessageContent { last_incoming_sequence_num: u64, } -/// A service message sent by the CDK to the client. +/// A service message sent by the CDK to the client or vice versa. #[derive(CandidType, Deserialize)] enum WebsocketServiceMessageContent { /// Message sent by the **canister** when a client opens a connection. @@ -500,6 +500,16 @@ enum WebsocketServiceMessageContent { KeepAliveMessage(ClientKeepAliveMessageContent), } +impl WebsocketServiceMessageContent { + fn from_candid_bytes(bytes: Vec) -> Result { + decode_one(&bytes).map_err(|e| { + let mut err = String::from("Error decoding service message content: "); + err.push_str(&e.to_string()); + err + }) + } +} + fn send_service_message_to_client( client_key: &ClientKey, message: WebsocketServiceMessageContent, @@ -557,6 +567,20 @@ fn _ws_send( Ok(()) } +fn handle_received_service_message(content: Vec) -> CanisterWsMessageResult { + let decoded = WebsocketServiceMessageContent::from_candid_bytes(content)?; + match decoded { + WebsocketServiceMessageContent::OpenMessage(_) + | WebsocketServiceMessageContent::AckMessage(_) => { + Err(String::from("Invalid received service message")) + }, + WebsocketServiceMessageContent::KeepAliveMessage(_) => { + custom_print!("Service message handling not implemented yet"); + Ok(()) + }, + } +} + /// Arguments passed to the `on_open` handler. pub struct OnOpenCallbackArgs { pub client_principal: ClientPrincipal, @@ -769,9 +793,8 @@ pub fn ws_message(args: CanisterWsMessageArguments) -> CanisterWsMessageResult { // increase the expected sequence number by 1 increment_expected_incoming_message_from_client_num(&client_key)?; - // TODO: test if is_service_message { - custom_print!("Service message handling not implemented yet"); + return handle_received_service_message(content); } // call the on_message handler initialized in init() From 5d6c31f8bb035d4bcfb90a937ed4a9e1a76f75ca Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Mon, 9 Oct 2023 14:05:09 +0200 Subject: [PATCH 4/4] chore: integration tests --- tests/integration/canister.test.ts | 564 ++++++++++------------------ tests/integration/utils/api.ts | 65 +--- tests/integration/utils/client.ts | 5 + tests/integration/utils/messages.ts | 78 +++- tests/src/lib.rs | 10 +- tests/test_canister.did | 2 +- 6 files changed, 302 insertions(+), 422 deletions(-) create mode 100644 tests/integration/utils/client.ts diff --git a/tests/integration/canister.test.ts b/tests/integration/canister.test.ts index ecf5eaa..f1f76a7 100644 --- a/tests/integration/canister.test.ts +++ b/tests/integration/canister.test.ts @@ -1,6 +1,4 @@ import { IDL } from "@dfinity/candid"; -import { Principal } from "@dfinity/principal"; -import { Cbor } from "@dfinity/agent"; import { anonymousClient, canisterId, @@ -13,8 +11,6 @@ import { gateway2, } from "./utils/actors"; import { - isMessageBodyValid, - isValidCertificate, reinitialize, wsClose, wsGetMessages, @@ -34,14 +30,37 @@ import type { WebsocketMessage, } from "../src/declarations/test_canister/test_canister.did"; import { generateClientKey, getRandomClientNonce } from "./utils/random"; -import { CanisterOpenMessageContent, WebsocketServiceMessageContent, encodeWebsocketServiceMessageContent, getServiceMessageFromCanisterMessage, isClientKeyEq } from "./utils/idl"; -import { createWebsocketMessage, decodeWebsocketMessage, filterServiceMessagesFromCanisterMessages } from "./utils/messages"; - -const MAX_NUMBER_OF_RETURNED_MESSAGES = 10; // set in the CDK +import { + CanisterOpenMessageContent, + WebsocketServiceMessageContent, + encodeWebsocketServiceMessageContent, + getServiceMessageFromCanisterMessage, + isClientKeyEq, +} from "./utils/idl"; +import { + isMessageBodyValid, + isValidCertificate, + createWebsocketMessage, + decodeWebsocketMessage, + filterServiceMessagesFromCanisterMessages, + getNextPollingNonceFromMessages, +} from "./utils/messages"; +import { formatClientKey } from "./utils/client"; + +/** + * The maximum number of messages returned by the **ws_get_messages** method. Set in the CDK. + * + * Value: `10` + */ +const MAX_NUMBER_OF_RETURNED_MESSAGES = 10; +/** + * @{@link MAX_NUMBER_OF_RETURNED_MESSAGES} + 2 + * + * Value: `12` + */ const SEND_MESSAGES_COUNT = MAX_NUMBER_OF_RETURNED_MESSAGES + 2; // test with more messages to check the indexes and limits -const MAX_GATEWAY_KEEP_ALIVE_TIME_MS = 15_000; // set in the CDK -const DEFAULT_TEST_SEND_ACK_INTERVAL_MS = 300_000; // 5 minutes to make sure the canister doesn't reset the client -const DEFAULT_TEST_KEEP_ALIVE_DELAY_MS = 300_000; // 5 minutes to make sure the canister doesn't reset the client +// const DEFAULT_TEST_SEND_ACK_INTERVAL_MS = 300_000; // 5 minutes to make sure the canister doesn't reset the client +// const DEFAULT_TEST_KEEP_ALIVE_DELAY_MS = 300_000; // 5 minutes to make sure the canister doesn't reset the client let client1Key: ClientKey; let client2Key: ClientKey; @@ -252,14 +271,14 @@ describe("Canister - ws_message", () => { clientActor: client1, }, true); - // wring content encoding + // wrong content encoding const res = await wsMessage({ message: createWebsocketMessage(client1Key, 1, true, new Uint8Array([1, 2, 3])), actor: client1, }); expect(res).toMatchObject({ - Err: expect.stringContaining("Error decoding service message from client:"), + Err: expect.stringContaining("Error decoding service message content:"), }); const wrongServiceMessage: WebsocketServiceMessageContent = { @@ -274,7 +293,7 @@ describe("Canister - ws_message", () => { }); expect(res2).toMatchObject({ - Err: "invalid keep alive message content", + Err: "Invalid received service message", }); }); @@ -333,233 +352,15 @@ describe("Canister - ws_get_messages (failures,empty)", () => { }); }); -// describe("Canister - ws_message (gateway status)", () => { -// beforeAll(async () => { -// await assignKeysToClients(); - -// await wsRegister({ -// clientActor: client1, -// clientKey: client1Key.publicKey, -// }, true); - -// await wsOpen({ -// clientPublicKey: client1Key.publicKey, -// clientSecretKey: client1Key.secretKey, -// canisterId, -// clientActor: gateway1, -// }, true); - -// await wsSend({ -// clientKey: client1Key.publicKey, -// actor: client1, -// message: { text: "test" }, -// }, true); -// }); - -// afterAll(async () => { -// await wsWipe(gateway1); -// }); - -// it("fails if a non registered gateway sends an IcWebSocketGatewayStatus message", async () => { -// const res = await wsMessage({ -// message: { -// IcWebSocketGatewayStatus: { -// status_index: BigInt(1), -// }, -// }, -// actor: gateway2, -// }); - -// expect(res).toMatchObject({ -// Err: "caller is not the gateway that has been registered during CDK initialization", -// }); -// }); - -// it("registered gateway should update the status index", async () => { -// const res = await wsMessage({ -// message: { -// IcWebSocketGatewayStatus: { -// status_index: BigInt(2), // set it high to test behavior for indexes behind the current one -// }, -// }, -// actor: gateway1, -// }); - -// expect(res).toMatchObject({ -// Ok: null, -// }); -// }); - -// it("fails if a registered gateway sends an IcWebSocketGatewayStatus with a wrong status index (equal to current)", async () => { -// const res = await wsMessage({ -// message: { -// IcWebSocketGatewayStatus: { -// status_index: BigInt(2), -// }, -// }, -// actor: gateway1, -// }); - -// expect(res).toMatchObject({ -// Err: "Gateway status index is equal to or behind the current one", -// }); -// }); - -// it("fails if a registered gateway sends an IcWebSocketGatewayStatus with a wrong status index (behind the current)", async () => { -// const res = await wsMessage({ -// message: { -// IcWebSocketGatewayStatus: { -// status_index: BigInt(1), -// }, -// }, -// actor: gateway1, -// }); - -// expect(res).toMatchObject({ -// Err: "Gateway status index is equal to or behind the current one", -// }); -// }); - -// it("registered gateway should disconnect after maximum time", async () => { -// let res = await gateway1.ws_get_messages({ -// nonce: BigInt(0), -// }); - -// expect(res).toMatchObject({ -// Ok: { -// messages: expect.any(Array), -// cert: expect.any(Uint8Array), -// tree: expect.any(Uint8Array), -// }, -// }); -// expect((res as { Ok: CanisterOutputCertifiedMessages }).Ok.messages.length).toEqual(1); - -// // wait for the maximum time the gateway can send a status message, -// // so that the internal canister state is reset -// // double the time to make sure the canister state is reset -// await new Promise((resolve) => setTimeout(resolve, 2 * MAX_GATEWAY_KEEP_ALIVE_TIME_MS)); - -// // check if messages have been deleted -// res = await gateway1.ws_get_messages({ -// nonce: BigInt(0), -// }); -// expect(res).toMatchObject({ -// Ok: { -// messages: [], -// cert: expect.any(Uint8Array), -// tree: expect.any(Uint8Array), -// }, -// }); - -// // check if registered client has been deleted -// const sendRes = await wsSend({ -// clientKey: client1Key.publicKey, -// actor: client1, -// message: { text: "test" }, -// }); -// expect(sendRes).toMatchObject({ -// Err: "client's public key has not been previously registered by client", -// }); -// }); - -// it("registered gateway should reconnect by resetting the status index", async () => { -// let res = await wsMessage({ -// message: { -// IcWebSocketGatewayStatus: { -// status_index: BigInt(0), -// }, -// }, -// actor: gateway1, -// }); - -// expect(res).toMatchObject({ -// Ok: null, -// }); - -// res = await wsMessage({ -// message: { -// IcWebSocketGatewayStatus: { -// status_index: BigInt(1), -// }, -// }, -// actor: gateway1, -// }); - -// expect(res).toMatchObject({ -// Ok: null, -// }); -// }); - -// it("registered gateway should reconnect before maximum time", async () => { -// // reconnect the client -// await wsRegister({ -// clientActor: client1, -// clientKey: client1Key.publicKey, -// }, true); - -// await wsOpen({ -// clientPublicKey: client1Key.publicKey, -// clientSecretKey: client1Key.secretKey, -// canisterId, -// clientActor: gateway1, -// }, true); - -// // send a test message from the canister to check if the internal state is reset -// await wsSend({ -// clientKey: client1Key.publicKey, -// actor: client1, -// message: { text: "test" }, -// }, true); - -// // check if the canister has the message in the queue -// let messagesRes = await gateway1.ws_get_messages({ -// nonce: BigInt(0), -// }); -// expect(messagesRes).toMatchObject({ -// Ok: { -// messages: expect.any(Array), -// cert: expect.any(Uint8Array), -// tree: expect.any(Uint8Array), -// }, -// }); -// expect((messagesRes as { Ok: CanisterOutputCertifiedMessages }).Ok.messages.length).toEqual(1); - -// // simulate a reconnection -// const res = await wsMessage({ -// message: { -// IcWebSocketGatewayStatus: { -// status_index: BigInt(0), -// }, -// }, -// actor: gateway1, -// }); -// expect(res).toMatchObject({ -// Ok: null, -// }); - -// // check if the canister reset the internal state -// messagesRes = await gateway1.ws_get_messages({ -// nonce: BigInt(0), -// }); -// expect(messagesRes).toMatchObject({ -// Ok: { -// messages: [], -// cert: expect.any(Uint8Array), -// tree: expect.any(Uint8Array), -// }, -// }); -// }); -// }); - -describe.only("Canister - ws_get_messages (receive)", () => { +describe("Canister - ws_get_messages (receive)", () => { beforeAll(async () => { await assignKeysToClients(); // reset the internal timers - await reinitialize({ - sendAckIntervalMs: DEFAULT_TEST_SEND_ACK_INTERVAL_MS, - keepAliveDelayMs: DEFAULT_TEST_KEEP_ALIVE_DELAY_MS, - }); + // await reinitialize({ + // sendAckIntervalMs: DEFAULT_TEST_SEND_ACK_INTERVAL_MS, + // keepAliveDelayMs: DEFAULT_TEST_KEEP_ALIVE_DELAY_MS, + // }); await wsOpen({ clientNonce: client1Key.client_nonce, @@ -568,14 +369,15 @@ describe.only("Canister - ws_get_messages (receive)", () => { }, true); // prepare the messages - for (let i = 0; i < SEND_MESSAGES_COUNT; i++) { - const appMessage = { text: `test${i}` }; - await wsSend({ - clientPrincipal: client1Key.client_principal, - actor: client1, - message: appMessage, - }, true); - } + const messages = Array.from({ length: SEND_MESSAGES_COUNT }, (_, i) => { + return { text: `test${i}` }; + }); + + await wsSend({ + clientPrincipal: client1Key.client_principal, + actor: client1, + messages, + }, true); await commonAgent.fetchRootKey(); }); @@ -586,7 +388,7 @@ describe.only("Canister - ws_get_messages (receive)", () => { it("registered gateway can receive correct amount of messages", async () => { // on open, the canister puts a service message in the queue - const messagesCount = SEND_MESSAGES_COUNT + 1; // +1 for the service message + const messagesCount = SEND_MESSAGES_COUNT + 1; // +1 for the open service message for (let i = 0; i < messagesCount; i++) { const res = await gateway1.ws_get_messages({ nonce: BigInt(i), @@ -625,21 +427,22 @@ describe.only("Canister - ws_get_messages (receive)", () => { it("registered gateway can receive certified messages", async () => { // first batch of messages const firstBatchRes = await gateway1.ws_get_messages({ - nonce: BigInt(1), + nonce: BigInt(1), // skip the case in which the gateway restarts polling from the beginning (tested below) }); const firstBatchMessagesResult = (firstBatchRes as { Ok: CanisterOutputCertifiedMessages }).Ok; - console.log(firstBatchMessagesResult.messages.map((msg) => msg.key)); - for (let i = 0; i < firstBatchMessagesResult.messages.length; i++) { - const message = firstBatchMessagesResult.messages[i]; + expect(firstBatchMessagesResult.messages.length).toBe(MAX_NUMBER_OF_RETURNED_MESSAGES); + + let expectedSequenceNumber = 2; // first is the service open message and the number is incremented before sending + let i = 0; + for (const message of firstBatchMessagesResult.messages) { expect(isClientKeyEq(message.client_key, client1Key)).toEqual(true); const websocketMessage = decodeWebsocketMessage(new Uint8Array(message.content)); - console.log(websocketMessage); expect(websocketMessage).toMatchObject({ client_key: expect.any(Object), content: expect.any(Uint8Array), - sequence_num: BigInt(i + 1), - timestamp: expect.any(Object), // weird cbor bigint deserialization + sequence_num: BigInt(expectedSequenceNumber), + timestamp: expect.any(BigInt), is_service_message: false, }); expect(isClientKeyEq(websocketMessage.client_key, client1Key)).toEqual(true); @@ -661,27 +464,33 @@ describe.only("Canister - ws_get_messages (receive)", () => { firstBatchMessagesResult.tree as Uint8Array, ) ).resolves.toBe(true); + + expectedSequenceNumber++; + i++; } + const nextPollingNonce = getNextPollingNonceFromMessages(firstBatchMessagesResult.messages); + // second batch of messages, starting from the last nonce of the first batch const secondBatchRes = await gateway1.ws_get_messages({ - nonce: BigInt(MAX_NUMBER_OF_RETURNED_MESSAGES), + nonce: BigInt(nextPollingNonce), }); const secondBatchMessagesResult = (secondBatchRes as { Ok: CanisterOutputCertifiedMessages }).Ok; - for (let i = 0; i < secondBatchMessagesResult.messages.length; i++) { - const message = secondBatchMessagesResult.messages[i]; + expect(secondBatchMessagesResult.messages.length).toBe(SEND_MESSAGES_COUNT - MAX_NUMBER_OF_RETURNED_MESSAGES); // remaining from SEND_MESSAGES_COUNT + + for (const message of secondBatchMessagesResult.messages) { expect(isClientKeyEq(message.client_key, client1Key)).toEqual(true); const websocketMessage = decodeWebsocketMessage(new Uint8Array(message.content)); expect(websocketMessage).toMatchObject({ client_key: expect.any(Object), content: expect.any(Uint8Array), - sequence_num: BigInt(i + MAX_NUMBER_OF_RETURNED_MESSAGES + 1), - timestamp: expect.any(Object), // weird cbor bigint deserialization + sequence_num: BigInt(expectedSequenceNumber), + timestamp: expect.any(BigInt), is_service_message: false, }); expect(isClientKeyEq(websocketMessage.client_key, client1Key)).toEqual(true); - expect(IDL.decode([IDL.Record({ 'text': IDL.Text })], websocketMessage.content as Uint8Array)).toEqual([{ text: `test${i + MAX_NUMBER_OF_RETURNED_MESSAGES}` }]); + expect(IDL.decode([IDL.Record({ 'text': IDL.Text })], websocketMessage.content as Uint8Array)).toEqual([{ text: `test${i}` }]); // check the certification await expect( @@ -699,107 +508,144 @@ describe.only("Canister - ws_get_messages (receive)", () => { secondBatchMessagesResult.tree as Uint8Array, ) ).resolves.toBe(true); + + expectedSequenceNumber++; + i++; + } + }); + + it("registered gateway can poll messages after restart", async () => { + const batchRes = await gateway1.ws_get_messages({ + nonce: BigInt(0), // start polling from the beginning, as if the gateway restarted + }); + + // we expect that the messages returned are the last MAX_NUMBER_OF_RETURNED_MESSAGES + const messagesResult = (batchRes as { Ok: CanisterOutputCertifiedMessages }).Ok; + expect(messagesResult.messages.length).toBe(MAX_NUMBER_OF_RETURNED_MESSAGES); + + let expectedSequenceNumber = SEND_MESSAGES_COUNT - MAX_NUMBER_OF_RETURNED_MESSAGES + 1 + 1; // +1 for the service open message +1 because the seq num is incremented before sending + let i = SEND_MESSAGES_COUNT - MAX_NUMBER_OF_RETURNED_MESSAGES; + for (const message of messagesResult.messages) { + expect(isClientKeyEq(message.client_key, client1Key)).toEqual(true); + const websocketMessage = decodeWebsocketMessage(new Uint8Array(message.content)); + expect(websocketMessage).toMatchObject({ + client_key: expect.any(Object), + content: expect.any(Uint8Array), + sequence_num: BigInt(expectedSequenceNumber), + timestamp: expect.any(BigInt), + is_service_message: false, + }); + expect(isClientKeyEq(websocketMessage.client_key, client1Key)).toEqual(true); + expect(IDL.decode([IDL.Record({ 'text': IDL.Text })], websocketMessage.content as Uint8Array)).toEqual([{ text: `test${i}` }]); + + // check the certification + await expect( + isValidCertificate( + canisterId, + messagesResult.cert as Uint8Array, + messagesResult.tree as Uint8Array, + commonAgent + ) + ).resolves.toBe(true); + await expect( + isMessageBodyValid( + message.key, + message.content as Uint8Array, + messagesResult.tree as Uint8Array, + ) + ).resolves.toBe(true); + + expectedSequenceNumber++; + i++; } }); }); -// describe("Canister - ws_close", () => { -// beforeAll(async () => { -// await assignKeysToClients(); - -// await wsRegister({ -// clientActor: client1, -// clientKey: client1Key.publicKey, -// }, true); - -// await wsOpen({ -// clientPublicKey: client1Key.publicKey, -// clientSecretKey: client1Key.secretKey, -// canisterId, -// clientActor: gateway1, -// }, true); -// }); - -// afterAll(async () => { -// await wsWipe(gateway1); -// }); - -// it("fails if gateway is not registered", async () => { -// const res = await wsClose({ -// clientPublicKey: client1Key.publicKey, -// gatewayActor: gateway2, -// }); - -// expect(res).toMatchObject({ -// Err: "caller is not the gateway that has been registered during CDK initialization", -// }); -// }); - -// it("fails if client is not registered", async () => { -// const res = await wsClose({ -// clientPublicKey: client2Key.publicKey, -// gatewayActor: gateway1, -// }); - -// expect(res).toMatchObject({ -// Err: "client's public key has not been previously registered by client", -// }); -// }); - -// it("should close the websocket for a registered client", async () => { -// const res = await wsClose({ -// clientPublicKey: client1Key.publicKey, -// gatewayActor: gateway1, -// }); - -// expect(res).toMatchObject({ -// Ok: null, -// }); -// }); -// }); - -// describe("Canister - ws_send", () => { -// beforeAll(async () => { -// await assignKeysToClients(); - -// await wsRegister({ -// clientActor: client1, -// clientKey: client1Key.publicKey, -// }, true); - -// await wsOpen({ -// clientPublicKey: client1Key.publicKey, -// clientSecretKey: client1Key.secretKey, -// canisterId, -// clientActor: gateway1, -// }, true); -// }); - -// afterAll(async () => { -// await wsWipe(gateway1); -// }); - -// it("fails if sending a message to a non registered client", async () => { -// const res = await wsSend({ -// clientKey: client2Key.publicKey, -// actor: client1, -// message: { text: "test" }, -// }); - -// expect(res).toMatchObject({ -// Err: "client's public key has not been previously registered by client", -// }); -// }); - -// it("should send a message to a registered client", async () => { -// const res = await wsSend({ -// clientKey: client1Key.publicKey, -// actor: client1, -// message: { text: "test" }, -// }); - -// expect(res).toMatchObject({ -// Ok: null, -// }); -// }); -// }); +describe("Canister - ws_close", () => { + beforeAll(async () => { + await assignKeysToClients(); + + await wsOpen({ + clientNonce: client1Key.client_nonce, + canisterId, + clientActor: client1, + }, true); + }); + + afterAll(async () => { + await wsWipe(); + }); + + it("fails if gateway is not registered", async () => { + const res = await wsClose({ + clientKey: client1Key, + gatewayActor: gateway2, + }); + + expect(res).toMatchObject({ + Err: "caller is not the gateway that has been registered during CDK initialization", + }); + }); + + it("fails if client is not registered", async () => { + const res = await wsClose({ + clientKey: client2Key, + gatewayActor: gateway1, + }); + + expect(res).toMatchObject({ + Err: `client with key ${formatClientKey(client2Key)} doesn't have an open connection`, + }); + }); + + it("should close the websocket for a registered client", async () => { + const res = await wsClose({ + clientKey: client1Key, + gatewayActor: gateway1, + }); + + expect(res).toMatchObject({ + Ok: null, + }); + }); +}); + +describe("Canister - ws_send", () => { + beforeAll(async () => { + await assignKeysToClients(); + + await wsOpen({ + clientNonce: client1Key.client_nonce, + canisterId, + clientActor: client1, + }, true); + }); + + afterAll(async () => { + await wsWipe(); + }); + + it("fails if sending a message to a non registered client", async () => { + const res = await wsSend({ + clientPrincipal: client2Key.client_principal, + actor: client1, + messages: [{ text: "test" }], + }); + + expect(res).toMatchObject({ + Err: `client with principal ${client2Key.client_principal.toText()} doesn't have an open connection`, + }); + }); + + it("should send a message to a registered client", async () => { + const res = await wsSend({ + clientPrincipal: client1Key.client_principal, + actor: client1, + messages: [{ text: "test" }], + }); + + expect(res).toMatchObject({ + Ok: null, + }); + }); +}); diff --git a/tests/integration/utils/api.ts b/tests/integration/utils/api.ts index a380a87..731ed54 100644 --- a/tests/integration/utils/api.ts +++ b/tests/integration/utils/api.ts @@ -1,8 +1,6 @@ // helpers for functions that are called frequently in tests -import { ActorSubclass, Cbor, Certificate, HashTree, HttpAgent, compare, lookup_path, reconstruct } from "@dfinity/agent"; -import { Secp256k1KeyIdentity } from "@dfinity/identity-secp256k1"; -import { Principal } from "@dfinity/principal"; +import { ActorSubclass } from "@dfinity/agent"; import { IDL } from "@dfinity/candid"; import { anonymousClient, gateway1Data } from "./actors"; import type { CanisterOutputCertifiedMessages, ClientKey, ClientPrincipal, WebsocketMessage, _SERVICE } from "../../src/declarations/test_canister/test_canister.did"; @@ -119,65 +117,16 @@ export const reinitialize = async (args: ReinitializeArgs) => { type WsSendArgs = { clientPrincipal: ClientPrincipal, actor: ActorSubclass<_SERVICE>, - message: { + messages: Array<{ text: string, - }, + }>, }; export const wsSend = async (args: WsSendArgs, throwIfError = false) => { - const msgBytes = IDL.encode([IDL.Record({ 'text': IDL.Text })], [args.message]); - const res = await args.actor.ws_send(args.clientPrincipal, new Uint8Array(msgBytes)); + const serializedMessages = args.messages.map((msg) => { + return new Uint8Array(IDL.encode([IDL.Record({ 'text': IDL.Text })], [msg])); + }); + const res = await args.actor.ws_send(args.clientPrincipal, serializedMessages); return resolveResult(res, throwIfError); }; - -export const getCertifiedMessageKey = async (gatewayIdentity: Promise, nonce: number) => { - const gatewayPrincipal = (await gatewayIdentity).getPrincipal().toText(); - return `${gatewayPrincipal}_${String(nonce).padStart(20, '0')}`; -}; - -export const isValidCertificate = async (canisterId: string, certificate: Uint8Array, tree: Uint8Array, agent: HttpAgent) => { - const canisterPrincipal = Principal.fromText(canisterId); - let cert: Certificate; - - try { - cert = await Certificate.create({ - certificate, - canisterId: canisterPrincipal, - rootKey: agent.rootKey! - }); - } catch (error) { - console.error("Error creating certificate:", error); - return false; - } - - const hashTree = Cbor.decode(tree); - const reconstructed = await reconstruct(hashTree); - const witness = cert.lookup([ - "canister", - canisterPrincipal.toUint8Array(), - "certified_data" - ]); - - if (!witness) { - throw new Error( - "Could not find certified data for this canister in the certificate." - ); - } - - // First validate that the Tree is as good as the certification. - return compare(witness, reconstructed) === 0; -}; - -export const isMessageBodyValid = async (path: string, body: Uint8Array | ArrayBuffer, tree: Uint8Array) => { - const hashTree = Cbor.decode(tree); - const sha = await crypto.subtle.digest("SHA-256", body); - let treeSha = lookup_path(["websocket", path], hashTree); - - if (!treeSha) { - // Allow fallback to index path. - treeSha = lookup_path(["websocket"], hashTree); - } - - return !!treeSha && (compare(sha, treeSha) === 0); -}; diff --git a/tests/integration/utils/client.ts b/tests/integration/utils/client.ts new file mode 100644 index 0000000..2a66f27 --- /dev/null +++ b/tests/integration/utils/client.ts @@ -0,0 +1,5 @@ +import { ClientKey } from "../../src/declarations/test_canister/test_canister.did"; + +export const formatClientKey = (clientKey: ClientKey): string => { + return `${clientKey.client_principal.toText()}_${clientKey.client_nonce.toString()}`; +}; diff --git a/tests/integration/utils/messages.ts b/tests/integration/utils/messages.ts index 993e36c..da2ac8e 100644 --- a/tests/integration/utils/messages.ts +++ b/tests/integration/utils/messages.ts @@ -1,6 +1,8 @@ -import { Cbor } from "@dfinity/agent"; +import { Cbor, Certificate, HashTree, HttpAgent, compare, lookup_path, reconstruct } from "@dfinity/agent"; import { CanisterOutputMessage, ClientKey, WebsocketMessage } from "../../src/declarations/test_canister/test_canister.did"; import { getWebsocketMessageFromCanisterMessage } from "./idl"; +import { Principal } from "@dfinity/principal"; +import { Secp256k1KeyIdentity } from "@dfinity/identity-secp256k1"; export const filterServiceMessagesFromCanisterMessages = (messages: CanisterOutputMessage[]): CanisterOutputMessage[] => { return messages.filter((msg) => { @@ -27,5 +29,77 @@ export const createWebsocketMessage = ( }; export const decodeWebsocketMessage = (bytes: Uint8Array): WebsocketMessage => { - return Cbor.decode(bytes); + const decoded: any = Cbor.decode(bytes); + + // normalize the decoded message + return { + client_key: { + client_principal: Principal.fromUint8Array(decoded.client_key.client_principal), + client_nonce: BigInt(decoded.client_key.client_nonce), + }, + sequence_num: BigInt(decoded.sequence_num), // not clear why cbor deserializes bigint as number + timestamp: BigInt(decoded.timestamp), + content: decoded.content, + is_service_message: decoded.is_service_message, + } +}; + +export const getPollingNonceFromMessage = (message: CanisterOutputMessage): number => { + const nonceStr = message.key.split("_")[1]; + return parseInt(nonceStr); +}; + +export const getNextPollingNonceFromMessages = (messages: CanisterOutputMessage[]): number => { + return getPollingNonceFromMessage(messages[messages.length - 1]) + 1; +}; + +export const getCertifiedMessageKey = async (gatewayIdentity: Promise, nonce: number) => { + const gatewayPrincipal = (await gatewayIdentity).getPrincipal().toText(); + return `${gatewayPrincipal}_${String(nonce).padStart(20, '0')}`; +}; + +export const isValidCertificate = async (canisterId: string, certificate: Uint8Array, tree: Uint8Array, agent: HttpAgent) => { + const canisterPrincipal = Principal.fromText(canisterId); + let cert: Certificate; + + try { + cert = await Certificate.create({ + certificate, + canisterId: canisterPrincipal, + rootKey: agent.rootKey! + }); + } catch (error) { + console.error("Error creating certificate:", error); + return false; + } + + const hashTree = Cbor.decode(tree); + const reconstructed = await reconstruct(hashTree); + const witness = cert.lookup([ + "canister", + canisterPrincipal.toUint8Array(), + "certified_data" + ]); + + if (!witness) { + throw new Error( + "Could not find certified data for this canister in the certificate." + ); + } + + // First validate that the Tree is as good as the certification. + return compare(witness, reconstructed) === 0; +}; + +export const isMessageBodyValid = async (path: string, body: Uint8Array | ArrayBuffer, tree: Uint8Array) => { + const hashTree = Cbor.decode(tree); + const sha = await crypto.subtle.digest("SHA-256", body); + let treeSha = lookup_path(["websocket", path], hashTree); + + if (!treeSha) { + // Allow fallback to index path. + treeSha = lookup_path(["websocket"], hashTree); + } + + return !!treeSha && (compare(sha, treeSha) === 0); }; diff --git a/tests/src/lib.rs b/tests/src/lib.rs index 26932a3..68fe993 100644 --- a/tests/src/lib.rs +++ b/tests/src/lib.rs @@ -66,8 +66,14 @@ fn ws_wipe() { // send a message to the client, usually called by the canister itself #[update] -fn ws_send(client_principal: ClientPrincipal, msg_bytes: Vec) -> CanisterWsSendResult { - ic_websocket_cdk::ws_send(client_principal, msg_bytes) +fn ws_send(client_principal: ClientPrincipal, messages: Vec>) -> CanisterWsSendResult { + for msg_bytes in messages { + match ic_websocket_cdk::ws_send(client_principal, msg_bytes) { + Ok(_) => {}, + Err(e) => return Err(e), + } + } + Ok(()) } // reinitialize the canister diff --git a/tests/test_canister.did b/tests/test_canister.did index da066d3..785e50f 100644 --- a/tests/test_canister.did +++ b/tests/test_canister.did @@ -13,6 +13,6 @@ service : (text, nat64, nat64) -> { // methods used just for debugging/testing "ws_wipe" : () -> (); - "ws_send" : (ClientPrincipal, blob) -> (CanisterWsSendResult); + "ws_send" : (ClientPrincipal, vec blob) -> (CanisterWsSendResult); "reinitialize" : (text, nat64, nat64) -> (); };