From f1bd4262bca1260850b31c1e91fdf3109315b80b Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Tue, 5 Dec 2023 14:12:33 +0100 Subject: [PATCH 01/18] fix: remove catch_unwind from callback handlers --- .../src/tests/unit_tests/mod.rs | 39 ------------------- src/ic-websocket-cdk/src/types.rs | 36 ++++++----------- 2 files changed, 11 insertions(+), 64 deletions(-) diff --git a/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs b/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs index 04fdac4..18405ad 100644 --- a/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs +++ b/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs @@ -106,45 +106,6 @@ fn test_ws_handlers_are_called() { assert!(CUSTOM_STATE.with(|h| h.borrow().is_on_close_called)); } -#[test] -fn test_ws_handlers_panic_is_handled() { - let h = WsHandlers { - on_open: Some(|_| { - panic!("on_open_panic"); - }), - on_message: Some(|_| { - panic!("on_close_panic"); - }), - on_close: Some(|_| { - panic!("on_close_panic"); - }), - }; - - set_params(WsInitParams::new(h)); - - let handlers = get_handlers_from_params(); - - let res = panic::catch_unwind(|| { - handlers.call_on_open(OnOpenCallbackArgs { - client_principal: common::generate_random_principal(), - }); - }); - assert!(res.is_ok()); - let res = panic::catch_unwind(|| { - handlers.call_on_message(OnMessageCallbackArgs { - client_principal: common::generate_random_principal(), - message: vec![], - }); - }); - assert!(res.is_ok()); - let res = panic::catch_unwind(|| { - handlers.call_on_close(OnCloseCallbackArgs { - client_principal: common::generate_random_principal(), - }); - }); - assert!(res.is_ok()); -} - #[test] fn test_ws_init_params() { let handlers = WsHandlers::default(); diff --git a/src/ic-websocket-cdk/src/types.rs b/src/ic-websocket-cdk/src/types.rs index 35a211c..70abc47 100644 --- a/src/ic-websocket-cdk/src/types.rs +++ b/src/ic-websocket-cdk/src/types.rs @@ -1,13 +1,13 @@ -use std::{collections::VecDeque, fmt, panic, time::Duration}; +use std::{collections::VecDeque, fmt, time::Duration}; use candid::{decode_one, CandidType, Principal}; use serde::{Deserialize, Serialize}; use serde_cbor::Serializer; use crate::{ - custom_print, custom_trap, errors::WsError, utils::get_current_time, - DEFAULT_CLIENT_KEEP_ALIVE_TIMEOUT_MS, DEFAULT_MAX_NUMBER_OF_RETURNED_MESSAGES, - DEFAULT_SEND_ACK_INTERVAL_MS, INITIAL_OUTGOING_MESSAGE_NONCE, + custom_trap, errors::WsError, utils::get_current_time, DEFAULT_CLIENT_KEEP_ALIVE_TIMEOUT_MS, + DEFAULT_MAX_NUMBER_OF_RETURNED_MESSAGES, DEFAULT_SEND_ACK_INTERVAL_MS, + INITIAL_OUTGOING_MESSAGE_NONCE, }; pub type ClientPrincipal = Principal; @@ -340,37 +340,23 @@ pub struct WsHandlers { impl WsHandlers { pub(crate) fn call_on_open(&self, args: OnOpenCallbackArgs) { if let Some(on_open) = self.on_open { - let res = panic::catch_unwind(|| { - on_open(args); - }); - - if let Err(e) = res { - custom_print!("Error calling on_open handler: {:?}", e); - } + // we don't have to recover from errors here, + // we just let the canister trap + on_open(args); } } pub(crate) fn call_on_message(&self, args: OnMessageCallbackArgs) { if let Some(on_message) = self.on_message { - let res = panic::catch_unwind(|| { - on_message(args); - }); - - if let Err(e) = res { - custom_print!("Error calling on_message handler: {:?}", e); - } + // see call_on_open + on_message(args); } } pub(crate) fn call_on_close(&self, args: OnCloseCallbackArgs) { if let Some(on_close) = self.on_close { - let res = panic::catch_unwind(|| { - on_close(args); - }); - - if let Err(e) = res { - custom_print!("Error calling on_close handler: {:?}", e); - } + // see call_on_open + on_close(args); } } } From 59c8eb5f95e8ab381252572bba3b098461043035 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Tue, 5 Dec 2023 14:21:03 +0100 Subject: [PATCH 02/18] fix: do not overflow on subtraction --- src/ic-websocket-cdk/src/types.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ic-websocket-cdk/src/types.rs b/src/ic-websocket-cdk/src/types.rs index 70abc47..1d2d874 100644 --- a/src/ic-websocket-cdk/src/types.rs +++ b/src/ic-websocket-cdk/src/types.rs @@ -173,7 +173,7 @@ impl RegisteredGateway { /// Decrements the connected clients count by 1, returning the new value. pub(crate) fn decrement_clients_count(&mut self) -> u64 { - self.connected_clients_count -= 1; + self.connected_clients_count = self.connected_clients_count.saturating_sub(1); self.connected_clients_count } From 022a549c111dd71f36612becf3dd2e537ea242da Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Tue, 5 Dec 2023 15:59:46 +0100 Subject: [PATCH 03/18] feat: send close message to client on some errors --- src/ic-websocket-cdk/service_messages.did | 12 +++ src/ic-websocket-cdk/src/lib.rs | 6 +- src/ic-websocket-cdk/src/state.rs | 57 +++++++++---- .../tests/integration_tests/b_ws_message.rs | 80 ++++++++++++++----- .../f_messages_acknowledgement.rs | 56 +++++++------ .../tests/integration_tests/utils/messages.rs | 17 +++- .../src/tests/unit_tests/mod.rs | 80 ++++++++++++++++++- src/ic-websocket-cdk/src/timers.rs | 2 +- src/ic-websocket-cdk/src/types.rs | 19 +++++ 9 files changed, 261 insertions(+), 68 deletions(-) diff --git a/src/ic-websocket-cdk/service_messages.did b/src/ic-websocket-cdk/service_messages.did index 81da189..7b9b763 100644 --- a/src/ic-websocket-cdk/service_messages.did +++ b/src/ic-websocket-cdk/service_messages.did @@ -16,8 +16,20 @@ type ClientKeepAliveMessageContent = record { last_incoming_sequence_num : nat64; }; +type CloseMessageReason = variant { + WrongSequenceNumber; + InvalidServiceMessage; + KeepAliveTimeout; + ClosedByApplication; +}; + +type CanisterCloseMessageContent = record { + reason : CloseMessageReason; +}; + type WebsocketServiceMessageContent = variant { OpenMessage : CanisterOpenMessageContent; AckMessage : CanisterAckMessageContent; KeepAliveMessage : ClientKeepAliveMessageContent; + CloseMessage : CanisterCloseMessageContent; }; diff --git a/src/ic-websocket-cdk/src/lib.rs b/src/ic-websocket-cdk/src/lib.rs index d0e4acc..7000fa0 100644 --- a/src/ic-websocket-cdk/src/lib.rs +++ b/src/ic-websocket-cdk/src/lib.rs @@ -88,7 +88,7 @@ pub fn ws_open(args: CanisterWsOpenArguments) -> CanisterWsOpenResult { // Do nothing }, Ok(old_client_key) => { - remove_client(&old_client_key); + remove_client(&old_client_key, None); }, }; @@ -124,7 +124,7 @@ pub fn ws_close(args: CanisterWsCloseArguments) -> CanisterWsCloseResult { // check if the client is registered to the gateway that is closing the connection check_client_registered_to_gateway(&args.client_key, &gateway_principal)?; - remove_client(&args.client_key); + remove_client(&args.client_key, None); Ok(()) } @@ -188,7 +188,7 @@ pub fn ws_message Deserialize<'a>>( .eq(&expected_sequence_num) .then_some(()) .ok_or_else(|| { - remove_client(&client_key); + remove_client(&client_key, Some(CloseMessageReason::WrongSequenceNumber)); WsError::IncomingSequenceNumberWrong { expected_sequence_num, diff --git a/src/ic-websocket-cdk/src/state.rs b/src/ic-websocket-cdk/src/state.rs index f521f57..7ef9041 100644 --- a/src/ic-websocket-cdk/src/state.rs +++ b/src/ic-websocket-cdk/src/state.rs @@ -5,6 +5,7 @@ use std::{ }; use candid::{encode_one, Principal}; +#[allow(unused_imports)] use ic_cdk::api::{data_certificate, set_certified_data}; use ic_certified_map::{labeled, labeled_hash, AsHashTree, Hash as ICHash, RbTree}; use serde::Serialize; @@ -44,7 +45,7 @@ pub(crate) fn reset_internal_state() { // for each client, call the on_close handler before clearing the map for client_key in client_keys_to_remove { - remove_client(&client_key); + remove_client(&client_key, None); } // make sure all the maps are cleared @@ -79,16 +80,22 @@ pub(crate) fn increment_gateway_clients_count(gateway_principal: GatewayPrincipa }); } -/// Decrements the clients connected count for the given gateway. -/// If there are no more clients connected, the gateway is removed from the list of registered gateways. -pub(crate) fn decrement_gateway_clients_count(gateway_principal: &GatewayPrincipal) { +/// Decrements the clients connected count for the given gateway, if it exists. +/// +/// If `remove_if_empty` is true, the gateway is removed from the list of registered gateways +/// if it has no clients connected. +pub(crate) fn decrement_gateway_clients_count( + gateway_principal: &GatewayPrincipal, + remove_if_empty: bool, +) { REGISTERED_GATEWAYS.with(|map| { let mut map = map.borrow_mut(); - let g = map.get_mut(gateway_principal).unwrap(); // gateway must be registered at this point - let clients_count = g.decrement_clients_count(); + if let Some(g) = map.get_mut(gateway_principal) { + let clients_count = g.decrement_clients_count(); - if clients_count == 0 { - map.remove(gateway_principal); + if remove_if_empty && clients_count == 0 { + map.remove(gateway_principal); + } } }); } @@ -266,7 +273,23 @@ pub(crate) fn add_client(client_key: ClientKey, new_client: RegisteredClient) { /// Removes a client from the internal state /// and call the on_close callback, /// if the client was registered in the state. -pub(crate) fn remove_client(client_key: &ClientKey) { +/// +/// If a `close_reason` is provided, it also sends a close message to the client, +/// so that the client can close the WS connection with the gateway. +/// +/// If a `close_reason` is **not** provided, it also removes the gateway from the state +/// if it has no clients connected anymore. +pub(crate) fn remove_client(client_key: &ClientKey, close_reason: Option) { + if let Some(close_reason) = close_reason.clone() { + // ignore the error + let _ = send_service_message_to_client( + client_key, + &WebsocketServiceMessageContent::CloseMessage(CanisterCloseMessageContent { + reason: close_reason, + }), + ); + } + CLIENTS_WAITING_FOR_KEEP_ALIVE.with(|set| { set.borrow_mut().remove(client_key); }); @@ -283,7 +306,10 @@ pub(crate) fn remove_client(client_key: &ClientKey) { if let Some(registered_client) = REGISTERED_CLIENTS.with(|map| map.borrow_mut().remove(client_key)) { - decrement_gateway_clients_count(®istered_client.gateway_principal); + decrement_gateway_clients_count( + ®istered_client.gateway_principal, + close_reason.is_none(), + ); let handlers = get_handlers_from_params(); handlers.call_on_close(OnCloseCallbackArgs { @@ -383,12 +409,15 @@ pub(crate) fn get_cert_messages_empty() -> CanisterWsGetMessagesResult { } fn put_cert_for_message(key: String, value: &Vec) { + #[allow(unused_variables)] let root_hash = CERT_TREE.with(|tree| { let mut tree = tree.borrow_mut(); tree.insert(key.clone(), Sha256::digest(value).into()); labeled_hash(LABEL_WEBSOCKET, &tree.root_hash()) }); + #[cfg(not(test))] + // executing this in tests fails because the tree is an IC-specific implementation set_certified_data(&root_hash); } @@ -419,9 +448,6 @@ pub(crate) fn delete_old_messages_for_gateway( ) -> Result<(), String> { let ack_interval_ms = get_params().send_ack_interval_ms; - // allow unused variables because sometimes the compiler complains about unused variables - // since it is only used in production code - #[allow(unused_variables)] let deleted_messages_keys = REGISTERED_GATEWAYS.with(|map| { map.borrow_mut() .get_mut(gateway_principal) @@ -429,8 +455,6 @@ pub(crate) fn delete_old_messages_for_gateway( .and_then(|g| Ok(g.delete_old_messages(MESSAGES_TO_DELETE_COUNT, ack_interval_ms))) })?; - #[cfg(not(test))] - // executing this in tests fails because the tree is an IC-specific implementation CERT_TREE.with(|tree| { for key in deleted_messages_keys { tree.borrow_mut().delete(key.as_ref()); @@ -487,7 +511,8 @@ pub(crate) fn handle_received_service_message( let decoded = WebsocketServiceMessageContent::from_candid_bytes(content)?; match decoded { WebsocketServiceMessageContent::OpenMessage(_) - | WebsocketServiceMessageContent::AckMessage(_) => { + | WebsocketServiceMessageContent::AckMessage(_) + | WebsocketServiceMessageContent::CloseMessage(_) => { WsError::InvalidServiceMessage.to_string_result() }, WebsocketServiceMessageContent::KeepAliveMessage(keep_alive_message) => { diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/b_ws_message.rs b/src/ic-websocket-cdk/src/tests/integration_tests/b_ws_message.rs index 5391b12..1a5aea9 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/b_ws_message.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/b_ws_message.rs @@ -3,9 +3,19 @@ use std::ops::Deref; use candid::encode_one; use crate::{ - errors::WsError, tests::common::generate_random_principal, CanisterAckMessageContent, - CanisterWsMessageArguments, CanisterWsMessageResult, ClientKeepAliveMessageContent, ClientKey, - WebsocketServiceMessageContent, + errors::WsError, + tests::{ + common::generate_random_principal, + integration_tests::utils::{ + actor::{ws_close::call_ws_close, ws_get_messages::call_ws_get_messages_with_panic}, + clients::GATEWAY_1, + messages::check_canister_message_has_close_reason, + }, + }, + types::CloseMessageReason, + CanisterAckMessageContent, CanisterWsCloseArguments, CanisterWsCloseResult, + CanisterWsGetMessagesArguments, CanisterWsMessageArguments, CanisterWsMessageResult, + ClientKeepAliveMessageContent, ClientKey, WebsocketServiceMessageContent, }; use super::utils::{ @@ -44,6 +54,10 @@ fn test_1_fails_if_client_is_not_registered() { #[test] fn test_2_fails_if_client_sends_a_message_with_a_different_client_key() { let client_1_key = CLIENT_1_KEY.deref(); + // first, reset the canister + get_test_env().reset_canister_with_default_params(); + // second, open a connection for client 1 + call_ws_open_for_client_key_with_panic(client_1_key); let wrong_client_key = ClientKey { client_principal: generate_random_principal(), @@ -93,6 +107,11 @@ fn test_2_fails_if_client_sends_a_message_with_a_different_client_key() { #[test] fn test_3_should_send_a_message_from_a_registered_client() { let client_1_key = CLIENT_1_KEY.deref(); + // first, reset the canister + get_test_env().reset_canister_with_default_params(); + // second, open a connection for client 1 + call_ws_open_for_client_key_with_panic(client_1_key); + let res = call_ws_message( &client_1_key.client_principal, CanisterWsMessageArguments { @@ -105,8 +124,13 @@ fn test_3_should_send_a_message_from_a_registered_client() { #[test] fn test_4_fails_if_client_sends_a_message_with_a_wrong_sequence_number() { let client_1_key = CLIENT_1_KEY.deref(); - let wrong_sequence_number = 1; // the message with sequence number 1 has already been sent in the previous test - let expected_sequence_number = 2; // the next valid sequence number + // first, reset the canister + get_test_env().reset_canister_with_default_params(); + // second, open a connection for client 1 + call_ws_open_for_client_key_with_panic(client_1_key); + + let wrong_sequence_number = 2; // the message with sequence number 1 has already been sent in the previous test + let expected_sequence_number = 1; // the next valid sequence number let res = call_ws_message( &client_1_key.client_principal, CanisterWsMessageArguments { @@ -124,28 +148,41 @@ fn test_4_fails_if_client_sends_a_message_with_a_wrong_sequence_number() { ) ); - // check if the client has been removed - let res = call_ws_message( - &client_1_key.client_principal, - CanisterWsMessageArguments { - msg: create_websocket_message(client_1_key, 1, None, false), // the sequence number doesn't matter here because the method fails before checking it + // check if the gateway put the close message in the queue + let msgs = call_ws_get_messages_with_panic( + GATEWAY_1.deref(), + CanisterWsGetMessagesArguments { nonce: 1 }, // skip the first open message + ); + check_canister_message_has_close_reason( + &msgs.messages[0], + CloseMessageReason::WrongSequenceNumber, + ); + + // the gateway should still be between the registered gateways + // so calling the ws_close endpoint should return the ClientKeyNotConnected error + let res = call_ws_close( + GATEWAY_1.deref(), + CanisterWsCloseArguments { + client_key: client_1_key.clone(), }, ); assert_eq!( res, - CanisterWsMessageResult::Err( - WsError::ClientPrincipalNotConnected { - client_principal: &client_1_key.client_principal + CanisterWsCloseResult::Err( + WsError::ClientKeyNotConnected { + client_key: &client_1_key } .to_string() ) - ) + ); } #[test] fn test_5_fails_if_client_sends_a_wrong_service_message() { let client_1_key = CLIENT_1_KEY.deref(); - // first, open the connection again for client 1 + // first, reset the canister + get_test_env().reset_canister_with_default_params(); + // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(client_1_key); // fail with wrong content encoding @@ -160,8 +197,10 @@ fn test_5_fails_if_client_sends_a_wrong_service_message() { ), }, ); - let err = res.err().unwrap(); - assert!(err.starts_with("Error decoding service message content:")); + assert!(res + .err() + .unwrap() + .starts_with("Error decoding service message content:")); // fail with wrong service message variant let wrong_service_message = @@ -190,6 +229,11 @@ fn test_5_fails_if_client_sends_a_wrong_service_message() { #[test] fn test_6_should_send_a_service_message_from_a_registered_client() { let client_1_key = CLIENT_1_KEY.deref(); + // first, reset the canister + get_test_env().reset_canister_with_default_params(); + // second, open a connection for client 1 + call_ws_open_for_client_key_with_panic(client_1_key); + let client_service_message = WebsocketServiceMessageContent::KeepAliveMessage(ClientKeepAliveMessageContent { last_incoming_sequence_num: 0, @@ -199,7 +243,7 @@ fn test_6_should_send_a_service_message_from_a_registered_client() { CanisterWsMessageArguments { msg: create_websocket_message( client_1_key, - 3, + 1, Some(encode_websocket_service_message_content( &client_service_message, )), diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/f_messages_acknowledgement.rs b/src/ic-websocket-cdk/src/tests/integration_tests/f_messages_acknowledgement.rs index 652e753..c4c79db 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/f_messages_acknowledgement.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/f_messages_acknowledgement.rs @@ -1,15 +1,19 @@ use std::ops::Deref; use crate::{ - errors::WsError, CanisterOutputCertifiedMessages, CanisterWsGetMessagesArguments, - CanisterWsMessageArguments, CanisterWsSendResult, ClientKeepAliveMessageContent, + errors::WsError, + tests::integration_tests::utils::{ + actor::ws_close::call_ws_close, messages::check_canister_message_has_close_reason, + }, + types::CloseMessageReason, + CanisterOutputCertifiedMessages, CanisterWsCloseArguments, CanisterWsCloseResult, + CanisterWsGetMessagesArguments, CanisterWsMessageArguments, ClientKeepAliveMessageContent, WebsocketServiceMessageContent, }; use super::utils::{ actor::{ - ws_get_messages::call_ws_get_messages_with_panic, - ws_message::{call_ws_message, call_ws_message_with_panic}, + ws_get_messages::call_ws_get_messages_with_panic, ws_message::call_ws_message_with_panic, ws_open::call_ws_open_for_client_key_with_panic, }, certification::{is_message_body_valid, is_valid_certificate}, @@ -61,41 +65,41 @@ fn test_2_client_is_removed_if_keep_alive_timeout_is_reached() { // advance the canister time to make sure the ack timer expires and an ack is sent get_test_env().advance_canister_time_ms(DEFAULT_TEST_SEND_ACK_INTERVAL_MS); // get messages to check if the ack message has been set - let res = call_ws_get_messages_with_panic( + let msgs = call_ws_get_messages_with_panic( GATEWAY_1.deref(), CanisterWsGetMessagesArguments { nonce: 1 }, ); - helpers::check_ack_message_result(&res, client_1_key, 0, 2); + helpers::check_ack_message_result(&msgs, client_1_key, 0, 2); // advance the canister time to make sure the keep alive timeout expires get_test_env().advance_canister_time_ms(DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS); - // to check if the client has been removed, we try to send the keep alive message late - let res = call_ws_message( - &client_1_key.client_principal, - CanisterWsMessageArguments { - msg: create_websocket_message( - client_1_key, - 1, - Some(encode_websocket_service_message_content( - &WebsocketServiceMessageContent::KeepAliveMessage( - ClientKeepAliveMessageContent { - last_incoming_sequence_num: 1, // ignored in the CDK - }, - ), - )), - true, - ), + // check if the gateway put the close message in the queue + let msgs = call_ws_get_messages_with_panic( + GATEWAY_1.deref(), + CanisterWsGetMessagesArguments { nonce: 1 }, // skip the first open message + ); + check_canister_message_has_close_reason( + &msgs.messages[1], + CloseMessageReason::KeepAliveTimeout, + ); + + // the gateway should still be between the registered gateways + // so calling the ws_close endpoint should return the ClientKeyNotConnected error + let res = call_ws_close( + GATEWAY_1.deref(), + CanisterWsCloseArguments { + client_key: client_1_key.clone(), }, ); assert_eq!( res, - CanisterWsSendResult::Err( - WsError::ClientPrincipalNotConnected { - client_principal: &client_1_key.client_principal + CanisterWsCloseResult::Err( + WsError::ClientKeyNotConnected { + client_key: &client_1_key } .to_string() - ), + ) ); } diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/utils/messages.rs b/src/ic-websocket-cdk/src/tests/integration_tests/utils/messages.rs index 0a43b39..de9891e 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/utils/messages.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/utils/messages.rs @@ -6,7 +6,10 @@ use super::{ get_current_timestamp_ns, test_env::get_test_env, }; -use crate::{CanisterOutputMessage, ClientKey, WebsocketMessage, WebsocketServiceMessageContent}; +use crate::{ + types::{CanisterCloseMessageContent, CloseMessageReason}, + CanisterOutputMessage, ClientKey, WebsocketMessage, WebsocketServiceMessageContent, +}; #[derive(CandidType, Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] pub(crate) struct AppMessage { @@ -136,3 +139,15 @@ fn verify_message( assert!(is_valid_certificate(&get_test_env(), cert, tree,)); assert!(is_message_body_valid(&message.key, &message.content, tree)); } + +pub(in crate::tests::integration_tests) fn check_canister_message_has_close_reason( + msg: &CanisterOutputMessage, + close_reason: CloseMessageReason, +) { + assert_eq!( + get_service_message_content_from_canister_message(msg), + WebsocketServiceMessageContent::CloseMessage(CanisterCloseMessageContent { + reason: close_reason + }), + ); +} diff --git a/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs b/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs index 18405ad..81fdda0 100644 --- a/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs +++ b/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs @@ -3,6 +3,7 @@ use std::{cell::RefCell, panic}; use super::common; use crate::utils::get_current_time; use crate::*; +use candid::decode_one; use proptest::prelude::*; mod utils; @@ -206,7 +207,29 @@ proptest! { while clients_count > 0 { let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_gateway_principal).cloned()).unwrap(); prop_assert_eq!(registered_gateway.connected_clients_count, clients_count); - decrement_gateway_clients_count(&test_gateway_principal); + decrement_gateway_clients_count(&test_gateway_principal, false); + clients_count -= 1; + } + + let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_gateway_principal).cloned()); + prop_assert!(registered_gateway.is_some()); + } + + #[test] + fn test_decrement_gateway_clients_count_and_remove_gateway(test_gateway_principal in any::().prop_map(|_| common::generate_random_principal()), test_connected_clients_count in (1..1000u64)) { + // Set up + REGISTERED_GATEWAYS.with(|n| { + let mut gw = RegisteredGateway::new(); + gw.connected_clients_count = test_connected_clients_count; + n.borrow_mut().insert(test_gateway_principal, gw); + }); + + // Test + let mut clients_count = test_connected_clients_count; + while clients_count > 0 { + let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_gateway_principal).cloned()).unwrap(); + prop_assert_eq!(registered_gateway.connected_clients_count, clients_count); + decrement_gateway_clients_count(&test_gateway_principal, true); clients_count -= 1; } @@ -538,7 +561,7 @@ proptest! { #[test] fn test_remove_client_nonexistent(test_client_key in any::().prop_map(|_| common::get_random_client_key())) { let res = panic::catch_unwind(|| { - remove_client(&test_client_key); + remove_client(&test_client_key, None); }); prop_assert!(res.is_ok()); } @@ -566,7 +589,7 @@ proptest! { map.borrow_mut().insert(test_client_key.clone(), INITIAL_CANISTER_SEQUENCE_NUM); }); - remove_client(&test_client_key); + remove_client(&test_client_key, None); let client_key = CURRENT_CLIENT_KEY_MAP.with(|map| map.borrow().get(&test_client_key.client_principal).cloned()); prop_assert!(client_key.is_none()); @@ -584,6 +607,57 @@ proptest! { prop_assert!(outgoing_seq_num.is_none()); } + #[test] + fn test_remove_client_with_reason(test_client_key in any::().prop_map(|_| common::get_random_client_key())) { + // Set up + CURRENT_CLIENT_KEY_MAP.with(|map| { + map.borrow_mut().insert(test_client_key.client_principal.clone(), test_client_key.clone()); + }); + let test_registered_client = utils::generate_random_registered_client(); + REGISTERED_CLIENTS.with(|map| { + map.borrow_mut().insert(test_client_key.clone(), test_registered_client.clone()); + }); + REGISTERED_GATEWAYS.with(|map| { + let mut gw = RegisteredGateway::new(); + gw.connected_clients_count = 1; + map.borrow_mut() + .insert(test_registered_client.gateway_principal, gw); + }); + INCOMING_MESSAGE_FROM_CLIENT_NUM_MAP.with(|map| { + map.borrow_mut().insert(test_client_key.clone(), INITIAL_CLIENT_SEQUENCE_NUM); + }); + OUTGOING_MESSAGE_TO_CLIENT_NUM_MAP.with(|map| { + map.borrow_mut().insert(test_client_key.clone(), INITIAL_CANISTER_SEQUENCE_NUM); + }); + + remove_client(&test_client_key, Some(CloseMessageReason::ClosedByApplication)); + + let client_key = CURRENT_CLIENT_KEY_MAP.with(|map| map.borrow().get(&test_client_key.client_principal).cloned()); + prop_assert!(client_key.is_none()); + + let registered_client = REGISTERED_CLIENTS.with(|map| map.borrow().get(&test_client_key).cloned()); + prop_assert!(registered_client.is_none()); + + let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_registered_client.gateway_principal).cloned()).unwrap(); + prop_assert!(registered_gateway.connected_clients_count == 0); + let expected_websocket_message: WebsocketMessage = serde_cbor::from_slice(®istered_gateway.messages_queue[0].content).unwrap(); + let expected_service_content = decode_one(&expected_websocket_message.content).unwrap(); + assert!( + matches!( + expected_service_content, + WebsocketServiceMessageContent::CloseMessage(CanisterCloseMessageContent{ + reason: CloseMessageReason::ClosedByApplication + }), + ) + ); + + let incoming_seq_num = INCOMING_MESSAGE_FROM_CLIENT_NUM_MAP.with(|map| map.borrow().get(&test_client_key).cloned()); + prop_assert!(incoming_seq_num.is_none()); + + let outgoing_seq_num = OUTGOING_MESSAGE_TO_CLIENT_NUM_MAP.with(|map| map.borrow().get(&test_client_key).cloned()); + prop_assert!(outgoing_seq_num.is_none()); + } + #[test] fn test_format_message_for_gateway_key(test_gateway_principal in any::().prop_map(|_| common::generate_random_principal()), test_nonce in any::()) { let message_key = format_message_for_gateway_key(&test_gateway_principal, test_nonce); diff --git a/src/ic-websocket-cdk/src/timers.rs b/src/ic-websocket-cdk/src/timers.rs index c76f89d..0851515 100644 --- a/src/ic-websocket-cdk/src/timers.rs +++ b/src/ic-websocket-cdk/src/timers.rs @@ -133,7 +133,7 @@ fn check_keep_alive_timer_callback(keep_alive_timeout_ms: u64) { .collect(); for client_key in client_keys_to_remove { - remove_client(&client_key); + remove_client(&client_key, Some(CloseMessageReason::KeepAliveTimeout)); custom_print!( "[check-keep-alive-timer-cb]: Client {} has not sent a keep alive message in the last {} ms and has been removed", diff --git a/src/ic-websocket-cdk/src/types.rs b/src/ic-websocket-cdk/src/types.rs index 1d2d874..60ad47f 100644 --- a/src/ic-websocket-cdk/src/types.rs +++ b/src/ic-websocket-cdk/src/types.rs @@ -263,6 +263,23 @@ pub(crate) struct ClientKeepAliveMessageContent { pub(crate) last_incoming_sequence_num: u64, } +#[derive(CandidType, Clone, Debug, Deserialize, PartialEq, Eq)] +pub(crate) enum CloseMessageReason { + /// When the canister receives a wrong sequence number from the client. + WrongSequenceNumber, + /// When the canister receives an invalid service message from the client. + InvalidServiceMessage, + /// When the canister doesn't receive the keep alive message from the client in time. + KeepAliveTimeout, + /// When the developer calls the `close` function. + ClosedByApplication, +} + +#[derive(CandidType, Debug, Deserialize, PartialEq, Eq)] +pub(crate) struct CanisterCloseMessageContent { + pub(crate) reason: CloseMessageReason, +} + /// A service message sent by the CDK to the client or vice versa. #[derive(CandidType, Debug, Deserialize, PartialEq, Eq)] pub(crate) enum WebsocketServiceMessageContent { @@ -272,6 +289,8 @@ pub(crate) enum WebsocketServiceMessageContent { AckMessage(CanisterAckMessageContent), /// Message sent by the **client** in response to an acknowledgement message from the canister. KeepAliveMessage(ClientKeepAliveMessageContent), + /// Message sent by the **canister** when it wants to close the connection. + CloseMessage(CanisterCloseMessageContent), } impl WebsocketServiceMessageContent { From 08963c95c543d35cf58b5bff307d7e0564fb9deb Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 8 Dec 2023 10:22:45 +0100 Subject: [PATCH 04/18] feat: delete keys from cert map if gateway empty --- src/ic-websocket-cdk/src/state.rs | 38 ++++++++++++++----- .../src/tests/unit_tests/mod.rs | 30 +++++++++++++++ .../src/tests/unit_tests/utils.rs | 4 ++ 3 files changed, 63 insertions(+), 9 deletions(-) diff --git a/src/ic-websocket-cdk/src/state.rs b/src/ic-websocket-cdk/src/state.rs index 7ef9041..846d516 100644 --- a/src/ic-websocket-cdk/src/state.rs +++ b/src/ic-websocket-cdk/src/state.rs @@ -29,7 +29,7 @@ thread_local! { /// Maps the client's key to the expected sequence number of the next incoming message (from that client). /* flexible */ pub(crate) static INCOMING_MESSAGE_FROM_CLIENT_NUM_MAP: RefCell> = RefCell::new(HashMap::new()); /// Keeps track of the Merkle tree used for certified queries - /* flexible */ static CERT_TREE: RefCell> = RefCell::new(RbTree::new()); + /* flexible */ pub(crate) static CERT_TREE: RefCell> = RefCell::new(RbTree::new()); /// Keeps track of the principals of the WS Gateways that poll the canister /* flexible */ pub(crate) static REGISTERED_GATEWAYS: RefCell> = RefCell::new(HashMap::new()); /// The parameters passed in the CDK initialization @@ -88,16 +88,24 @@ pub(crate) fn decrement_gateway_clients_count( gateway_principal: &GatewayPrincipal, remove_if_empty: bool, ) { - REGISTERED_GATEWAYS.with(|map| { + let messages_keys_to_delete = REGISTERED_GATEWAYS.with(|map| { let mut map = map.borrow_mut(); if let Some(g) = map.get_mut(gateway_principal) { let clients_count = g.decrement_clients_count(); if remove_if_empty && clients_count == 0 { - map.remove(gateway_principal); + let g = map.remove(gateway_principal).unwrap(); + + return Some(g.messages_queue.iter().map(|m| m.key.clone()).collect()); } } + + None }); + + if let Some(messages_keys_to_delete) = messages_keys_to_delete { + delete_keys_from_cert_tree(messages_keys_to_delete); + } } pub(crate) fn get_registered_gateway( @@ -408,7 +416,7 @@ pub(crate) fn get_cert_messages_empty() -> CanisterWsGetMessagesResult { Ok(CanisterOutputCertifiedMessages::empty()) } -fn put_cert_for_message(key: String, value: &Vec) { +pub(crate) fn put_cert_for_message(key: String, value: &Vec) { #[allow(unused_variables)] let root_hash = CERT_TREE.with(|tree| { let mut tree = tree.borrow_mut(); @@ -417,7 +425,7 @@ fn put_cert_for_message(key: String, value: &Vec) { }); #[cfg(not(test))] - // executing this in tests fails because the tree is an IC-specific implementation + // executing this in unit tests fails because it's an IC-specific API set_certified_data(&root_hash); } @@ -455,13 +463,25 @@ pub(crate) fn delete_old_messages_for_gateway( .and_then(|g| Ok(g.delete_old_messages(MESSAGES_TO_DELETE_COUNT, ack_interval_ms))) })?; - CERT_TREE.with(|tree| { - for key in deleted_messages_keys { - tree.borrow_mut().delete(key.as_ref()); + delete_keys_from_cert_tree(deleted_messages_keys); + + Ok(()) +} + +pub(crate) fn delete_keys_from_cert_tree(keys: Vec) { + #[allow(unused_variables)] + let root_hash = CERT_TREE.with(|tree| { + let mut tree = tree.borrow_mut(); + for key in keys { + tree.delete(key.as_ref()); } + labeled_hash(LABEL_WEBSOCKET, &tree.root_hash()) }); - Ok(()) + // certify data with the new root hash + #[cfg(not(test))] + // executing this in unit tests fails because it's an IC-specific API + set_certified_data(&root_hash); } fn get_cert_for_range(first: &String, last: &String) -> (Vec, Vec) { diff --git a/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs b/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs index 81fdda0..8b87481 100644 --- a/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs +++ b/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs @@ -4,6 +4,7 @@ use super::common; use crate::utils::get_current_time; use crate::*; use candid::decode_one; +use ic_certified_map::RbTree; use proptest::prelude::*; mod utils; @@ -867,6 +868,19 @@ proptest! { utils::clean_messages_for_gateway(&gateway_principal); } + #[test] + fn test_put_cert_for_message(test_key in any::().prop_map(|_| utils::generate_random_message_key(&common::generate_random_principal()))) { + // Set up + CERT_TREE.with(|tree| tree.replace(RbTree::new())); + + // Test + put_cert_for_message(test_key.clone(), &vec![1,2,3]); // we don't care about the value here + prop_assert!(CERT_TREE.with(|tree| tree.borrow().get(&test_key.as_bytes()).is_some())); + + // Clean up + CERT_TREE.with(|tree| tree.replace(RbTree::new())); + } + #[test] fn test_delete_old_messages_for_gateway_nonexistent_gateway(gateway_principal in any::().prop_map(|_| common::generate_random_principal())) { let res = delete_old_messages_for_gateway(&gateway_principal); @@ -945,4 +959,20 @@ proptest! { prop_assert_eq!(registered_gateway.messages_queue.len(), expected_messages_left_count); prop_assert_eq!(registered_gateway.messages_to_delete.len(), expected_messages_left_count); } + + #[test] + fn test_delete_keys_from_cert_tree(test_keys in any::>().prop_map(|v| v.iter().map(|_| utils::generate_random_message_key(&common::generate_random_principal())).collect::>())) { + // Set up + CERT_TREE.with(|tree| { + let mut tree = tree.borrow_mut(); + for key in test_keys.clone() { + tree.insert(key.clone(), [0u8; 32]); // we don't care about the value here + } + }); + prop_assert_eq!(CERT_TREE.with(|tree| tree.borrow().iter().count()), test_keys.len()); + + // Test + delete_keys_from_cert_tree(test_keys); + prop_assert_eq!(CERT_TREE.with(|tree| tree.borrow().iter().count()), 0); + } } diff --git a/src/ic-websocket-cdk/src/tests/unit_tests/utils.rs b/src/ic-websocket-cdk/src/tests/unit_tests/utils.rs index 38e593d..fd6b6da 100644 --- a/src/ic-websocket-cdk/src/tests/unit_tests/utils.rs +++ b/src/ic-websocket-cdk/src/tests/unit_tests/utils.rs @@ -44,3 +44,7 @@ pub fn initialize_params() { ..Default::default() }); } + +pub fn generate_random_message_key(gateway_principal: &GatewayPrincipal) -> String { + format_message_for_gateway_key(gateway_principal, rand::random()) +} From 972b13cb0fe4d84975addcf5995c1b39e0bf45f2 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 8 Dec 2023 10:45:07 +0100 Subject: [PATCH 05/18] feat: close method --- src/ic-websocket-cdk/src/lib.rs | 21 +++++-- .../src/tests/integration_tests/e_ws_send.rs | 5 ++ .../src/tests/integration_tests/f_close.rs | 56 +++++++++++++++++++ ...ement.rs => g_messages_acknowledgement.rs} | 0 ...ple_gateways.rs => h_multiple_gateways.rs} | 0 .../src/tests/integration_tests/mod.rs | 5 +- .../tests/integration_tests/utils/actor.rs | 26 +++++++++ src/ic-websocket-cdk/src/types.rs | 2 + src/test_canister/src/lib.rs | 14 +++-- src/test_canister/test_canister.did | 6 ++ 10 files changed, 125 insertions(+), 10 deletions(-) create mode 100644 src/ic-websocket-cdk/src/tests/integration_tests/f_close.rs rename src/ic-websocket-cdk/src/tests/integration_tests/{f_messages_acknowledgement.rs => g_messages_acknowledgement.rs} (100%) rename src/ic-websocket-cdk/src/tests/integration_tests/{g_multiple_gateways.rs => h_multiple_gateways.rs} (100%) diff --git a/src/ic-websocket-cdk/src/lib.rs b/src/ic-websocket-cdk/src/lib.rs index 7000fa0..ac88de5 100644 --- a/src/ic-websocket-cdk/src/lib.rs +++ b/src/ic-websocket-cdk/src/lib.rs @@ -15,10 +15,11 @@ use state::*; use timers::*; use types::*; pub use types::{ - CanisterWsCloseArguments, CanisterWsCloseResult, CanisterWsGetMessagesArguments, - CanisterWsGetMessagesResult, CanisterWsMessageArguments, CanisterWsMessageResult, - CanisterWsOpenArguments, CanisterWsOpenResult, CanisterWsSendResult, ClientPrincipal, - OnCloseCallbackArgs, OnMessageCallbackArgs, OnOpenCallbackArgs, WsHandlers, WsInitParams, + CanisterCloseResult, CanisterWsCloseArguments, CanisterWsCloseResult, + CanisterWsGetMessagesArguments, CanisterWsGetMessagesResult, CanisterWsMessageArguments, + CanisterWsMessageResult, CanisterWsOpenArguments, CanisterWsOpenResult, CanisterWsSendResult, + ClientPrincipal, OnCloseCallbackArgs, OnMessageCallbackArgs, OnOpenCallbackArgs, WsHandlers, + WsInitParams, }; /// The label used when constructing the certification tree. @@ -112,6 +113,9 @@ pub fn ws_open(args: CanisterWsOpenArguments) -> CanisterWsOpenResult { } /// Handles the WS connection close event received from the WS Gateway. +/// +/// If you want to close the connection with the client in your logic, +/// use the [close] function instead. pub fn ws_close(args: CanisterWsCloseArguments) -> CanisterWsCloseResult { let gateway_principal = caller(); @@ -254,6 +258,15 @@ pub fn ws_send(client_principal: ClientPrincipal, msg_bytes: Vec) -> Caniste _ws_send(&client_key, msg_bytes, false) } +/// Closes the connection with the client. +pub fn close(client_principal: ClientPrincipal) -> CanisterCloseResult { + let client_key = get_client_key_from_principal(&client_principal)?; + + remove_client(&client_key, Some(CloseMessageReason::ClosedByApplication)); + + Ok(()) +} + /// Resets the internal state of the IC WebSocket CDK. /// /// **Note:** You should only call this function in tests. diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/e_ws_send.rs b/src/ic-websocket-cdk/src/tests/integration_tests/e_ws_send.rs index 92e40b5..b40c7b0 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/e_ws_send.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/e_ws_send.rs @@ -37,6 +37,11 @@ fn test_1_fails_if_sending_a_message_to_a_non_registered_client() { #[test] fn test_2_should_send_a_message_to_a_registered_client() { + // first, reset the canister + get_test_env().reset_canister_with_default_params(); + // second, open a connection for client 1 + call_ws_open_for_client_key_with_panic(CLIENT_1_KEY.deref()); + let res = call_ws_send( &CLIENT_1_KEY.client_principal, vec![AppMessage { diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/f_close.rs b/src/ic-websocket-cdk/src/tests/integration_tests/f_close.rs new file mode 100644 index 0000000..0aa647e --- /dev/null +++ b/src/ic-websocket-cdk/src/tests/integration_tests/f_close.rs @@ -0,0 +1,56 @@ +use std::ops::Deref; + +use crate::{ + errors::WsError, tests::integration_tests::utils::actor::close::call_close, CanisterCloseResult, +}; + +use super::utils::{ + actor::ws_open::call_ws_open_for_client_key_with_panic, + clients::{CLIENT_1_KEY, CLIENT_2_KEY}, + test_env::get_test_env, +}; + +#[test] +fn test_1_fails_closing_for_non_registered_client() { + // first, reset the canister + get_test_env().reset_canister_with_default_params(); + // second, open a connection for client 1 + call_ws_open_for_client_key_with_panic(CLIENT_1_KEY.deref()); + + // finally, we can start testing + let client_2_principal = &CLIENT_2_KEY.client_principal; + let res = call_close(client_2_principal); + assert_eq!( + res, + CanisterCloseResult::Err( + WsError::ClientPrincipalNotConnected { + client_principal: client_2_principal + } + .to_string() + ), + ); +} + +#[test] +fn test_2_should_close_connection_for_registered_client() { + // first, reset the canister + get_test_env().reset_canister_with_default_params(); + // second, open a connection for client 1 + call_ws_open_for_client_key_with_panic(CLIENT_1_KEY.deref()); + + // finally, we can start testing + let client_1_principal = &CLIENT_1_KEY.client_principal; + let res = call_close(client_1_principal); + assert_eq!(res, CanisterCloseResult::Ok(())); + // call the same function again, and expect it to fail + let res = call_close(client_1_principal); + assert_eq!( + res, + CanisterCloseResult::Err( + WsError::ClientPrincipalNotConnected { + client_principal: client_1_principal + } + .to_string() + ), + ); +} diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/f_messages_acknowledgement.rs b/src/ic-websocket-cdk/src/tests/integration_tests/g_messages_acknowledgement.rs similarity index 100% rename from src/ic-websocket-cdk/src/tests/integration_tests/f_messages_acknowledgement.rs rename to src/ic-websocket-cdk/src/tests/integration_tests/g_messages_acknowledgement.rs diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/g_multiple_gateways.rs b/src/ic-websocket-cdk/src/tests/integration_tests/h_multiple_gateways.rs similarity index 100% rename from src/ic-websocket-cdk/src/tests/integration_tests/g_multiple_gateways.rs rename to src/ic-websocket-cdk/src/tests/integration_tests/h_multiple_gateways.rs diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/mod.rs b/src/ic-websocket-cdk/src/tests/integration_tests/mod.rs index 8188d2a..a261e8d 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/mod.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/mod.rs @@ -12,5 +12,6 @@ mod b_ws_message; mod c_ws_get_messages; mod d_ws_close; mod e_ws_send; -mod f_messages_acknowledgement; -mod g_multiple_gateways; +mod f_close; +mod g_messages_acknowledgement; +mod h_multiple_gateways; diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs b/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs index 0b62428..05fff54 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs @@ -228,3 +228,29 @@ pub mod ws_send { } } } + +pub mod close { + use crate::{CanisterCloseResult, ClientPrincipal}; + use candid::encode_args; + + use super::*; + + /// # Panics + /// if the call returns a [WasmResult::Reject]. + pub(crate) fn call_close(client_principal: &ClientPrincipal) -> CanisterCloseResult { + let canister_id = get_test_env().canister_id; + let res = get_test_env() + .pic + .update_call( + canister_id, + Principal::anonymous(), + "close", + encode_args((client_principal,)).unwrap(), + ) + .expect("Failed to call canister"); + match res { + WasmResult::Reply(bytes) => decode_one(&bytes).unwrap(), + _ => panic!("Expected reply"), + } + } +} diff --git a/src/ic-websocket-cdk/src/types.rs b/src/ic-websocket-cdk/src/types.rs index 60ad47f..2dc7090 100644 --- a/src/ic-websocket-cdk/src/types.rs +++ b/src/ic-websocket-cdk/src/types.rs @@ -43,6 +43,8 @@ pub type CanisterWsMessageResult = Result<(), String>; pub type CanisterWsGetMessagesResult = Result; /// The result of [ws_send]. pub type CanisterWsSendResult = Result<(), String>; +/// The result of [close](crate::close). +pub type CanisterCloseResult = Result<(), String>; /// The arguments for [ws_open]. #[derive(CandidType, Clone, Deserialize, Serialize, Eq, PartialEq, Debug)] diff --git a/src/test_canister/src/lib.rs b/src/test_canister/src/lib.rs index d559fba..622c1a3 100644 --- a/src/test_canister/src/lib.rs +++ b/src/test_canister/src/lib.rs @@ -2,10 +2,10 @@ use ic_cdk_macros::*; use canister::{on_close, on_message, on_open, AppMessage}; use ic_websocket_cdk::{ - CanisterWsCloseArguments, CanisterWsCloseResult, CanisterWsGetMessagesArguments, - CanisterWsGetMessagesResult, CanisterWsMessageArguments, CanisterWsMessageResult, - CanisterWsOpenArguments, CanisterWsOpenResult, CanisterWsSendResult, ClientPrincipal, - WsHandlers, WsInitParams, + CanisterCloseResult, CanisterWsCloseArguments, CanisterWsCloseResult, + CanisterWsGetMessagesArguments, CanisterWsGetMessagesResult, CanisterWsMessageArguments, + CanisterWsMessageResult, CanisterWsOpenArguments, CanisterWsOpenResult, CanisterWsSendResult, + ClientPrincipal, WsHandlers, WsInitParams, }; mod canister; @@ -84,3 +84,9 @@ fn ws_send(client_principal: ClientPrincipal, messages: Vec>) -> Caniste } Ok(()) } + +// close the connection with a client, usually called by the canister itself +#[update] +fn close(client_principal: ClientPrincipal) -> CanisterCloseResult { + ic_websocket_cdk::close(client_principal) +} diff --git a/src/test_canister/test_canister.did b/src/test_canister/test_canister.did index 11f9c11..3044dc6 100644 --- a/src/test_canister/test_canister.did +++ b/src/test_canister/test_canister.did @@ -5,6 +5,11 @@ type CanisterWsSendResult = variant { Err : text; }; +type CanisterCloseResult = variant { + Ok : null; + Err : text; +}; + type AppMessage = record { text : text; }; @@ -17,4 +22,5 @@ service : (text, nat64, nat64, nat64) -> { // methods used just for debugging/testing "ws_send" : (ClientPrincipal, vec blob) -> (CanisterWsSendResult); + "close" : (ClientPrincipal) -> (CanisterCloseResult); }; From da04a905a66de0c94d22a78af454715714431c12 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 8 Dec 2023 10:46:26 +0100 Subject: [PATCH 06/18] docs: fix links --- src/ic-websocket-cdk/src/types.rs | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/ic-websocket-cdk/src/types.rs b/src/ic-websocket-cdk/src/types.rs index 2dc7090..aee198e 100644 --- a/src/ic-websocket-cdk/src/types.rs +++ b/src/ic-websocket-cdk/src/types.rs @@ -33,39 +33,39 @@ impl fmt::Display for ClientKey { } } -/// The result of [ws_open]. +/// The result of [ws_open](crate::ws_open). pub type CanisterWsOpenResult = Result<(), String>; -/// The result of [ws_close]. +/// The result of [ws_close](crate::ws_close). pub type CanisterWsCloseResult = Result<(), String>; -/// The result of [ws_message]. +/// The result of [ws_message](crate::ws_message). pub type CanisterWsMessageResult = Result<(), String>; -/// The result of [ws_get_messages]. +/// The result of [ws_get_messages](crate::ws_get_messages). pub type CanisterWsGetMessagesResult = Result; -/// The result of [ws_send]. +/// The result of [ws_send](crate::ws_send). pub type CanisterWsSendResult = Result<(), String>; /// The result of [close](crate::close). pub type CanisterCloseResult = Result<(), String>; -/// The arguments for [ws_open]. +/// The arguments for [ws_open](crate::ws_open). #[derive(CandidType, Clone, Deserialize, Serialize, Eq, PartialEq, Debug)] pub struct CanisterWsOpenArguments { pub(crate) client_nonce: u64, pub(crate) gateway_principal: GatewayPrincipal, } -/// The arguments for [ws_close]. +/// The arguments for [ws_close](crate::ws_close). #[derive(CandidType, Clone, Deserialize, Serialize, Eq, PartialEq, Debug)] pub struct CanisterWsCloseArguments { pub(crate) client_key: ClientKey, } -/// The arguments for [ws_message]. +/// The arguments for [ws_message](crate::ws_message). #[derive(CandidType, Clone, Deserialize, Serialize, Eq, PartialEq, Debug)] pub struct CanisterWsMessageArguments { pub(crate) msg: WebsocketMessage, } -/// The arguments for [ws_get_messages]. +/// The arguments for [ws_get_messages](crate::ws_get_messages). #[derive(CandidType, Clone, Deserialize, Serialize, Eq, PartialEq, Debug)] pub struct CanisterWsGetMessagesArguments { pub(crate) nonce: u64, From 95e31186566f831573e4e7ebcbf91fe949344309 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 8 Dec 2023 11:21:33 +0100 Subject: [PATCH 07/18] refactor: rename ws_send to send to be consistent with the newly introduced `close` method --- src/ic-websocket-cdk/src/lib.rs | 22 +++++++++++----- src/ic-websocket-cdk/src/state.rs | 4 +-- .../integration_tests/c_ws_get_messages.rs | 10 +++---- .../{e_ws_send.rs => e_send.rs} | 12 ++++----- .../integration_tests/h_multiple_gateways.rs | 8 +++--- .../src/tests/integration_tests/mod.rs | 2 +- .../tests/integration_tests/utils/actor.rs | 26 +++++++++---------- src/ic-websocket-cdk/src/types.rs | 6 +++-- src/test_canister/src/lib.rs | 10 +++---- src/test_canister/test_canister.did | 4 +-- 10 files changed, 57 insertions(+), 47 deletions(-) rename src/ic-websocket-cdk/src/tests/integration_tests/{e_ws_send.rs => e_send.rs} (81%) diff --git a/src/ic-websocket-cdk/src/lib.rs b/src/ic-websocket-cdk/src/lib.rs index ac88de5..16f1142 100644 --- a/src/ic-websocket-cdk/src/lib.rs +++ b/src/ic-websocket-cdk/src/lib.rs @@ -13,13 +13,14 @@ mod utils; use state::*; use timers::*; +#[allow(deprecated)] +pub use types::CanisterWsSendResult; use types::*; pub use types::{ - CanisterCloseResult, CanisterWsCloseArguments, CanisterWsCloseResult, + CanisterCloseResult, CanisterSendResult, CanisterWsCloseArguments, CanisterWsCloseResult, CanisterWsGetMessagesArguments, CanisterWsGetMessagesResult, CanisterWsMessageArguments, - CanisterWsMessageResult, CanisterWsOpenArguments, CanisterWsOpenResult, CanisterWsSendResult, - ClientPrincipal, OnCloseCallbackArgs, OnMessageCallbackArgs, OnOpenCallbackArgs, WsHandlers, - WsInitParams, + CanisterWsMessageResult, CanisterWsOpenArguments, CanisterWsOpenResult, ClientPrincipal, + OnCloseCallbackArgs, OnMessageCallbackArgs, OnOpenCallbackArgs, WsHandlers, WsInitParams, }; /// The label used when constructing the certification tree. @@ -235,7 +236,7 @@ pub fn ws_get_messages(args: CanisterWsGetMessagesArguments) -> CanisterWsGetMes /// This example is the serialize equivalent of the [OnMessageCallbackArgs's example](struct.OnMessageCallbackArgs.html#example) deserialize one. /// ```rust /// use candid::{encode_one, CandidType, Principal}; -/// use ic_websocket_cdk::ws_send; +/// use ic_websocket_cdk::send; /// use serde::Deserialize; /// /// #[derive(CandidType, Deserialize)] @@ -251,13 +252,20 @@ pub fn ws_get_messages(args: CanisterWsGetMessagesArguments) -> CanisterWsGetMes /// }; /// /// let msg_bytes = encode_one(&my_message).unwrap(); -/// ws_send(my_client_principal, msg_bytes); +/// send(my_client_principal, msg_bytes); /// ``` -pub fn ws_send(client_principal: ClientPrincipal, msg_bytes: Vec) -> CanisterWsSendResult { +pub fn send(client_principal: ClientPrincipal, msg_bytes: Vec) -> CanisterSendResult { let client_key = get_client_key_from_principal(&client_principal)?; _ws_send(&client_key, msg_bytes, false) } +#[deprecated(since = "0.3.2", note = "use `ic_websocket_cdk::send` instead")] +#[allow(deprecated)] +/// Deprecated: use [send] instead. +pub fn ws_send(client_principal: ClientPrincipal, msg_bytes: Vec) -> CanisterWsSendResult { + send(client_principal, msg_bytes) +} + /// Closes the connection with the client. pub fn close(client_principal: ClientPrincipal) -> CanisterCloseResult { let client_key = get_client_key_from_principal(&client_principal)?; diff --git a/src/ic-websocket-cdk/src/state.rs b/src/ic-websocket-cdk/src/state.rs index 846d516..2ba4d0f 100644 --- a/src/ic-websocket-cdk/src/state.rs +++ b/src/ic-websocket-cdk/src/state.rs @@ -437,7 +437,7 @@ pub(crate) fn push_message_in_gateway_queue( ) -> Result<(), String> { REGISTERED_GATEWAYS.with(|map| { // messages in the queue are inserted with contiguous and increasing nonces - // (from beginning to end of the queue) as ws_send is called sequentially, the nonce + // (from beginning to end of the queue) as `send` is called sequentially, the nonce // is incremented by one in each call, and the message is pushed at the end of the queue map.borrow_mut() .get_mut(gateway_principal) @@ -555,7 +555,7 @@ pub(crate) fn _ws_send( client_key: &ClientKey, msg_bytes: Vec, is_service_message: bool, -) -> CanisterWsSendResult { +) -> CanisterSendResult { // get the registered client if it exists let registered_client = get_registered_client(client_key)?; diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs b/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs index 314486a..9a2c1a6 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs @@ -8,8 +8,8 @@ use crate::{ use super::utils::{ actor::{ - ws_get_messages::call_ws_get_messages_with_panic, - ws_open::call_ws_open_for_client_key_with_panic, ws_send::call_ws_send_with_panic, + send::call_send_with_panic, ws_get_messages::call_ws_get_messages_with_panic, + ws_open::call_ws_open_for_client_key_with_panic, }, clients::{generate_random_client_key, CLIENT_1_KEY, GATEWAY_1}, messages::{get_next_polling_nonce_from_messages, verify_messages, AppMessage}, @@ -70,7 +70,7 @@ proptest! { text: format!("test{}", i), }) .collect(); - call_ws_send_with_panic( + call_send_with_panic( &client_1_key.client_principal, messages_to_send.clone(), ); @@ -131,7 +131,7 @@ proptest! { text: format!("test{}", i), }) .collect(); - call_ws_send_with_panic(&client_1_key.client_principal, messages_to_send.clone()); + call_send_with_panic(&client_1_key.client_principal, messages_to_send.clone()); let mut next_polling_nonce = 0; let mut expected_sequence_number = 1; // `1` because the seq number is incremented before sending on the canister @@ -179,7 +179,7 @@ proptest! { text: format!("test{}", i), }) .collect(); - call_ws_send_with_panic(&client_1_key.client_principal, messages_to_send.clone()); + call_send_with_panic(&client_1_key.client_principal, messages_to_send.clone()); // advance canister time because the messages are deleted only // if they're older than the send ack interval; diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/e_ws_send.rs b/src/ic-websocket-cdk/src/tests/integration_tests/e_send.rs similarity index 81% rename from src/ic-websocket-cdk/src/tests/integration_tests/e_ws_send.rs rename to src/ic-websocket-cdk/src/tests/integration_tests/e_send.rs index b40c7b0..1e59a3f 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/e_ws_send.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/e_send.rs @@ -1,9 +1,9 @@ use std::ops::Deref; -use crate::{errors::WsError, CanisterWsSendResult}; +use crate::{errors::WsError, CanisterSendResult}; use super::utils::{ - actor::{ws_open::call_ws_open_for_client_key_with_panic, ws_send::call_ws_send}, + actor::{send::call_send, ws_open::call_ws_open_for_client_key_with_panic}, clients::{CLIENT_1_KEY, CLIENT_2_KEY}, messages::AppMessage, test_env::get_test_env, @@ -18,7 +18,7 @@ fn test_1_fails_if_sending_a_message_to_a_non_registered_client() { // finally, we can start testing let client_2_principal = &CLIENT_2_KEY.client_principal; - let res = call_ws_send( + let res = call_send( client_2_principal, vec![AppMessage { text: String::from("test"), @@ -26,7 +26,7 @@ fn test_1_fails_if_sending_a_message_to_a_non_registered_client() { ); assert_eq!( res, - CanisterWsSendResult::Err( + CanisterSendResult::Err( WsError::ClientPrincipalNotConnected { client_principal: client_2_principal } @@ -42,11 +42,11 @@ fn test_2_should_send_a_message_to_a_registered_client() { // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(CLIENT_1_KEY.deref()); - let res = call_ws_send( + let res = call_send( &CLIENT_1_KEY.client_principal, vec![AppMessage { text: String::from("test"), }], ); - assert_eq!(res, CanisterWsSendResult::Ok(())); + assert_eq!(res, CanisterSendResult::Ok(())); } diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/h_multiple_gateways.rs b/src/ic-websocket-cdk/src/tests/integration_tests/h_multiple_gateways.rs index 9b675a6..6926554 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/h_multiple_gateways.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/h_multiple_gateways.rs @@ -7,9 +7,9 @@ use crate::{ use super::utils::{ actor::{ - ws_close::call_ws_close_with_panic, ws_get_messages::call_ws_get_messages_with_panic, + send::call_send_with_panic, ws_close::call_ws_close_with_panic, + ws_get_messages::call_ws_get_messages_with_panic, ws_open::call_ws_open_for_client_key_and_gateway_with_panic, - ws_send::call_ws_send_with_panic, }, messages::{get_service_message_content_from_canister_message, verify_messages, AppMessage}, test_env::get_test_env, @@ -34,7 +34,7 @@ proptest! { text: format!("test{}", i), }) .collect(); - call_ws_send_with_panic( + call_send_with_panic( &client_key.client_principal, messages_to_send.clone(), ); @@ -108,7 +108,7 @@ proptest! { }) .collect(); // simulate canister sending other messages to client - call_ws_send_with_panic( + call_send_with_panic( &client_key.client_principal, messages_to_send.clone(), ); diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/mod.rs b/src/ic-websocket-cdk/src/tests/integration_tests/mod.rs index a261e8d..ed6481f 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/mod.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/mod.rs @@ -11,7 +11,7 @@ mod a_ws_open; mod b_ws_message; mod c_ws_get_messages; mod d_ws_close; -mod e_ws_send; +mod e_send; mod f_close; mod g_messages_acknowledgement; mod h_multiple_gateways; diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs b/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs index 05fff54..be1eee2 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs @@ -181,30 +181,30 @@ pub mod ws_get_messages { } } -pub mod ws_send { - use crate::{CanisterWsSendResult, ClientPrincipal}; +pub mod send { + use crate::{CanisterSendResult, ClientPrincipal}; use candid::encode_args; use super::*; /// (`ClientPrincipal`, `Vec>`) - type WsSendArguments = (ClientPrincipal, Vec>); + type SendArguments = (ClientPrincipal, Vec>); /// # Panics /// if the call returns a [WasmResult::Reject]. - pub(crate) fn call_ws_send( + pub(crate) fn call_send( send_to_principal: &ClientPrincipal, messages: Vec, - ) -> CanisterWsSendResult { + ) -> CanisterSendResult { let messages: Vec> = messages.iter().map(|m| encode_one(m).unwrap()).collect(); - let args: WsSendArguments = (send_to_principal.clone(), messages); + let args: SendArguments = (send_to_principal.clone(), messages); let canister_id = get_test_env().canister_id; let res = get_test_env() .pic .update_call( canister_id, Principal::anonymous(), - "ws_send", + "send", encode_args(args).unwrap(), ) .expect("Failed to call canister"); @@ -214,17 +214,17 @@ pub mod ws_send { } } - /// Same as [call_ws_send]. + /// Same as [call_send]. /// /// # Panics - /// If [call_ws_send] panics or if the call returns an error variant. - pub(crate) fn call_ws_send_with_panic( + /// If [call_send] panics or if the call returns an error variant. + pub(crate) fn call_send_with_panic( send_to_principal: &ClientPrincipal, messages: Vec, ) { - match call_ws_send(send_to_principal, messages) { - CanisterWsSendResult::Ok(_) => {}, - CanisterWsSendResult::Err(err) => panic!("failed ws_send: {:?}", err), + match call_send(send_to_principal, messages) { + CanisterSendResult::Ok(_) => {}, + CanisterSendResult::Err(err) => panic!("failed send: {:?}", err), } } } diff --git a/src/ic-websocket-cdk/src/types.rs b/src/ic-websocket-cdk/src/types.rs index aee198e..1abad3a 100644 --- a/src/ic-websocket-cdk/src/types.rs +++ b/src/ic-websocket-cdk/src/types.rs @@ -41,7 +41,9 @@ pub type CanisterWsCloseResult = Result<(), String>; pub type CanisterWsMessageResult = Result<(), String>; /// The result of [ws_get_messages](crate::ws_get_messages). pub type CanisterWsGetMessagesResult = Result; -/// The result of [ws_send](crate::ws_send). +/// The result of [send](crate::send). +pub type CanisterSendResult = Result<(), String>; +#[deprecated(since = "0.3.2", note = "use `CanisterSendResult` instead")] pub type CanisterWsSendResult = Result<(), String>; /// The result of [close](crate::close). pub type CanisterCloseResult = Result<(), String>; @@ -314,7 +316,7 @@ type OnOpenCallback = fn(OnOpenCallbackArgs); /// To deserialize the message, use [candid::decode_one]. /// /// # Example -/// This example is the deserialize equivalent of the [ws_send's example](fn.ws_send.html#example) serialize one. +/// This example is the deserialize equivalent of the [send's example](fn.send.html#example) serialize one. /// ```rust /// use candid::{decode_one, CandidType}; /// use ic_websocket_cdk::OnMessageCallbackArgs; diff --git a/src/test_canister/src/lib.rs b/src/test_canister/src/lib.rs index 622c1a3..a8e90c2 100644 --- a/src/test_canister/src/lib.rs +++ b/src/test_canister/src/lib.rs @@ -2,10 +2,10 @@ use ic_cdk_macros::*; use canister::{on_close, on_message, on_open, AppMessage}; use ic_websocket_cdk::{ - CanisterCloseResult, CanisterWsCloseArguments, CanisterWsCloseResult, + CanisterCloseResult, CanisterSendResult, CanisterWsCloseArguments, CanisterWsCloseResult, CanisterWsGetMessagesArguments, CanisterWsGetMessagesResult, CanisterWsMessageArguments, - CanisterWsMessageResult, CanisterWsOpenArguments, CanisterWsOpenResult, CanisterWsSendResult, - ClientPrincipal, WsHandlers, WsInitParams, + CanisterWsMessageResult, CanisterWsOpenArguments, CanisterWsOpenResult, ClientPrincipal, + WsHandlers, WsInitParams, }; mod canister; @@ -75,9 +75,9 @@ fn ws_get_messages(args: CanisterWsGetMessagesArguments) -> CanisterWsGetMessage //// Debug/tests methods // send a message to the client, usually called by the canister itself #[update] -fn ws_send(client_principal: ClientPrincipal, messages: Vec>) -> CanisterWsSendResult { +fn send(client_principal: ClientPrincipal, messages: Vec>) -> CanisterSendResult { for msg_bytes in messages { - match ic_websocket_cdk::ws_send(client_principal, msg_bytes) { + match ic_websocket_cdk::send(client_principal, msg_bytes) { Ok(_) => {}, Err(e) => return Err(e), } diff --git a/src/test_canister/test_canister.did b/src/test_canister/test_canister.did index 3044dc6..39b9e82 100644 --- a/src/test_canister/test_canister.did +++ b/src/test_canister/test_canister.did @@ -1,6 +1,6 @@ import "../src/ic-websocket-cdk/ws_types.did"; -type CanisterWsSendResult = variant { +type CanisterSendResult = variant { Ok : null; Err : text; }; @@ -21,6 +21,6 @@ service : (text, nat64, nat64, nat64) -> { "ws_get_messages" : (CanisterWsGetMessagesArguments) -> (CanisterWsGetMessagesResult) query; // methods used just for debugging/testing - "ws_send" : (ClientPrincipal, vec blob) -> (CanisterWsSendResult); + "send" : (ClientPrincipal, vec blob) -> (CanisterSendResult); "close" : (ClientPrincipal) -> (CanisterCloseResult); }; From 323fb83bead5a30aed313a16841e6e2b0f6152ba Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 8 Dec 2023 11:40:19 +0100 Subject: [PATCH 08/18] docs: caveats --- src/ic-websocket-cdk/src/lib.rs | 2 ++ src/ic-websocket-cdk/src/types.rs | 7 +++++++ 2 files changed, 9 insertions(+) diff --git a/src/ic-websocket-cdk/src/lib.rs b/src/ic-websocket-cdk/src/lib.rs index 16f1142..0ef32d4 100644 --- a/src/ic-websocket-cdk/src/lib.rs +++ b/src/ic-websocket-cdk/src/lib.rs @@ -267,6 +267,8 @@ pub fn ws_send(client_principal: ClientPrincipal, msg_bytes: Vec) -> Caniste } /// Closes the connection with the client. +/// +/// This function **must not** be called in the `on_close` callback. pub fn close(client_principal: ClientPrincipal) -> CanisterCloseResult { let client_key = get_client_key_from_principal(&client_principal)?; diff --git a/src/ic-websocket-cdk/src/types.rs b/src/ic-websocket-cdk/src/types.rs index 1abad3a..966f1b9 100644 --- a/src/ic-websocket-cdk/src/types.rs +++ b/src/ic-websocket-cdk/src/types.rs @@ -350,9 +350,16 @@ pub struct OnCloseCallbackArgs { /// Handler initialized by the canister /// and triggered by the CDK once the WS Gateway closes the IC WebSocket connection /// for that client. +/// +/// Make sure you **don't** call the [close](crate::close) function in this callback. type OnCloseCallback = fn(OnCloseCallbackArgs); /// Handlers initialized by the canister and triggered by the CDK. +/// +/// **Note**: if the callbacks that you define here trap for some reason, +/// the CDK will disconnect the client with principal `args.client_principal`. +/// However, the client **won't** be notified +/// until at least the next time it will try to send a message to the canister. #[derive(Clone, Debug, Default, PartialEq)] pub struct WsHandlers { pub on_open: Option, From fdbe450fc91b40c4a901be122fb00f01b5908695 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 8 Dec 2023 12:27:21 +0100 Subject: [PATCH 09/18] feat: MAX_ALLOWED_CONNECTION_LATENCY const --- src/ic-websocket-cdk/src/lib.rs | 6 +++- .../integration_tests/c_ws_get_messages.rs | 5 ++- .../g_messages_acknowledgement.rs | 22 +++++-------- .../tests/integration_tests/utils/test_env.rs | 19 ++--------- .../src/tests/unit_tests/mod.rs | 32 ++++++------------ src/ic-websocket-cdk/src/timers.rs | 20 ++++++----- src/ic-websocket-cdk/src/types.rs | 33 +++++++++---------- src/test_canister/src/lib.rs | 19 ++--------- src/test_canister/test_canister.did | 2 +- 9 files changed, 59 insertions(+), 99 deletions(-) diff --git a/src/ic-websocket-cdk/src/lib.rs b/src/ic-websocket-cdk/src/lib.rs index 0ef32d4..561cefd 100644 --- a/src/ic-websocket-cdk/src/lib.rs +++ b/src/ic-websocket-cdk/src/lib.rs @@ -30,8 +30,12 @@ const LABEL_WEBSOCKET: &[u8] = b"websocket"; const DEFAULT_MAX_NUMBER_OF_RETURNED_MESSAGES: usize = 50; /// The default interval at which to send acknowledgements to the client. const DEFAULT_SEND_ACK_INTERVAL_MS: u64 = 300_000; // 5 minutes +/// The maximum latency allowed between the client and the canister. +const MAX_ALLOWED_CONNECTION_LATENCY_MS: u64 = 30_000; // 30 seconds /// The default timeout to wait for the client to send a keep alive after receiving an acknowledgement. -const DEFAULT_CLIENT_KEEP_ALIVE_TIMEOUT_MS: u64 = 60_000; // 1 minute +const CLIENT_KEEP_ALIVE_TIMEOUT_MS: u64 = 2 * MAX_ALLOWED_CONNECTION_LATENCY_MS; +/// Same as [CLIENT_KEEP_ALIVE_TIMEOUT_MS], but in nanoseconds. +const CLIENT_KEEP_ALIVE_TIMEOUT_NS: u64 = CLIENT_KEEP_ALIVE_TIMEOUT_MS * 1_000_000; /// The initial nonce for outgoing messages. const INITIAL_OUTGOING_MESSAGE_NONCE: u64 = 0; diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs b/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs index 9a2c1a6..314d48d 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs @@ -14,8 +14,8 @@ use super::utils::{ clients::{generate_random_client_key, CLIENT_1_KEY, GATEWAY_1}, messages::{get_next_polling_nonce_from_messages, verify_messages, AppMessage}, test_env::{ - get_test_env, DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS, - DEFAULT_TEST_MAX_NUMBER_OF_RETURNED_MESSAGES, DEFAULT_TEST_SEND_ACK_INTERVAL_MS, + get_test_env, DEFAULT_TEST_MAX_NUMBER_OF_RETURNED_MESSAGES, + DEFAULT_TEST_SEND_ACK_INTERVAL_MS, }, }; @@ -168,7 +168,6 @@ proptest! { get_test_env().reset_canister( 1_000, // avoid the queue size limit DEFAULT_TEST_SEND_ACK_INTERVAL_MS, - DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS, ); // second, register client 1 let client_1_key = CLIENT_1_KEY.deref(); diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/g_messages_acknowledgement.rs b/src/ic-websocket-cdk/src/tests/integration_tests/g_messages_acknowledgement.rs index c4c79db..24084a3 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/g_messages_acknowledgement.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/g_messages_acknowledgement.rs @@ -8,7 +8,7 @@ use crate::{ types::CloseMessageReason, CanisterOutputCertifiedMessages, CanisterWsCloseArguments, CanisterWsCloseResult, CanisterWsGetMessagesArguments, CanisterWsMessageArguments, ClientKeepAliveMessageContent, - WebsocketServiceMessageContent, + WebsocketServiceMessageContent, CLIENT_KEEP_ALIVE_TIMEOUT_MS, }; use super::utils::{ @@ -22,9 +22,7 @@ use super::utils::{ create_websocket_message, decode_websocket_service_message_content, encode_websocket_service_message_content, get_websocket_message_from_canister_message, }, - test_env::{ - get_test_env, DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS, DEFAULT_TEST_SEND_ACK_INTERVAL_MS, - }, + test_env::{get_test_env, DEFAULT_TEST_SEND_ACK_INTERVAL_MS}, }; #[test] @@ -72,7 +70,7 @@ fn test_2_client_is_removed_if_keep_alive_timeout_is_reached() { helpers::check_ack_message_result(&msgs, client_1_key, 0, 2); // advance the canister time to make sure the keep alive timeout expires - get_test_env().advance_canister_time_ms(DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS); + get_test_env().advance_canister_time_ms(CLIENT_KEEP_ALIVE_TIMEOUT_MS); // check if the gateway put the close message in the queue let msgs = call_ws_get_messages_with_panic( @@ -137,7 +135,7 @@ fn test_3_client_is_not_removed_if_it_sends_a_keep_alive_before_timeout() { }, ); // advance the canister time to make sure the keep alive timeout expires and the canister checks the keep alive - get_test_env().advance_canister_time_ms(DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS); + get_test_env().advance_canister_time_ms(CLIENT_KEEP_ALIVE_TIMEOUT_MS); // send a message to the canister to see the sequence number increasing in the ack message // and be sure that the client has not been removed call_ws_message_with_panic( @@ -147,9 +145,8 @@ fn test_3_client_is_not_removed_if_it_sends_a_keep_alive_before_timeout() { }, ); // wait for the canister to send the next ack - get_test_env().advance_canister_time_ms( - DEFAULT_TEST_SEND_ACK_INTERVAL_MS - DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS, - ); + get_test_env() + .advance_canister_time_ms(DEFAULT_TEST_SEND_ACK_INTERVAL_MS - CLIENT_KEEP_ALIVE_TIMEOUT_MS); let res = call_ws_get_messages_with_panic( GATEWAY_1.deref(), CanisterWsGetMessagesArguments { nonce: 2 }, // skip the service open message and the fist ack message @@ -183,11 +180,10 @@ fn test_4_client_is_not_removed_if_it_connects_while_canister_is_waiting_for_kee ); // wait for the keep alive timeout to expire - get_test_env().advance_canister_time_ms(DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS); + get_test_env().advance_canister_time_ms(CLIENT_KEEP_ALIVE_TIMEOUT_MS); // wait for the canister to send the next ack - get_test_env().advance_canister_time_ms( - DEFAULT_TEST_SEND_ACK_INTERVAL_MS - DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS, - ); + get_test_env() + .advance_canister_time_ms(DEFAULT_TEST_SEND_ACK_INTERVAL_MS - CLIENT_KEEP_ALIVE_TIMEOUT_MS); let res = call_ws_get_messages_with_panic( GATEWAY_1.deref(), diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs b/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs index d255ec6..bb5c2b8 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs @@ -19,12 +19,6 @@ pub const DEFAULT_TEST_MAX_NUMBER_OF_RETURNED_MESSAGES: u64 = 50; /// Value: `300_000` = 5 minutes pub const DEFAULT_TEST_SEND_ACK_INTERVAL_MS: u64 = 300_000; -/// The interval between keep alive checks in the canister. -/// Set to a high value to make sure the canister doesn't reset the client while testing other functions. -/// -/// Value: `120_000` (2 minutes) -pub const DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS: u64 = 120_000; - lazy_static! { pub static ref TEST_ENV: Mutex = Mutex::new(TestEnv::new()); } @@ -41,8 +35,8 @@ pub struct TestEnv { root_ic_key: Vec, } -/// (`max_number_or_returned_messages`, `send_ack_interval_ms`, `send_ack_timeout_ms`) -type CanisterInitArgs = (u64, u64, u64); +/// (`max_number_or_returned_messages`, `send_ack_interval_ms`) +type CanisterInitArgs = (u64, u64); impl TestEnv { pub fn new() -> Self { @@ -62,7 +56,6 @@ impl TestEnv { let arguments: CanisterInitArgs = ( DEFAULT_TEST_MAX_NUMBER_OF_RETURNED_MESSAGES, DEFAULT_TEST_SEND_ACK_INTERVAL_MS, - DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS, ); pic.install_canister( canister_id, @@ -86,13 +79,8 @@ impl TestEnv { &mut self, max_number_or_returned_messages: u64, send_ack_interval_ms: u64, - keep_alive_delay_ms: u64, ) { - let arguments: CanisterInitArgs = ( - max_number_or_returned_messages, - send_ack_interval_ms, - keep_alive_delay_ms, - ); + let arguments: CanisterInitArgs = (max_number_or_returned_messages, send_ack_interval_ms); let res = self.pic.reinstall_canister( self.canister_id, self.wasm_module.to_owned(), @@ -115,7 +103,6 @@ impl TestEnv { self.reset_canister( DEFAULT_TEST_MAX_NUMBER_OF_RETURNED_MESSAGES, DEFAULT_TEST_SEND_ACK_INTERVAL_MS, - DEFAULT_TEST_KEEP_ALIVE_TIMEOUT_MS, ); } diff --git a/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs b/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs index 8b87481..337644f 100644 --- a/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs +++ b/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs @@ -119,38 +119,25 @@ fn test_ws_init_params() { DEFAULT_MAX_NUMBER_OF_RETURNED_MESSAGES ); assert_eq!(params.send_ack_interval_ms, DEFAULT_SEND_ACK_INTERVAL_MS); - assert_eq!( - params.keep_alive_timeout_ms, - DEFAULT_CLIENT_KEEP_ALIVE_TIMEOUT_MS - ); let params = WsInitParams::new(handlers.clone()) .with_max_number_of_returned_messages(5) - .with_send_ack_interval_ms(10) - .with_keep_alive_timeout_ms(2); + .with_send_ack_interval_ms(120_000); assert_eq!(params.max_number_of_returned_messages, 5); - assert_eq!(params.send_ack_interval_ms, 10); - assert_eq!(params.keep_alive_timeout_ms, 2); + assert_eq!(params.send_ack_interval_ms, 120_000); } #[test] -#[should_panic = "send_ack_interval_ms must be greater than keep_alive_timeout_ms"] +#[should_panic = "send_ack_interval_ms must be greater than CLIENT_KEEP_ALIVE_TIMEOUT_MS"] fn test_ws_init_params_keep_alive_greater() { - let params = WsInitParams::new(WsHandlers::default()) - .with_send_ack_interval_ms(5) - .with_keep_alive_timeout_ms(10); - - params.check_validity(); + WsInitParams::new(WsHandlers::default()).with_send_ack_interval_ms(5); } #[test] -#[should_panic = "send_ack_interval_ms must be greater than keep_alive_timeout_ms"] +#[should_panic = "send_ack_interval_ms must be greater than CLIENT_KEEP_ALIVE_TIMEOUT_MS"] fn test_ws_init_params_keep_alive_equal() { - let params = WsInitParams::new(WsHandlers::default()) - .with_send_ack_interval_ms(10) - .with_keep_alive_timeout_ms(10); - - params.check_validity(); + WsInitParams::new(WsHandlers::default()) + .with_send_ack_interval_ms(CLIENT_KEEP_ALIVE_TIMEOUT_MS); } #[test] @@ -901,13 +888,14 @@ proptest! { } #[test] - fn test_delete_old_messages_for_gateway_queue( + fn test_delete_old_messages_for_gateway( gateway_principal in any::().prop_map(|_| common::generate_random_principal()), old_messages_count in 0..100usize, new_messages_count in 0..100usize, - test_ack_interval_ms in (1..100u64), + test_ack_interval_delta_ms in (1..100u64), ) { // Set up + let test_ack_interval_ms = CLIENT_KEEP_ALIVE_TIMEOUT_MS + test_ack_interval_delta_ms; PARAMS.with(|p| *p.borrow_mut() = WsInitParams::new(WsHandlers::default()) .with_send_ack_interval_ms(test_ack_interval_ms) diff --git a/src/ic-websocket-cdk/src/timers.rs b/src/ic-websocket-cdk/src/timers.rs index 0851515..29d4b3d 100644 --- a/src/ic-websocket-cdk/src/timers.rs +++ b/src/ic-websocket-cdk/src/timers.rs @@ -4,10 +4,10 @@ use std::cell::RefCell; use std::rc::Rc; use std::time::Duration; -use crate::custom_print; use crate::state::*; use crate::types::*; use crate::utils::*; +use crate::{custom_print, CLIENT_KEEP_ALIVE_TIMEOUT_MS, CLIENT_KEEP_ALIVE_TIMEOUT_NS}; thread_local! { /// The acknowledgement active timer. @@ -61,10 +61,12 @@ pub(crate) fn schedule_send_ack_to_clients() { /// /// The timer callback is [check_keep_alive_timer_callback]. fn schedule_check_keep_alive() { - let keep_alive_timeout_ms = get_params().keep_alive_timeout_ms; - let timer_id = set_timer(Duration::from_millis(keep_alive_timeout_ms), move || { - check_keep_alive_timer_callback(keep_alive_timeout_ms); - }); + let timer_id = set_timer( + Duration::from_millis(CLIENT_KEEP_ALIVE_TIMEOUT_MS), + move || { + check_keep_alive_timer_callback(); + }, + ); put_keep_alive_timer_id(timer_id); } @@ -110,7 +112,7 @@ fn send_ack_to_clients_timer_callback() { /// Checks if the clients for which we are waiting for keep alive have sent a keep alive message. /// If a client has not sent a keep alive message, it is removed from the connected clients. -fn check_keep_alive_timer_callback(keep_alive_timeout_ms: u64) { +fn check_keep_alive_timer_callback() { let client_keys_to_remove: Vec = CLIENTS_WAITING_FOR_KEEP_ALIVE .with(Rc::clone) .borrow() @@ -121,7 +123,7 @@ fn check_keep_alive_timer_callback(keep_alive_timeout_ms: u64) { REGISTERED_CLIENTS.with(Rc::clone).borrow().get(client_key) { let last_keep_alive = client_metadata.get_last_keep_alive_timestamp(); - if get_current_time() - last_keep_alive > (keep_alive_timeout_ms * 1_000_000) { + if get_current_time() - last_keep_alive > CLIENT_KEEP_ALIVE_TIMEOUT_NS { Some(client_key.to_owned()) } else { None @@ -136,9 +138,9 @@ fn check_keep_alive_timer_callback(keep_alive_timeout_ms: u64) { remove_client(&client_key, Some(CloseMessageReason::KeepAliveTimeout)); custom_print!( - "[check-keep-alive-timer-cb]: Client {} has not sent a keep alive message in the last {} ms and has been removed", + "[check-keep-alive-timer-cb]: Client {} has not sent a keep alive message in the last {}ms and has been removed", client_key, - keep_alive_timeout_ms + CLIENT_KEEP_ALIVE_TIMEOUT_MS ); } diff --git a/src/ic-websocket-cdk/src/types.rs b/src/ic-websocket-cdk/src/types.rs index 966f1b9..0a547fc 100644 --- a/src/ic-websocket-cdk/src/types.rs +++ b/src/ic-websocket-cdk/src/types.rs @@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize}; use serde_cbor::Serializer; use crate::{ - custom_trap, errors::WsError, utils::get_current_time, DEFAULT_CLIENT_KEEP_ALIVE_TIMEOUT_MS, + custom_trap, errors::WsError, utils::get_current_time, CLIENT_KEEP_ALIVE_TIMEOUT_MS, DEFAULT_MAX_NUMBER_OF_RETURNED_MESSAGES, DEFAULT_SEND_ACK_INTERVAL_MS, INITIAL_OUTGOING_MESSAGE_NONCE, }; @@ -397,21 +397,16 @@ pub struct WsInitParams { /// The callback handlers for the WebSocket. pub handlers: WsHandlers, /// The maximum number of messages to be returned in a polling iteration. + /// /// Defaults to `50`. pub max_number_of_returned_messages: usize, /// The interval at which to send an acknowledgement message to the client, /// so that the client knows that all the messages it sent have been received by the canister (in milliseconds). /// - /// Must be greater than `keep_alive_timeout_ms`. + /// Must be greater than [`CLIENT_KEEP_ALIVE_TIMEOUT_MS`] (1 minute). /// /// Defaults to `300_000` (5 minutes). pub send_ack_interval_ms: u64, - /// The delay to wait for the client to send a keep alive after receiving an acknowledgement (in milliseconds). - /// - /// Must be lower than `send_ack_interval_ms`. - /// - /// Defaults to `60_000` (1 minute). - pub keep_alive_timeout_ms: u64, } impl WsInitParams { @@ -428,13 +423,13 @@ impl WsInitParams { } /// Checks the validity of the timer parameters. - /// `send_ack_interval_ms` must be greater than `keep_alive_timeout_ms`. + /// `send_ack_interval_ms` must be greater than [`CLIENT_KEEP_ALIVE_TIMEOUT_MS`]. /// /// # Traps - /// If `send_ack_interval_ms` <= `keep_alive_timeout_ms`. + /// If `send_ack_interval_ms` <= [`CLIENT_KEEP_ALIVE_TIMEOUT_MS`]. pub(crate) fn check_validity(&self) { - if self.keep_alive_timeout_ms >= self.send_ack_interval_ms { - custom_trap!("send_ack_interval_ms must be greater than keep_alive_timeout_ms"); + if self.send_ack_interval_ms <= CLIENT_KEEP_ALIVE_TIMEOUT_MS { + custom_trap!("send_ack_interval_ms must be greater than CLIENT_KEEP_ALIVE_TIMEOUT_MS"); } } @@ -446,13 +441,16 @@ impl WsInitParams { self } + /// Sets the interval (in milliseconds) at which to send an acknowledgement message + /// to the connected clients. + /// + /// Must be greater than [`CLIENT_KEEP_ALIVE_TIMEOUT_MS`] (1 minute). + /// + /// # Traps + /// If `send_ack_interval_ms` <= [`CLIENT_KEEP_ALIVE_TIMEOUT_MS`]. See [WsInitParams::check_validity]. pub fn with_send_ack_interval_ms(mut self, send_ack_interval_ms: u64) -> Self { self.send_ack_interval_ms = send_ack_interval_ms; - self - } - - pub fn with_keep_alive_timeout_ms(mut self, keep_alive_timeout_ms: u64) -> Self { - self.keep_alive_timeout_ms = keep_alive_timeout_ms; + self.check_validity(); self } } @@ -463,7 +461,6 @@ impl Default for WsInitParams { handlers: WsHandlers::default(), max_number_of_returned_messages: DEFAULT_MAX_NUMBER_OF_RETURNED_MESSAGES, send_ack_interval_ms: DEFAULT_SEND_ACK_INTERVAL_MS, - keep_alive_timeout_ms: DEFAULT_CLIENT_KEEP_ALIVE_TIMEOUT_MS, } } } diff --git a/src/test_canister/src/lib.rs b/src/test_canister/src/lib.rs index a8e90c2..a343590 100644 --- a/src/test_canister/src/lib.rs +++ b/src/test_canister/src/lib.rs @@ -11,11 +11,7 @@ use ic_websocket_cdk::{ mod canister; #[init] -fn init( - max_number_of_returned_messages: usize, - send_ack_interval_ms: u64, - keep_alive_timeout_ms: u64, -) { +fn init(max_number_of_returned_messages: usize, send_ack_interval_ms: u64) { let handlers = WsHandlers { on_open: Some(on_open), on_message: Some(on_message), @@ -26,23 +22,14 @@ fn init( handlers, max_number_of_returned_messages, send_ack_interval_ms, - keep_alive_timeout_ms, }; ic_websocket_cdk::init(params) } #[post_upgrade] -fn post_upgrade( - max_number_of_returned_messages: usize, - send_ack_interval_ms: u64, - keep_alive_timeout_ms: u64, -) { - init( - max_number_of_returned_messages, - send_ack_interval_ms, - keep_alive_timeout_ms, - ); +fn post_upgrade(max_number_of_returned_messages: usize, send_ack_interval_ms: u64) { + init(max_number_of_returned_messages, send_ack_interval_ms); } // method called by the WS Gateway after receiving FirstMessage from the client diff --git a/src/test_canister/test_canister.did b/src/test_canister/test_canister.did index 39b9e82..2a04b7d 100644 --- a/src/test_canister/test_canister.did +++ b/src/test_canister/test_canister.did @@ -14,7 +14,7 @@ type AppMessage = record { text : text; }; -service : (text, nat64, nat64, nat64) -> { +service : (nat64, nat64) -> { "ws_open" : (CanisterWsOpenArguments) -> (CanisterWsOpenResult); "ws_close" : (CanisterWsCloseArguments) -> (CanisterWsCloseResult); "ws_message" : (CanisterWsMessageArguments, opt AppMessage) -> (CanisterWsMessageResult); From 6add8d1d5a3f6a45006af2c0f18b42de95605ff0 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 8 Dec 2023 16:01:04 +0100 Subject: [PATCH 10/18] perf: rename constant --- src/ic-websocket-cdk/src/lib.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ic-websocket-cdk/src/lib.rs b/src/ic-websocket-cdk/src/lib.rs index 561cefd..dacada1 100644 --- a/src/ic-websocket-cdk/src/lib.rs +++ b/src/ic-websocket-cdk/src/lib.rs @@ -30,10 +30,10 @@ const LABEL_WEBSOCKET: &[u8] = b"websocket"; const DEFAULT_MAX_NUMBER_OF_RETURNED_MESSAGES: usize = 50; /// The default interval at which to send acknowledgements to the client. const DEFAULT_SEND_ACK_INTERVAL_MS: u64 = 300_000; // 5 minutes -/// The maximum latency allowed between the client and the canister. -const MAX_ALLOWED_CONNECTION_LATENCY_MS: u64 = 30_000; // 30 seconds +/// The maximum network latency allowed between the client and the canister. +const MAX_ALLOWED_NETWORK_LATENCY_MS: u64 = 30_000; // 30 seconds /// The default timeout to wait for the client to send a keep alive after receiving an acknowledgement. -const CLIENT_KEEP_ALIVE_TIMEOUT_MS: u64 = 2 * MAX_ALLOWED_CONNECTION_LATENCY_MS; +const CLIENT_KEEP_ALIVE_TIMEOUT_MS: u64 = 2 * MAX_ALLOWED_NETWORK_LATENCY_MS; /// Same as [CLIENT_KEEP_ALIVE_TIMEOUT_MS], but in nanoseconds. const CLIENT_KEEP_ALIVE_TIMEOUT_NS: u64 = CLIENT_KEEP_ALIVE_TIMEOUT_MS * 1_000_000; From 55ad12130d0c4f8b2e94971a0e4cd4c4fef8653f Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Sat, 9 Dec 2023 19:11:48 +0100 Subject: [PATCH 11/18] chore: use with_... to init params (test canister) --- src/test_canister/src/lib.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/test_canister/src/lib.rs b/src/test_canister/src/lib.rs index a343590..b1ea705 100644 --- a/src/test_canister/src/lib.rs +++ b/src/test_canister/src/lib.rs @@ -18,11 +18,9 @@ fn init(max_number_of_returned_messages: usize, send_ack_interval_ms: u64) { on_close: Some(on_close), }; - let params = WsInitParams { - handlers, - max_number_of_returned_messages, - send_ack_interval_ms, - }; + let params = WsInitParams::new(handlers) + .with_max_number_of_returned_messages(max_number_of_returned_messages) + .with_send_ack_interval_ms(send_ack_interval_ms); ic_websocket_cdk::init(params) } From 67b1bdd9e26b748897f8ba1b4a4abe84a28476ea Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 15 Dec 2023 09:45:50 +0100 Subject: [PATCH 12/18] perf: rename const --- src/ic-websocket-cdk/src/lib.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ic-websocket-cdk/src/lib.rs b/src/ic-websocket-cdk/src/lib.rs index dacada1..2857c09 100644 --- a/src/ic-websocket-cdk/src/lib.rs +++ b/src/ic-websocket-cdk/src/lib.rs @@ -30,10 +30,10 @@ const LABEL_WEBSOCKET: &[u8] = b"websocket"; const DEFAULT_MAX_NUMBER_OF_RETURNED_MESSAGES: usize = 50; /// The default interval at which to send acknowledgements to the client. const DEFAULT_SEND_ACK_INTERVAL_MS: u64 = 300_000; // 5 minutes -/// The maximum network latency allowed between the client and the canister. -const MAX_ALLOWED_NETWORK_LATENCY_MS: u64 = 30_000; // 30 seconds +/// The maximum communication latency allowed between the client and the canister. +const COMMUNICATION_LATENCY_BOUND_MS: u64 = 30_000; // 30 seconds /// The default timeout to wait for the client to send a keep alive after receiving an acknowledgement. -const CLIENT_KEEP_ALIVE_TIMEOUT_MS: u64 = 2 * MAX_ALLOWED_NETWORK_LATENCY_MS; +const CLIENT_KEEP_ALIVE_TIMEOUT_MS: u64 = 2 * COMMUNICATION_LATENCY_BOUND_MS; /// Same as [CLIENT_KEEP_ALIVE_TIMEOUT_MS], but in nanoseconds. const CLIENT_KEEP_ALIVE_TIMEOUT_NS: u64 = CLIENT_KEEP_ALIVE_TIMEOUT_MS * 1_000_000; From 24fe40b10b5e196e8ffe014bc2c17629a192b121 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 15 Dec 2023 11:06:46 +0100 Subject: [PATCH 13/18] chore: update pocket-ic to v2 --- .github/workflows/tests.yml | 2 +- Cargo.lock | 520 +++++++++++++++++- scripts/download-pocket-ic.sh | 2 +- scripts/test.sh | 2 +- src/ic-websocket-cdk/Cargo.toml | 2 +- .../src/tests/integration_tests/a_ws_open.rs | 10 +- .../tests/integration_tests/b_ws_message.rs | 18 +- .../integration_tests/c_ws_get_messages.rs | 19 +- .../src/tests/integration_tests/d_ws_close.rs | 8 +- .../src/tests/integration_tests/e_send.rs | 7 +- .../src/tests/integration_tests/f_close.rs | 11 +- .../g_messages_acknowledgement.rs | 12 +- .../integration_tests/h_multiple_gateways.rs | 5 +- .../tests/integration_tests/utils/actor.rs | 41 +- .../integration_tests/utils/certification.rs | 3 +- .../tests/integration_tests/utils/test_env.rs | 80 +-- src/test_canister/src/lib.rs | 8 + src/test_canister/test_canister.did | 1 + 18 files changed, 639 insertions(+), 112 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 5c590bc..23716df 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -35,4 +35,4 @@ jobs: ./scripts/download-pocket-ic.sh ./scripts/build-test-canister.sh - name: Run integration tests - run: POCKET_IC_BIN="$(pwd)/bin/pocket-ic" cargo test --package ic-websocket-cdk --lib -- tests::integration_tests --test-threads 1 + run: POCKET_IC_MUTE_SERVER=1 POCKET_IC_BIN="$(pwd)/bin/pocket-ic" cargo test --package ic-websocket-cdk --lib -- tests::integration_tests --test-threads 1 diff --git a/Cargo.lock b/Cargo.lock index 7a1bf8a..e4edbad 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,15 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" +[[package]] +name = "aho-corasick" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2969dcb958b36655471fc61f7e416fa76033bdd4bfed0678d8fee1e2d07a1f0" +dependencies = [ + "memchr", +] + [[package]] name = "anyhow" version = "1.0.75" @@ -29,6 +38,15 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "23b62fc65de8e4e7f52534fb52b0f3ed04746ae267519eef2a83941e8085068b" +[[package]] +name = "ascii-canvas" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8824ecca2e851cec16968d54a01dd372ef8f95b244fb84b84e70128be347c3c6" +dependencies = [ + "term", +] + [[package]] name = "async-trait" version = "0.1.74" @@ -96,6 +114,12 @@ version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8c3c1a368f70d6cf7302d78f8f7093da241fb8e8807c05cc9e51a125895a6d5b" +[[package]] +name = "beef" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a8241f3ebb85c056b509d4327ad0358fbbba6ffb340bf388f26350aeda225b1" + [[package]] name = "binread" version = "2.2.0" @@ -207,10 +231,14 @@ dependencies = [ "byteorder", "candid_derive", "codespan-reporting", + "convert_case", "crc32fast", "data-encoding", "hex", + "lalrpop", + "lalrpop-util", "leb128", + "logos", "num-bigint", "num-traits", "num_enum", @@ -266,6 +294,15 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28c122c3980598d243d63d9a704629a2d748d101f278052ff068be5a4423ab6f" +[[package]] +name = "convert_case" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec182b0ca2f35d8fc196cf3404988fd8b8c739a4d270ff118a398feb0cbec1ca" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "core-foundation" version = "0.9.3" @@ -300,6 +337,31 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a33c2bf77f2df06183c3aa30d1e96c0695a313d4f9c453cc3762a6db39f99200" +dependencies = [ + "cfg-if", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a22b2d63d4d1dc0b7f1b6b2747dd0088008a9be28b6ddf0b1e7d335e3037294" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crunchy" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" + [[package]] name = "crypto-bigint" version = "0.5.3" @@ -348,6 +410,12 @@ dependencies = [ "powerfmt", ] +[[package]] +name = "diff" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56254986775e3233ffa9c4d7d3faaf6d36a2c09d30b20687e9f88bc8bafc16c8" + [[package]] name = "digest" version = "0.9.0" @@ -369,6 +437,27 @@ dependencies = [ "subtle", ] +[[package]] +name = "dirs-next" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b98cf8ebf19c3d1b223e151f99a4f9f0690dca41414773390fc824184ac833e1" +dependencies = [ + "cfg-if", + "dirs-sys-next", +] + +[[package]] +name = "dirs-sys-next" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ebda144c4fe02d1f7ea1a7d9641b6fc6b580adcfa024ae48797ecdeb6825b4d" +dependencies = [ + "libc", + "redox_users", + "winapi", +] + [[package]] name = "ecdsa" version = "0.16.8" @@ -409,6 +498,15 @@ dependencies = [ "zeroize", ] +[[package]] +name = "ena" +version = "0.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c533630cf40e9caa44bd91aadc88a75d75a4c3a12b4cfde353cbed41daa1e1f1" +dependencies = [ + "log", +] + [[package]] name = "encoding_rs" version = "0.8.33" @@ -460,6 +558,12 @@ dependencies = [ "subtle", ] +[[package]] +name = "fixedbitset" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" + [[package]] name = "fnv" version = "1.0.7" @@ -970,6 +1074,26 @@ version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f518f335dce6725a761382244631d86cf0ccb2863413590b31338feb467f9c3" +[[package]] +name = "is-terminal" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb0889898416213fab133e1d33a0e5858a48177452750691bde3666d0fdbaf8b" +dependencies = [ + "hermit-abi", + "rustix", + "windows-sys", +] + +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.9" @@ -999,6 +1123,38 @@ dependencies = [ "signature", ] +[[package]] +name = "lalrpop" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da4081d44f4611b66c6dd725e6de3169f9f63905421e8626fcb86b6a898998b8" +dependencies = [ + "ascii-canvas", + "bit-set", + "diff", + "ena", + "is-terminal", + "itertools", + "lalrpop-util", + "petgraph", + "pico-args", + "regex", + "regex-syntax 0.7.5", + "string_cache", + "term", + "tiny-keccak", + "unicode-xid", +] + +[[package]] +name = "lalrpop-util" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f35c735096c0293d313e8f2a641627472b83d01b937177fe76e5e2708d31e0d" +dependencies = [ + "regex", +] + [[package]] name = "lazy_static" version = "1.4.0" @@ -1023,18 +1179,80 @@ version = "0.2.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" +[[package]] +name = "libredox" +version = "0.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85c833ca1e66078851dba29046874e38f08b2c883700aa29a03ddd3b23814ee8" +dependencies = [ + "bitflags 2.4.1", + "libc", + "redox_syscall", +] + [[package]] name = "linux-raw-sys" version = "0.4.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da2479e8c062e40bf0066ffa0bc823de0a9368974af99c9f6df941d2c231e03f" +[[package]] +name = "lock_api" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c168f8615b12bc01f9c17e2eb0cc07dcae1940121185446edc3744920e8ef45" +dependencies = [ + "autocfg", + "scopeguard", +] + [[package]] name = "log" version = "0.4.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f" +[[package]] +name = "logos" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c000ca4d908ff18ac99b93a062cb8958d331c3220719c52e77cb19cc6ac5d2c1" +dependencies = [ + "logos-derive", +] + +[[package]] +name = "logos-codegen" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc487311295e0002e452025d6b580b77bb17286de87b57138f3b5db711cded68" +dependencies = [ + "beef", + "fnv", + "proc-macro2", + "quote", + "regex-syntax 0.6.29", + "syn 2.0.38", +] + +[[package]] +name = "logos-derive" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbfc0d229f1f42d790440136d941afd806bc9e949e2bcb8faa813b0f00d1267e" +dependencies = [ + "logos-codegen", +] + +[[package]] +name = "matchers" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8263075bb86c5a1b1427b5ae862e8889656f126e9f77c484496e8b47cf5c5558" +dependencies = [ + "regex-automata 0.1.10", +] + [[package]] name = "memchr" version = "2.6.4" @@ -1047,6 +1265,16 @@ version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "mime_guess" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4192263c238a5f0d0c6bfd21f336a313a4ce1c450542449ca191bb657b4642ef" +dependencies = [ + "mime", + "unicase", +] + [[package]] name = "minimal-lexical" version = "0.2.1" @@ -1079,6 +1307,12 @@ version = "4.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d07cbe42e2a8dd41df582fb8e00fc24d920b5561cc301fcb6d14e2e0434b500f" +[[package]] +name = "new_debug_unreachable" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4a24736216ec316047a1fc4252e27dabb04218aa4a3f37c6e7ddbf1f9782b54" + [[package]] name = "nom" version = "7.1.3" @@ -1089,6 +1323,16 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "nu-ansi-term" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +dependencies = [ + "overload", + "winapi", +] + [[package]] name = "num-bigint" version = "0.4.4" @@ -1173,6 +1417,12 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "624a8340c38c1b80fd549087862da4ba43e08858af025b236e509b6649fc13d5" +[[package]] +name = "overload" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" + [[package]] name = "pairing" version = "0.22.0" @@ -1182,6 +1432,29 @@ dependencies = [ "group 0.12.1", ] +[[package]] +name = "parking_lot" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c42a9226546d68acdd9c0a280d17ce19bfe27a46bf68784e4066115788d008e" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-targets", +] + [[package]] name = "paste" version = "1.0.14" @@ -1213,6 +1486,31 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b2a4787296e9989611394c33f193f676704af1686e70b8f8033ab5ba9a35a94" +[[package]] +name = "petgraph" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d3afd2628e69da2be385eb6f2fd57c8ac7977ceeff6dc166ff1657b0e386a9" +dependencies = [ + "fixedbitset", + "indexmap 2.1.0", +] + +[[package]] +name = "phf_shared" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6796ad771acdc0123d2a88dc428b5e38ef24456743ddb1744ed628f9815c096" +dependencies = [ + "siphasher", +] + +[[package]] +name = "pico-args" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5be167a7af36ee22fe3115051bc51f6e6c7054c9348e28deb4f49bd6f705a315" + [[package]] name = "pin-project-lite" version = "0.2.13" @@ -1237,9 +1535,9 @@ dependencies = [ [[package]] name = "pocket-ic" -version = "1.0.0" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3ec0d4fbef9949836203390b819d2ad944f405030f93c724f78d20184d59f85" +checksum = "ce4c698117a4cc5fce56a909db4fce965003f041c1bf916b7cd8a6757b017586" dependencies = [ "async-trait", "base64 0.13.1", @@ -1250,6 +1548,9 @@ dependencies = [ "serde", "serde_bytes", "serde_json", + "tracing", + "tracing-appender", + "tracing-subscriber", ] [[package]] @@ -1264,6 +1565,12 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" +[[package]] +name = "precomputed-hash" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "925383efa346730478fb4838dbe9137d2a47675ad789c546d150a6e1dd4ab31c" + [[package]] name = "pretty" version = "0.12.3" @@ -1308,7 +1615,7 @@ dependencies = [ "rand", "rand_chacha", "rand_xorshift", - "regex-syntax", + "regex-syntax 0.7.5", "rusty-fork", "tempfile", "unarray", @@ -1386,12 +1693,67 @@ dependencies = [ "bitflags 1.3.2", ] +[[package]] +name = "redox_users" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a18479200779601e498ada4e8c1e1f50e3ee19deb0259c25825a98b5603b2cb4" +dependencies = [ + "getrandom", + "libredox", + "thiserror", +] + +[[package]] +name = "regex" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "380b951a9c5e80ddfd6136919eef32310721aa4aacd4889a8d39124b026ab343" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata 0.4.3", + "regex-syntax 0.8.2", +] + +[[package]] +name = "regex-automata" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" +dependencies = [ + "regex-syntax 0.6.29", +] + +[[package]] +name = "regex-automata" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f804c7828047e88b2d32e2d7fe5a105da8ee3264f01902f796c8e067dc2483f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax 0.8.2", +] + +[[package]] +name = "regex-syntax" +version = "0.6.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" + [[package]] name = "regex-syntax" version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dbb5fb1acd8a1a18b3dd5be62d25485eb770e05afb408a9627d14d451bae12da" +[[package]] +name = "regex-syntax" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08c74e62047bb2de4ff487b251e4a92e24f48745648451635cec7d591162d9f" + [[package]] name = "reqwest" version = "0.11.22" @@ -1412,6 +1774,7 @@ dependencies = [ "js-sys", "log", "mime", + "mime_guess", "once_cell", "percent-encoding", "pin-project-lite", @@ -1559,6 +1922,12 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ad4cc8da4ef723ed60bced201181d83791ad433213d8c24efffda1eec85d741" +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "sct" version = "0.7.1" @@ -1691,6 +2060,15 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "signature" version = "2.1.0" @@ -1713,6 +2091,12 @@ dependencies = [ "time", ] +[[package]] +name = "siphasher" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38b58827f4464d87d377d175e90bf58eb00fd8716ff0a62f80356b5e61555d0d" + [[package]] name = "slab" version = "0.4.9" @@ -1731,6 +2115,12 @@ dependencies = [ "version_check", ] +[[package]] +name = "smallvec" +version = "1.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4dccd0940a2dcdf68d092b8cbab7dc0ad8fa938bf95787e1b916b0e3d0e8e970" + [[package]] name = "socket2" version = "0.4.10" @@ -1786,6 +2176,19 @@ dependencies = [ "winapi", ] +[[package]] +name = "string_cache" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f91138e76242f575eb1d3b38b4f1362f10d3a43f47d182a5b359af488a02293b" +dependencies = [ + "new_debug_unreachable", + "once_cell", + "parking_lot", + "phf_shared", + "precomputed-hash", +] + [[package]] name = "subtle" version = "2.5.0" @@ -1848,6 +2251,17 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "term" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c59df8ac95d96ff9bede18eb7300b0fda5e5d8d90960e76f8e14ae765eedbf1f" +dependencies = [ + "dirs-next", + "rustversion", + "winapi", +] + [[package]] name = "termcolor" version = "1.3.0" @@ -1889,6 +2303,16 @@ dependencies = [ "syn 2.0.38", ] +[[package]] +name = "thread_local" +version = "1.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fdd6f064ccff2d6567adcb3873ca630700f00b5ad3f060c25b5dcfd9a4ce152" +dependencies = [ + "cfg-if", + "once_cell", +] + [[package]] name = "time" version = "0.3.30" @@ -1918,6 +2342,15 @@ dependencies = [ "time-core", ] +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + [[package]] name = "tinyvec" version = "1.6.0" @@ -2003,9 +2436,33 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef" dependencies = [ "pin-project-lite", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-appender" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3566e8ce28cc0a3fe42519fc80e6b4c943cc4c8cef275620eb8dac2d3d4e06cf" +dependencies = [ + "crossbeam-channel", + "thiserror", + "time", + "tracing-subscriber", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.38", +] + [[package]] name = "tracing-core" version = "0.1.32" @@ -2013,6 +2470,36 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" dependencies = [ "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", ] [[package]] @@ -2039,6 +2526,15 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" +[[package]] +name = "unicase" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7d2d4dafb69621809a81864c9c1b864479e1235c0dd4e199924b9742439ed89" +dependencies = [ + "version_check", +] + [[package]] name = "unicode-bidi" version = "0.3.13" @@ -2060,12 +2556,24 @@ 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" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85" +[[package]] +name = "unicode-xid" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f962df74c8c05a667b5ee8bcf162993134c104e96440b663c8daa176dc772d8c" + [[package]] name = "untrusted" version = "0.7.1" @@ -2089,6 +2597,12 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "valuable" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" + [[package]] name = "version_check" version = "0.9.4" diff --git a/scripts/download-pocket-ic.sh b/scripts/download-pocket-ic.sh index 56687ec..43afc18 100755 --- a/scripts/download-pocket-ic.sh +++ b/scripts/download-pocket-ic.sh @@ -8,7 +8,7 @@ if [ -f "$POCKET_IC_BIN" ]; then else echo "$POCKET_IC_BIN does not exist." echo "Downloading Pocket IC binary..." - curl -sLO https://download.dfinity.systems/ic/307d5847c1d2fe1f5e19181c7d0fcec23f4658b3/openssl-static-binaries/x86_64-linux/pocket-ic.gz + curl -sLO https://download.dfinity.systems/ic/69e1408347723dbaa7a6cd2faa9b65c42abbe861/openssl-static-binaries/x86_64-linux/pocket-ic.gz echo "Extracting Pocket IC binary..." gzip -d $POCKET_IC_BIN.gz diff --git a/scripts/test.sh b/scripts/test.sh index 2bf9891..51000c8 100755 --- a/scripts/test.sh +++ b/scripts/test.sh @@ -12,4 +12,4 @@ cargo test --package ic-websocket-cdk --doc ./scripts/build-test-canister.sh -POCKET_IC_BIN="$(pwd)/bin/pocket-ic" cargo test --package ic-websocket-cdk --lib -- tests::integration_tests --test-threads 1 +POCKET_IC_MUTE_SERVER=1 POCKET_IC_BIN="$(pwd)/bin/pocket-ic" cargo test --package ic-websocket-cdk --lib -- tests::integration_tests --test-threads 1 diff --git a/src/ic-websocket-cdk/Cargo.toml b/src/ic-websocket-cdk/Cargo.toml index b5aa2ca..af582e2 100644 --- a/src/ic-websocket-cdk/Cargo.toml +++ b/src/ic-websocket-cdk/Cargo.toml @@ -29,7 +29,7 @@ proptest = "1.2.0" rand = "0.8.5" ring = "0.16.20" lazy_static = "1.4.0" -pocket-ic = "1.0.0" +pocket-ic = "2.0.1" ic-certificate-verification = "1.2.0" ic-certification = "1.2.0" diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/a_ws_open.rs b/src/ic-websocket-cdk/src/tests/integration_tests/a_ws_open.rs index 6f75fef..6ef926d 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/a_ws_open.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/a_ws_open.rs @@ -2,9 +2,10 @@ use proptest::prelude::*; use std::ops::Deref; use crate::{ - errors::WsError, CanisterOutputCertifiedMessages, CanisterOutputMessage, - CanisterWsGetMessagesArguments, CanisterWsOpenArguments, CanisterWsOpenResult, ClientKey, - WebsocketServiceMessageContent, DEFAULT_MAX_NUMBER_OF_RETURNED_MESSAGES, + errors::WsError, tests::integration_tests::utils::actor::wipe::call_wipe, + CanisterOutputCertifiedMessages, CanisterOutputMessage, CanisterWsGetMessagesArguments, + CanisterWsOpenArguments, CanisterWsOpenResult, ClientKey, WebsocketServiceMessageContent, + DEFAULT_MAX_NUMBER_OF_RETURNED_MESSAGES, }; use candid::Principal; @@ -12,7 +13,6 @@ use super::utils::{ actor::{ws_get_messages::call_ws_get_messages_with_panic, ws_open::call_ws_open}, clients::{generate_random_client_nonce, CLIENT_1_KEY, GATEWAY_1}, messages::get_service_message_content_from_canister_message, - test_env::get_test_env, }; #[test] @@ -73,7 +73,7 @@ fn test_3_fails_for_a_client_with_the_same_nonce() { ); // reset canister for the next test - get_test_env().reset_canister_with_default_params(); + call_wipe(None); } proptest! { diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/b_ws_message.rs b/src/ic-websocket-cdk/src/tests/integration_tests/b_ws_message.rs index 1a5aea9..40d73c6 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/b_ws_message.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/b_ws_message.rs @@ -7,7 +7,10 @@ use crate::{ tests::{ common::generate_random_principal, integration_tests::utils::{ - actor::{ws_close::call_ws_close, ws_get_messages::call_ws_get_messages_with_panic}, + actor::{ + wipe::call_wipe, ws_close::call_ws_close, + ws_get_messages::call_ws_get_messages_with_panic, + }, clients::GATEWAY_1, messages::check_canister_message_has_close_reason, }, @@ -22,13 +25,12 @@ use super::utils::{ actor::{ws_message::call_ws_message, ws_open::call_ws_open_for_client_key_with_panic}, clients::{generate_random_client_nonce, CLIENT_1_KEY, CLIENT_2_KEY}, messages::{create_websocket_message, encode_websocket_service_message_content}, - test_env::get_test_env, }; #[test] fn test_1_fails_if_client_is_not_registered() { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(CLIENT_1_KEY.deref()); @@ -55,7 +57,7 @@ fn test_1_fails_if_client_is_not_registered() { fn test_2_fails_if_client_sends_a_message_with_a_different_client_key() { let client_1_key = CLIENT_1_KEY.deref(); // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(client_1_key); @@ -108,7 +110,7 @@ fn test_2_fails_if_client_sends_a_message_with_a_different_client_key() { fn test_3_should_send_a_message_from_a_registered_client() { let client_1_key = CLIENT_1_KEY.deref(); // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(client_1_key); @@ -125,7 +127,7 @@ fn test_3_should_send_a_message_from_a_registered_client() { fn test_4_fails_if_client_sends_a_message_with_a_wrong_sequence_number() { let client_1_key = CLIENT_1_KEY.deref(); // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(client_1_key); @@ -181,7 +183,7 @@ fn test_4_fails_if_client_sends_a_message_with_a_wrong_sequence_number() { fn test_5_fails_if_client_sends_a_wrong_service_message() { let client_1_key = CLIENT_1_KEY.deref(); // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(client_1_key); @@ -230,7 +232,7 @@ fn test_5_fails_if_client_sends_a_wrong_service_message() { fn test_6_should_send_a_service_message_from_a_registered_client() { let client_1_key = CLIENT_1_KEY.deref(); // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(client_1_key); diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs b/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs index 314d48d..9242355 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/c_ws_get_messages.rs @@ -8,7 +8,8 @@ use crate::{ use super::utils::{ actor::{ - send::call_send_with_panic, ws_get_messages::call_ws_get_messages_with_panic, + send::call_send_with_panic, wipe::call_wipe, + ws_get_messages::call_ws_get_messages_with_panic, ws_open::call_ws_open_for_client_key_with_panic, }, clients::{generate_random_client_key, CLIENT_1_KEY, GATEWAY_1}, @@ -25,7 +26,7 @@ proptest! { #[test] fn test_1_non_registered_gateway_should_receive_empty_messages(ref test_gateway_principal in any::().prop_map(|_| common::generate_random_principal())) { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); let res = call_ws_get_messages_with_panic( test_gateway_principal, @@ -45,7 +46,7 @@ proptest! { #[test] fn test_2_registered_gateway_should_receive_only_open_message_if_no_messages_sent_initial_nonce(ref test_client_key in any::().prop_map(|_| generate_random_client_key())) { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, register client 1 call_ws_open_for_client_key_with_panic(test_client_key); @@ -60,7 +61,7 @@ proptest! { #[test] fn test_3_registered_gateway_can_receive_correct_amount_of_messages(test_send_messages_count in 1..100u64) { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, register client 1 let client_1_key = CLIENT_1_KEY.deref(); call_ws_open_for_client_key_with_panic(client_1_key); @@ -121,7 +122,7 @@ proptest! { #[test] fn test_4_registered_gateway_can_receive_certified_messages(test_send_messages_count in 1..100u64) { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, register client 1 let client_1_key = CLIENT_1_KEY.deref(); call_ws_open_for_client_key_with_panic(client_1_key); @@ -165,9 +166,11 @@ proptest! { #[test] fn test_5_messages_for_gateway_are_deleted_if_old(test_send_messages_count in 1..100usize) { // first, reset the canister - get_test_env().reset_canister( - 1_000, // avoid the queue size limit - DEFAULT_TEST_SEND_ACK_INTERVAL_MS, + call_wipe( + Some(( + 1_000, // avoid the queue size limit + DEFAULT_TEST_SEND_ACK_INTERVAL_MS, + )) ); // second, register client 1 let client_1_key = CLIENT_1_KEY.deref(); diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/d_ws_close.rs b/src/ic-websocket-cdk/src/tests/integration_tests/d_ws_close.rs index 73df174..ec2c496 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/d_ws_close.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/d_ws_close.rs @@ -4,6 +4,7 @@ use crate::{errors::WsError, CanisterWsCloseArguments, CanisterWsCloseResult}; use super::utils::{ actor::{ + wipe::call_wipe, ws_close::call_ws_close, ws_open::{ call_ws_open_for_client_key_and_gateway_with_panic, @@ -11,13 +12,12 @@ use super::utils::{ }, }, clients::{CLIENT_1_KEY, CLIENT_2_KEY, GATEWAY_1, GATEWAY_2}, - test_env::get_test_env, }; #[test] fn test_1_fails_if_gateway_is_not_registered() { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(CLIENT_1_KEY.deref()); @@ -44,7 +44,7 @@ fn test_1_fails_if_gateway_is_not_registered() { #[test] fn test_2_fails_if_client_is_not_registered() { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(CLIENT_1_KEY.deref()); @@ -69,7 +69,7 @@ fn test_2_fails_if_client_is_not_registered() { #[test] fn test_3_fails_if_client_is_not_registered_to_gateway() { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for both clients call_ws_open_for_client_key_and_gateway_with_panic(CLIENT_1_KEY.deref(), *GATEWAY_1.deref()); call_ws_open_for_client_key_and_gateway_with_panic(CLIENT_2_KEY.deref(), *GATEWAY_2.deref()); diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/e_send.rs b/src/ic-websocket-cdk/src/tests/integration_tests/e_send.rs index 1e59a3f..cbe81cf 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/e_send.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/e_send.rs @@ -3,16 +3,15 @@ use std::ops::Deref; use crate::{errors::WsError, CanisterSendResult}; use super::utils::{ - actor::{send::call_send, ws_open::call_ws_open_for_client_key_with_panic}, + actor::{send::call_send, wipe::call_wipe, ws_open::call_ws_open_for_client_key_with_panic}, clients::{CLIENT_1_KEY, CLIENT_2_KEY}, messages::AppMessage, - test_env::get_test_env, }; #[test] fn test_1_fails_if_sending_a_message_to_a_non_registered_client() { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(CLIENT_1_KEY.deref()); @@ -38,7 +37,7 @@ fn test_1_fails_if_sending_a_message_to_a_non_registered_client() { #[test] fn test_2_should_send_a_message_to_a_registered_client() { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(CLIENT_1_KEY.deref()); diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/f_close.rs b/src/ic-websocket-cdk/src/tests/integration_tests/f_close.rs index 0aa647e..61c80b1 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/f_close.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/f_close.rs @@ -1,19 +1,16 @@ use std::ops::Deref; -use crate::{ - errors::WsError, tests::integration_tests::utils::actor::close::call_close, CanisterCloseResult, -}; +use crate::{errors::WsError, CanisterCloseResult}; use super::utils::{ - actor::ws_open::call_ws_open_for_client_key_with_panic, + actor::{close::call_close, wipe::call_wipe, ws_open::call_ws_open_for_client_key_with_panic}, clients::{CLIENT_1_KEY, CLIENT_2_KEY}, - test_env::get_test_env, }; #[test] fn test_1_fails_closing_for_non_registered_client() { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(CLIENT_1_KEY.deref()); @@ -34,7 +31,7 @@ fn test_1_fails_closing_for_non_registered_client() { #[test] fn test_2_should_close_connection_for_registered_client() { // first, reset the canister - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // second, open a connection for client 1 call_ws_open_for_client_key_with_panic(CLIENT_1_KEY.deref()); diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/g_messages_acknowledgement.rs b/src/ic-websocket-cdk/src/tests/integration_tests/g_messages_acknowledgement.rs index 24084a3..76d1a11 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/g_messages_acknowledgement.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/g_messages_acknowledgement.rs @@ -13,8 +13,8 @@ use crate::{ use super::utils::{ actor::{ - ws_get_messages::call_ws_get_messages_with_panic, ws_message::call_ws_message_with_panic, - ws_open::call_ws_open_for_client_key_with_panic, + wipe::call_wipe, ws_get_messages::call_ws_get_messages_with_panic, + ws_message::call_ws_message_with_panic, ws_open::call_ws_open_for_client_key_with_panic, }, certification::{is_message_body_valid, is_valid_certificate}, clients::{CLIENT_1_KEY, GATEWAY_1}, @@ -27,7 +27,7 @@ use super::utils::{ #[test] fn test_1_client_should_receive_ack_messages() { - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // open a connection for client 1 let client_1_key = CLIENT_1_KEY.deref(); call_ws_open_for_client_key_with_panic(client_1_key); @@ -58,7 +58,7 @@ fn test_1_client_should_receive_ack_messages() { fn test_2_client_is_removed_if_keep_alive_timeout_is_reached() { let client_1_key = CLIENT_1_KEY.deref(); // open a connection for client 1 - get_test_env().reset_canister_with_default_params(); + call_wipe(None); call_ws_open_for_client_key_with_panic(client_1_key); // advance the canister time to make sure the ack timer expires and an ack is sent get_test_env().advance_canister_time_ms(DEFAULT_TEST_SEND_ACK_INTERVAL_MS); @@ -104,7 +104,7 @@ fn test_2_client_is_removed_if_keep_alive_timeout_is_reached() { #[test] fn test_3_client_is_not_removed_if_it_sends_a_keep_alive_before_timeout() { let client_1_key = CLIENT_1_KEY.deref(); - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // open a connection for client 1 call_ws_open_for_client_key_with_panic(client_1_key); // advance the canister time to make sure the ack timer expires and an ack is sent @@ -157,7 +157,7 @@ fn test_3_client_is_not_removed_if_it_sends_a_keep_alive_before_timeout() { #[test] fn test_4_client_is_not_removed_if_it_connects_while_canister_is_waiting_for_keep_alive() { let client_1_key = CLIENT_1_KEY.deref(); - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // advance the canister time to make sure the ack timer expires and the canister started the keep alive timer get_test_env().advance_canister_time_ms(DEFAULT_TEST_SEND_ACK_INTERVAL_MS); // open a connection for client 1 diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/h_multiple_gateways.rs b/src/ic-websocket-cdk/src/tests/integration_tests/h_multiple_gateways.rs index 6926554..32317b4 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/h_multiple_gateways.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/h_multiple_gateways.rs @@ -7,12 +7,11 @@ use crate::{ use super::utils::{ actor::{ - send::call_send_with_panic, ws_close::call_ws_close_with_panic, + send::call_send_with_panic, wipe::call_wipe, ws_close::call_ws_close_with_panic, ws_get_messages::call_ws_get_messages_with_panic, ws_open::call_ws_open_for_client_key_and_gateway_with_panic, }, messages::{get_service_message_content_from_canister_message, verify_messages, AppMessage}, - test_env::get_test_env, }; proptest! { @@ -25,7 +24,7 @@ proptest! { ) { let first_gateway = &gateways[0]; let second_gateway = &gateways[1]; - get_test_env().reset_canister_with_default_params(); + call_wipe(None); // open a connection for client call_ws_open_for_client_key_and_gateway_with_panic(&client_key, *first_gateway); // simulate canister sending messages to client diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs b/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs index be1eee2..47512d6 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/utils/actor.rs @@ -13,7 +13,7 @@ pub mod ws_open { /// # Panics /// if the call returns a [WasmResult::Reject]. pub fn call_ws_open(caller: &Principal, args: CanisterWsOpenArguments) -> CanisterWsOpenResult { - let canister_id = get_test_env().canister_id; + let canister_id = get_test_env().get_test_canister_id(); let res = get_test_env() .pic .update_call(canister_id, *caller, "ws_open", encode_one(args).unwrap()) @@ -71,7 +71,7 @@ pub mod ws_message { caller: &Principal, args: CanisterWsMessageArguments, ) -> CanisterWsMessageResult { - let canister_id = get_test_env().canister_id; + let canister_id = get_test_env().get_test_canister_id(); let res = get_test_env() .pic .update_call( @@ -111,7 +111,7 @@ pub mod ws_close { caller: &Principal, args: CanisterWsCloseArguments, ) -> CanisterWsCloseResult { - let canister_id = get_test_env().canister_id; + let canister_id = get_test_env().get_test_canister_id(); let res = get_test_env() .pic .update_call(canister_id, *caller, "ws_close", encode_one(args).unwrap()) @@ -149,7 +149,7 @@ pub mod ws_get_messages { caller: &Principal, args: CanisterWsGetMessagesArguments, ) -> CanisterWsGetMessagesResult { - let canister_id = get_test_env().canister_id; + let canister_id = get_test_env().get_test_canister_id(); let res = get_test_env() .pic .query_call( @@ -198,7 +198,7 @@ pub mod send { ) -> CanisterSendResult { let messages: Vec> = messages.iter().map(|m| encode_one(m).unwrap()).collect(); let args: SendArguments = (send_to_principal.clone(), messages); - let canister_id = get_test_env().canister_id; + let canister_id = get_test_env().get_test_canister_id(); let res = get_test_env() .pic .update_call( @@ -238,7 +238,7 @@ pub mod close { /// # Panics /// if the call returns a [WasmResult::Reject]. pub(crate) fn call_close(client_principal: &ClientPrincipal) -> CanisterCloseResult { - let canister_id = get_test_env().canister_id; + let canister_id = get_test_env().get_test_canister_id(); let res = get_test_env() .pic .update_call( @@ -254,3 +254,32 @@ pub mod close { } } } + +pub mod wipe { + use candid::encode_args; + + use crate::tests::integration_tests::utils::test_env::{ + CanisterInitArgs, DEFAULT_CANISTER_INIT_ARGS, + }; + + use super::*; + + /// # Panics + /// if the call returns a [WasmResult::Reject]. + pub(crate) fn call_wipe(init_args: Option) { + let canister_id = get_test_env().get_test_canister_id(); + let res = get_test_env() + .pic + .update_call( + canister_id, + Principal::anonymous(), + "wipe", + encode_args(init_args.unwrap_or(DEFAULT_CANISTER_INIT_ARGS)).unwrap(), + ) + .expect("Failed to call canister"); + match res { + WasmResult::Reply(_) => {}, + _ => panic!("Expected reply"), + } + } +} diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/utils/certification.rs b/src/ic-websocket-cdk/src/tests/integration_tests/utils/certification.rs index edfded0..45a56e1 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/utils/certification.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/utils/certification.rs @@ -6,7 +6,8 @@ use super::test_env::TestEnv; pub fn is_valid_certificate(test_env: &TestEnv, certificate: &[u8], tree: &[u8]) -> bool { let cert: Certificate = serde_cbor::from_slice(certificate).unwrap(); - let canister_id_bytes = test_env.canister_id.as_slice(); + let canister_id = test_env.get_test_canister_id(); + let canister_id_bytes = canister_id.as_slice(); let verify_res = cert.verify(canister_id_bytes, &test_env.get_root_ic_key()); match verify_res { Ok(_) => { diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs b/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs index bb5c2b8..310e3bd 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs @@ -6,7 +6,7 @@ use std::{ use candid::Principal; use lazy_static::lazy_static; -use pocket_ic::PocketIc; +use pocket_ic::{PocketIc, PocketIcBuilder}; use super::wasm::{load_canister_wasm_from_bin, load_canister_wasm_from_path}; @@ -19,8 +19,20 @@ pub const DEFAULT_TEST_MAX_NUMBER_OF_RETURNED_MESSAGES: u64 = 50; /// Value: `300_000` = 5 minutes pub const DEFAULT_TEST_SEND_ACK_INTERVAL_MS: u64 = 300_000; +/// (`max_number_or_returned_messages`, `send_ack_interval_ms`) +pub type CanisterInitArgs = (u64, u64); + +pub const DEFAULT_CANISTER_INIT_ARGS: CanisterInitArgs = ( + DEFAULT_TEST_MAX_NUMBER_OF_RETURNED_MESSAGES, + DEFAULT_TEST_SEND_ACK_INTERVAL_MS, +); + lazy_static! { pub static ref TEST_ENV: Mutex = Mutex::new(TestEnv::new()); + static ref TEST_CANISTER_WASM_MODULE: Vec = match std::env::var("TEST_CANISTER_WASM_PATH") { + Ok(path) => load_canister_wasm_from_path(&PathBuf::from(path)), + Err(_) => load_canister_wasm_from_bin("test_canister.wasm"), + }; } pub fn get_test_env<'a>() -> MutexGuard<'a, TestEnv> { @@ -29,81 +41,43 @@ pub fn get_test_env<'a>() -> MutexGuard<'a, TestEnv> { pub struct TestEnv { pub pic: PocketIc, - pub canister_id: Principal, - canister_init_args: CanisterInitArgs, - wasm_module: Vec, + test_canister_id: Principal, root_ic_key: Vec, } -/// (`max_number_or_returned_messages`, `send_ack_interval_ms`) -type CanisterInitArgs = (u64, u64); - impl TestEnv { pub fn new() -> Self { - let pic = PocketIc::new(); + let pic = PocketIcBuilder::new() + // NNS subnet needed to retrieve the root key + .with_nns_subnet() + .with_application_subnet() + .build(); // set ic time to current time pic.set_time(SystemTime::now()); - let canister_id = pic.create_canister(None); + let app_subnet = pic.topology().get_app_subnets()[0]; + let canister_id = pic.create_canister_on_subnet(None, None, app_subnet); pic.add_cycles(canister_id, 1_000_000_000_000_000); - let wasm_bytes = match std::env::var("TEST_CANISTER_WASM_PATH") { - Ok(path) => load_canister_wasm_from_path(&PathBuf::from(path)), - Err(_) => load_canister_wasm_from_bin("test_canister.wasm"), - }; - - let arguments: CanisterInitArgs = ( - DEFAULT_TEST_MAX_NUMBER_OF_RETURNED_MESSAGES, - DEFAULT_TEST_SEND_ACK_INTERVAL_MS, - ); pic.install_canister( canister_id, - wasm_bytes.clone(), - candid::encode_args(arguments.clone()).unwrap(), + TEST_CANISTER_WASM_MODULE.clone(), + candid::encode_args(DEFAULT_CANISTER_INIT_ARGS).unwrap(), None, ); - let root_ic_key = pic.root_key(); + let root_ic_key = pic.root_key().unwrap(); Self { pic, - canister_id, - canister_init_args: arguments, - wasm_module: wasm_bytes, + test_canister_id: canister_id, root_ic_key, } } - pub fn reset_canister( - &mut self, - max_number_or_returned_messages: u64, - send_ack_interval_ms: u64, - ) { - let arguments: CanisterInitArgs = (max_number_or_returned_messages, send_ack_interval_ms); - let res = self.pic.reinstall_canister( - self.canister_id, - self.wasm_module.to_owned(), - candid::encode_args(arguments.clone()).unwrap(), - None, - ); - - match res { - Ok(_) => { - self.canister_init_args = arguments; - }, - Err(err) => { - panic!("Failed to reset canister: {:?}", err); - }, - } - } - - /// Resets the canister using the default parameters. See [reset_canister]. - pub fn reset_canister_with_default_params(&mut self) { - self.reset_canister( - DEFAULT_TEST_MAX_NUMBER_OF_RETURNED_MESSAGES, - DEFAULT_TEST_SEND_ACK_INTERVAL_MS, - ); + pub fn get_test_canister_id(&self) -> Principal { + self.test_canister_id } /// Returns the current time of the canister in nanoseconds. diff --git a/src/test_canister/src/lib.rs b/src/test_canister/src/lib.rs index b1ea705..01ac921 100644 --- a/src/test_canister/src/lib.rs +++ b/src/test_canister/src/lib.rs @@ -75,3 +75,11 @@ fn send(client_principal: ClientPrincipal, messages: Vec>) -> CanisterSe fn close(client_principal: ClientPrincipal) -> CanisterCloseResult { ic_websocket_cdk::close(client_principal) } + +// wipes the internal state +#[update] +fn wipe(max_number_of_returned_messages: usize, send_ack_interval_ms: u64) { + ic_websocket_cdk::wipe(); + + init(max_number_of_returned_messages, send_ack_interval_ms); +} diff --git a/src/test_canister/test_canister.did b/src/test_canister/test_canister.did index 2a04b7d..956c48a 100644 --- a/src/test_canister/test_canister.did +++ b/src/test_canister/test_canister.did @@ -23,4 +23,5 @@ service : (nat64, nat64) -> { // methods used just for debugging/testing "send" : (ClientPrincipal, vec blob) -> (CanisterSendResult); "close" : (ClientPrincipal) -> (CanisterCloseResult); + "wipe" : (nat64, nat64) -> (); }; From af5aa9aedf7f2b1ec85263434329e108248c8d19 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 15 Dec 2023 12:51:04 +0100 Subject: [PATCH 14/18] chore: fix integration tests' advance canister time --- .../src/tests/integration_tests/utils/test_env.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs b/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs index 310e3bd..4191078 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/utils/test_env.rs @@ -95,7 +95,10 @@ impl TestEnv { pub fn advance_canister_time_ms(&self, ms: u64) { self.pic.advance_time(Duration::from_millis(ms)); - // produce and advance by one block to fire eventual timers - self.pic.tick(); + // produce and advance by some blocks to fire eventual timers + // see https://forum.dfinity.org/t/pocketic-multi-subnet-canister-testing/24901/4 + for _ in 0..100 { + self.pic.tick(); + } } } From 1ad39163e965f88853591cc69175a5e6290b8677 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 15 Dec 2023 13:12:32 +0100 Subject: [PATCH 15/18] fix: map option instead of unwrap --- src/ic-websocket-cdk/src/state.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ic-websocket-cdk/src/state.rs b/src/ic-websocket-cdk/src/state.rs index 2ba4d0f..2b9c495 100644 --- a/src/ic-websocket-cdk/src/state.rs +++ b/src/ic-websocket-cdk/src/state.rs @@ -94,9 +94,9 @@ pub(crate) fn decrement_gateway_clients_count( let clients_count = g.decrement_clients_count(); if remove_if_empty && clients_count == 0 { - let g = map.remove(gateway_principal).unwrap(); - - return Some(g.messages_queue.iter().map(|m| m.key.clone()).collect()); + return map + .remove(gateway_principal) + .map(|g| g.messages_queue.iter().map(|m| m.key.clone()).collect()); } } From 858ebbcb622c3ddf80a4ee8a0ef4881674626ea0 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 15 Dec 2023 21:39:18 +0100 Subject: [PATCH 16/18] fix: remove gateways when empty and expired --- src/ic-websocket-cdk/src/state.rs | 83 ++++++++---- .../integration_tests/c_ws_get_messages.rs | 128 +++++++++++++++++- .../src/tests/integration_tests/d_ws_close.rs | 4 +- .../src/tests/unit_tests/mod.rs | 69 +++++++--- .../src/tests/unit_tests/utils.rs | 9 +- src/ic-websocket-cdk/src/timers.rs | 4 + src/ic-websocket-cdk/src/types.rs | 12 +- src/ic-websocket-cdk/src/utils.rs | 6 +- 8 files changed, 252 insertions(+), 63 deletions(-) diff --git a/src/ic-websocket-cdk/src/state.rs b/src/ic-websocket-cdk/src/state.rs index 2b9c495..b4ad4d2 100644 --- a/src/ic-websocket-cdk/src/state.rs +++ b/src/ic-websocket-cdk/src/state.rs @@ -2,6 +2,7 @@ use std::{ cell::RefCell, collections::{HashMap, HashSet}, rc::Rc, + time::Duration, }; use candid::{encode_one, Principal}; @@ -32,6 +33,8 @@ thread_local! { /* flexible */ pub(crate) static CERT_TREE: RefCell> = RefCell::new(RbTree::new()); /// Keeps track of the principals of the WS Gateways that poll the canister /* flexible */ pub(crate) static REGISTERED_GATEWAYS: RefCell> = RefCell::new(HashMap::new()); + /// Keeps track of the gateways that must be removed from the list of registered gateways in the next ack interval + /* flexible */ pub(crate) static GATEWAYS_TO_REMOVE: RefCell> = RefCell::new(HashMap::new()); /// The parameters passed in the CDK initialization /* flexible */ pub(crate) static PARAMS: RefCell = RefCell::new(WsInitParams::default()); } @@ -72,6 +75,10 @@ pub(crate) fn reset_internal_state() { /// Increments the clients connected count for the given gateway. /// If the gateway is not registered, a new entry is created with a clients connected count of 1. pub(crate) fn increment_gateway_clients_count(gateway_principal: GatewayPrincipal) { + GATEWAYS_TO_REMOVE.with(|state| { + state.borrow_mut().remove(&gateway_principal); + }); + REGISTERED_GATEWAYS.with(|map| { map.borrow_mut() .entry(gateway_principal) @@ -82,29 +89,54 @@ pub(crate) fn increment_gateway_clients_count(gateway_principal: GatewayPrincipa /// Decrements the clients connected count for the given gateway, if it exists. /// -/// If `remove_if_empty` is true, the gateway is removed from the list of registered gateways -/// if it has no clients connected. -pub(crate) fn decrement_gateway_clients_count( - gateway_principal: &GatewayPrincipal, - remove_if_empty: bool, -) { - let messages_keys_to_delete = REGISTERED_GATEWAYS.with(|map| { - let mut map = map.borrow_mut(); - if let Some(g) = map.get_mut(gateway_principal) { - let clients_count = g.decrement_clients_count(); - - if remove_if_empty && clients_count == 0 { - return map - .remove(gateway_principal) - .map(|g| g.messages_queue.iter().map(|m| m.key.clone()).collect()); - } - } +/// If the gateway has no more clients connected, it is added to the [GATEWAYS_TO_REMOVE] map, +/// in order to remove it in the next keep alive check. +pub(crate) fn decrement_gateway_clients_count(gateway_principal: &GatewayPrincipal) { + let is_empty = REGISTERED_GATEWAYS.with(|map| { + map.borrow_mut() + .get_mut(gateway_principal) + .is_some_and(|g| { + let clients_count = g.decrement_clients_count(); + clients_count == 0 + }) + }); + + if is_empty { + GATEWAYS_TO_REMOVE.with(|state| { + state + .borrow_mut() + .insert(gateway_principal.clone(), get_current_time()); + }); + } +} + +/// Removes the gateways that were added to the [GATEWAYS_TO_REMOVE] map +/// more than the ack interval ms time ago from the list of registered gateways +pub(crate) fn remove_empty_expired_gateways() { + let ack_interval_ms = get_params().send_ack_interval_ms; + let time = get_current_time(); + + let mut gateway_principals_to_remove: Vec = vec![]; - None + GATEWAYS_TO_REMOVE.with(|state| { + state.borrow_mut().retain(|gp, added_at| { + if Duration::from_nanos(time - *added_at) > Duration::from_millis(ack_interval_ms) { + gateway_principals_to_remove.push(gp.clone()); + false + } else { + true + } + }) }); - if let Some(messages_keys_to_delete) = messages_keys_to_delete { - delete_keys_from_cert_tree(messages_keys_to_delete); + for gateway_principal in &gateway_principals_to_remove { + if let Some(messages_keys_to_delete) = REGISTERED_GATEWAYS.with(|map| { + map.borrow_mut() + .remove(gateway_principal) + .map(|g| g.messages_queue.iter().map(|m| m.key.clone()).collect()) + }) { + delete_keys_from_cert_tree(messages_keys_to_delete); + } } } @@ -278,15 +310,11 @@ pub(crate) fn add_client(client_key: ClientKey, new_client: RegisteredClient) { increment_gateway_clients_count(new_client.gateway_principal); } -/// Removes a client from the internal state -/// and call the on_close callback, +/// Removes a client from the internal state and call the on_close callback, /// if the client was registered in the state. /// /// If a `close_reason` is provided, it also sends a close message to the client, /// so that the client can close the WS connection with the gateway. -/// -/// If a `close_reason` is **not** provided, it also removes the gateway from the state -/// if it has no clients connected anymore. pub(crate) fn remove_client(client_key: &ClientKey, close_reason: Option) { if let Some(close_reason) = close_reason.clone() { // ignore the error @@ -314,10 +342,7 @@ pub(crate) fn remove_client(client_key: &ClientKey, close_reason: Option = (1..=send_messages_count) + .map(|i| AppMessage { + text: format!("test{}", i), + }) + .collect(); + call_send_with_panic(&client_1_key.client_principal, messages_to_send.clone()); + + // check that gateway can receive the messages + helpers::assert_gateway_has_messages(send_messages_count); + + // disconnect the client and check that gateway can still receive the messages + call_ws_close_with_panic( + &GATEWAY_1, + CanisterWsCloseArguments { + client_key: client_1_key.clone(), + }, + ); + + // check that gateway can still receive the messages + helpers::assert_gateway_has_messages(send_messages_count); + + // wait for the ack interval to fire + get_test_env().advance_canister_time_ms(DEFAULT_TEST_SEND_ACK_INTERVAL_MS); + + // check that gateway can still receive the messages, even after the ack interval has fired + helpers::assert_gateway_has_messages(send_messages_count); + + // wait for the keep alive timeout to expire + get_test_env().advance_canister_time_ms(CLIENT_KEEP_ALIVE_TIMEOUT_MS); + + helpers::assert_gateway_has_no_messages(); +} + +#[test] +fn test_7_empty_gateway_can_get_messages_until_next_keep_alive_check_if_removed_before_ack_interval( +) { + let send_messages_count = 10; + // first, reset the canister + call_wipe(None); + // second, register client 1 + let client_1_key = CLIENT_1_KEY.deref(); + call_ws_open_for_client_key_with_panic(client_1_key); + // third, send a batch of messages to the client + let messages_to_send: Vec = (1..=send_messages_count) + .map(|i| AppMessage { + text: format!("test{}", i), + }) + .collect(); + call_send_with_panic(&client_1_key.client_principal, messages_to_send.clone()); + + // check that gateway can receive the messages + helpers::assert_gateway_has_messages(send_messages_count); + + // wait for the ack interval to fire + get_test_env().advance_canister_time_ms(DEFAULT_TEST_SEND_ACK_INTERVAL_MS); + + // disconnect the client and check that gateway can still receive the messages + call_ws_close_with_panic( + &GATEWAY_1, + CanisterWsCloseArguments { + client_key: client_1_key.clone(), + }, + ); + + let expected_messages_len = send_messages_count + 1; // +1 for the ack message + + // check that gateway can still receive the messages, even after the ack interval has fired + helpers::assert_gateway_has_messages(expected_messages_len); + + // wait for the keep alive timeout to expire + get_test_env().advance_canister_time_ms(CLIENT_KEEP_ALIVE_TIMEOUT_MS); + + // the gateway can still receive the messages, because it was emptied + // less than an ack interval ago + helpers::assert_gateway_has_messages(expected_messages_len); + + // wait for next ack interval to expire + get_test_env().advance_canister_time_ms(DEFAULT_TEST_SEND_ACK_INTERVAL_MS); + + // the gateway can still receive the messages, because empty expired gateways + // are removed only in the keep alive timeout callback + helpers::assert_gateway_has_messages(expected_messages_len); + + // wait for the keep alive timeout to expire + get_test_env().advance_canister_time_ms(CLIENT_KEEP_ALIVE_TIMEOUT_MS); + + helpers::assert_gateway_has_no_messages(); +} + +mod helpers { + use super::*; + + pub(super) fn assert_gateway_has_messages(send_messages_count: usize) { + let CanisterOutputCertifiedMessages { messages, .. } = call_ws_get_messages_with_panic( + &GATEWAY_1, + CanisterWsGetMessagesArguments { nonce: 0 }, + ); + assert_eq!( + messages.len(), + // + 1 for the open service message + send_messages_count + 1, + ); + } + + pub(super) fn assert_gateway_has_no_messages() { + let CanisterOutputCertifiedMessages { messages, .. } = call_ws_get_messages_with_panic( + &GATEWAY_1, + CanisterWsGetMessagesArguments { nonce: 0 }, + ); + assert_eq!(messages.len(), 0); + } +} diff --git a/src/ic-websocket-cdk/src/tests/integration_tests/d_ws_close.rs b/src/ic-websocket-cdk/src/tests/integration_tests/d_ws_close.rs index ec2c496..7abc148 100644 --- a/src/ic-websocket-cdk/src/tests/integration_tests/d_ws_close.rs +++ b/src/ic-websocket-cdk/src/tests/integration_tests/d_ws_close.rs @@ -115,8 +115,8 @@ fn test_4_should_close_the_websocket_for_a_registered_client() { assert_eq!( res, CanisterWsCloseResult::Err( - WsError::GatewayNotRegistered { - gateway_principal: GATEWAY_1.deref() + WsError::ClientKeyNotConnected { + client_key: &CLIENT_1_KEY, } .to_string() ) diff --git a/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs b/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs index 337644f..326fe77 100644 --- a/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs +++ b/src/ic-websocket-cdk/src/tests/unit_tests/mod.rs @@ -150,6 +150,14 @@ fn test_current_time() { assert!(get_current_time() >= timestamp_nanos); } +#[test] +fn test_remove_empty_expired_gateways_empty() { + let res = panic::catch_unwind(|| { + remove_empty_expired_gateways(); + }); + assert!(res.is_ok()); +} + proptest! { #[test] fn test_increment_gateway_clients_count_gateway_nonexistent(test_gateway_principal in any::().prop_map(|_| common::generate_random_principal())) { @@ -163,11 +171,14 @@ proptest! { fn test_increment_gateway_clients_count(test_gateway_principal in any::().prop_map(|_| common::generate_random_principal())) { // Set up REGISTERED_GATEWAYS.with(|n| n.borrow_mut().insert(test_gateway_principal, RegisteredGateway::new())); + GATEWAYS_TO_REMOVE.with(|state| state.borrow_mut().insert(test_gateway_principal, 0)); increment_gateway_clients_count(test_gateway_principal); let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_gateway_principal).cloned()).unwrap(); prop_assert_eq!(registered_gateway.connected_clients_count, 1); + let gateway_to_remove = GATEWAYS_TO_REMOVE.with(|map| map.borrow().get(&test_gateway_principal).cloned()); + prop_assert!(gateway_to_remove.is_none()); // change the registered gateway to see if its not replaced REGISTERED_GATEWAYS.with(|map| { @@ -179,6 +190,8 @@ proptest! { let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_gateway_principal).cloned()).unwrap(); prop_assert_eq!(registered_gateway.outgoing_message_nonce, 5); prop_assert_eq!(registered_gateway.connected_clients_count, 2); + let gateway_to_remove = GATEWAYS_TO_REMOVE.with(|map| map.borrow().get(&test_gateway_principal).cloned()); + prop_assert!(gateway_to_remove.is_none()); } #[test] @@ -195,34 +208,54 @@ proptest! { while clients_count > 0 { let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_gateway_principal).cloned()).unwrap(); prop_assert_eq!(registered_gateway.connected_clients_count, clients_count); - decrement_gateway_clients_count(&test_gateway_principal, false); + decrement_gateway_clients_count(&test_gateway_principal); clients_count -= 1; } let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_gateway_principal).cloned()); prop_assert!(registered_gateway.is_some()); + let gateway_timestamp = GATEWAYS_TO_REMOVE.with(|map| map.borrow().get(&test_gateway_principal).cloned()); + prop_assert!(gateway_timestamp.is_some()); } #[test] - fn test_decrement_gateway_clients_count_and_remove_gateway(test_gateway_principal in any::().prop_map(|_| common::generate_random_principal()), test_connected_clients_count in (1..1000u64)) { + fn test_remove_empty_expired_gateways_not_expired(test_gateway_principals in any::>().prop_map(|v| v.iter().map(|_|common::generate_random_principal()).collect::>())) { // Set up - REGISTERED_GATEWAYS.with(|n| { - let mut gw = RegisteredGateway::new(); - gw.connected_clients_count = test_connected_clients_count; - n.borrow_mut().insert(test_gateway_principal, gw); - }); + let gateway_timestamp = get_current_time() - (get_params().send_ack_interval_ms * 1_000_000) + 100_000_000; + for gateway_principal in test_gateway_principals.clone() { + REGISTERED_GATEWAYS.with(|n| n.borrow_mut().insert(gateway_principal, RegisteredGateway::new())); + GATEWAYS_TO_REMOVE.with(|state| state.borrow_mut().insert(gateway_principal, gateway_timestamp)); + } // Test - let mut clients_count = test_connected_clients_count; - while clients_count > 0 { - let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_gateway_principal).cloned()).unwrap(); - prop_assert_eq!(registered_gateway.connected_clients_count, clients_count); - decrement_gateway_clients_count(&test_gateway_principal, true); - clients_count -= 1; + remove_empty_expired_gateways(); + + for gateway_principal in test_gateway_principals { + let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&gateway_principal).cloned()); + prop_assert!(registered_gateway.is_some()); + let gateway_timestamp = GATEWAYS_TO_REMOVE.with(|map| map.borrow().get(&gateway_principal).cloned()); + prop_assert!(gateway_timestamp.is_some()); } + } - let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_gateway_principal).cloned()); - prop_assert!(registered_gateway.is_none()); + #[test] + fn test_remove_empty_expired_gateways(test_gateway_principals in any::>().prop_map(|v| v.iter().map(|_|common::generate_random_principal()).collect::>())) { + // Set up + let gateway_timestamp = get_current_time() - (get_params().send_ack_interval_ms * 1_000_000) - 100_000_000; + for gateway_principal in test_gateway_principals.clone() { + REGISTERED_GATEWAYS.with(|n| n.borrow_mut().insert(gateway_principal, RegisteredGateway::new())); + GATEWAYS_TO_REMOVE.with(|state| state.borrow_mut().insert(gateway_principal, gateway_timestamp)); + } + + // Test + remove_empty_expired_gateways(); + + for gateway_principal in test_gateway_principals { + let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&gateway_principal).cloned()); + prop_assert!(registered_gateway.is_none()); + let gateway_timestamp = GATEWAYS_TO_REMOVE.with(|map| map.borrow().get(&gateway_principal).cloned()); + prop_assert!(gateway_timestamp.is_none()); + } } #[test] @@ -585,8 +618,10 @@ proptest! { let registered_client = REGISTERED_CLIENTS.with(|map| map.borrow().get(&test_client_key).cloned()); prop_assert!(registered_client.is_none()); - let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_registered_client.gateway_principal).cloned()); - prop_assert!(registered_gateway.is_none()); + let registered_gateway = REGISTERED_GATEWAYS.with(|map| map.borrow().get(&test_registered_client.gateway_principal).cloned()).unwrap(); + prop_assert_eq!(registered_gateway.connected_clients_count, 0); + let gateway_timestamp = GATEWAYS_TO_REMOVE.with(|map| map.borrow().get(&test_registered_client.gateway_principal).cloned()); + prop_assert!(gateway_timestamp.is_some()); let incoming_seq_num = INCOMING_MESSAGE_FROM_CLIENT_NUM_MAP.with(|map| map.borrow().get(&test_client_key).cloned()); prop_assert!(incoming_seq_num.is_none()); diff --git a/src/ic-websocket-cdk/src/tests/unit_tests/utils.rs b/src/ic-websocket-cdk/src/tests/unit_tests/utils.rs index fd6b6da..0650c1f 100644 --- a/src/ic-websocket-cdk/src/tests/unit_tests/utils.rs +++ b/src/ic-websocket-cdk/src/tests/unit_tests/utils.rs @@ -1,12 +1,11 @@ -use candid::Principal; - use crate::{ - format_message_for_gateway_key, set_params, CanisterOutputMessage, ClientKey, GatewayPrincipal, - RegisteredClient, WsInitParams, REGISTERED_GATEWAYS, + format_message_for_gateway_key, set_params, tests::common::generate_random_principal, + CanisterOutputMessage, ClientKey, GatewayPrincipal, RegisteredClient, WsInitParams, + REGISTERED_GATEWAYS, }; pub(super) fn generate_random_registered_client() -> RegisteredClient { - RegisteredClient::new(Principal::anonymous()) + RegisteredClient::new(generate_random_principal()) } pub(super) fn add_messages_for_gateway( diff --git a/src/ic-websocket-cdk/src/timers.rs b/src/ic-websocket-cdk/src/timers.rs index 29d4b3d..e13f7aa 100644 --- a/src/ic-websocket-cdk/src/timers.rs +++ b/src/ic-websocket-cdk/src/timers.rs @@ -112,7 +112,11 @@ fn send_ack_to_clients_timer_callback() { /// Checks if the clients for which we are waiting for keep alive have sent a keep alive message. /// If a client has not sent a keep alive message, it is removed from the connected clients. +/// +/// Before checking the clients, it removes all the empty expired gateways from the list of registered gateways. fn check_keep_alive_timer_callback() { + remove_empty_expired_gateways(); + let client_keys_to_remove: Vec = CLIENTS_WAITING_FOR_KEEP_ALIVE .with(Rc::clone) .borrow() diff --git a/src/ic-websocket-cdk/src/types.rs b/src/ic-websocket-cdk/src/types.rs index 0a547fc..9d563be 100644 --- a/src/ic-websocket-cdk/src/types.rs +++ b/src/ic-websocket-cdk/src/types.rs @@ -78,7 +78,7 @@ pub struct CanisterWsGetMessagesArguments { pub(crate) struct WebsocketMessage { pub(crate) client_key: ClientKey, // The client that the gateway will forward the message to or that sent the message. pub(crate) sequence_num: u64, // Both ways, messages should arrive with sequence numbers 0, 1, 2... - pub(crate) timestamp: u64, // Timestamp of when the message was made for the recipient to inspect. + pub(crate) timestamp: TimestampNs, // Timestamp of when the message was made for the recipient to inspect. pub(crate) is_service_message: bool, // Whether the message is a service message sent by the CDK to the client or vice versa. #[serde(with = "serde_bytes")] pub(crate) content: Vec, // Application message encoded in binary. @@ -132,9 +132,11 @@ pub(crate) struct MessagesForGatewayRange { pub is_end_of_queue: bool, } +pub(crate) type TimestampNs = u64; + #[derive(Clone, Debug, Default, Eq, PartialEq)] pub(crate) struct MessageToDelete { - timestamp: u64, + timestamp: TimestampNs, } pub(crate) type GatewayPrincipal = Principal; @@ -185,7 +187,7 @@ impl RegisteredGateway { pub(crate) fn add_message_to_queue( &mut self, message: CanisterOutputMessage, - message_timestamp: u64, + message_timestamp: TimestampNs, ) { self.messages_queue.push_back(message.clone()); self.messages_to_delete.push_back(MessageToDelete { @@ -228,7 +230,7 @@ impl RegisteredGateway { /// The metadata about a registered client. #[derive(Clone, Debug, Eq, PartialEq)] pub(crate) struct RegisteredClient { - pub(crate) last_keep_alive_timestamp: u64, + pub(crate) last_keep_alive_timestamp: TimestampNs, pub(crate) gateway_principal: GatewayPrincipal, } @@ -242,7 +244,7 @@ impl RegisteredClient { } /// Gets the last keep alive timestamp. - pub(crate) fn get_last_keep_alive_timestamp(&self) -> u64 { + pub(crate) fn get_last_keep_alive_timestamp(&self) -> TimestampNs { self.last_keep_alive_timestamp } diff --git a/src/ic-websocket-cdk/src/utils.rs b/src/ic-websocket-cdk/src/utils.rs index bde2e72..6d0a2ff 100644 --- a/src/ic-websocket-cdk/src/utils.rs +++ b/src/ic-websocket-cdk/src/utils.rs @@ -1,6 +1,8 @@ #[cfg(not(test))] use ic_cdk::api::time; +use crate::types::TimestampNs; + #[macro_export] macro_rules! custom_print { ($($arg:tt)*) => { @@ -29,7 +31,7 @@ macro_rules! custom_trap { } } -pub(crate) fn get_current_time() -> u64 { +pub(crate) fn get_current_time() -> TimestampNs { #[cfg(test)] { use std::time::SystemTime; @@ -37,7 +39,7 @@ pub(crate) fn get_current_time() -> u64 { .duration_since(SystemTime::UNIX_EPOCH) .unwrap(); let timestamp_nanos = duration_since_epoch.as_nanos(); - timestamp_nanos as u64 + timestamp_nanos as TimestampNs } #[cfg(not(test))] { From 27ef65e5e6fc662b661c596955f2fac380d9b9f4 Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 15 Dec 2023 21:44:40 +0100 Subject: [PATCH 17/18] fix: clear map on reset state --- src/ic-websocket-cdk/src/state.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/ic-websocket-cdk/src/state.rs b/src/ic-websocket-cdk/src/state.rs index b4ad4d2..1af4f75 100644 --- a/src/ic-websocket-cdk/src/state.rs +++ b/src/ic-websocket-cdk/src/state.rs @@ -70,6 +70,9 @@ pub(crate) fn reset_internal_state() { REGISTERED_GATEWAYS.with(|map| { map.borrow_mut().clear(); }); + GATEWAYS_TO_REMOVE.with(|map| { + map.borrow_mut().clear(); + }); } /// Increments the clients connected count for the given gateway. From aaf57614ea39eab6e0af6525d009b97cb120312c Mon Sep 17 00:00:00 2001 From: Luca8991 Date: Fri, 15 Dec 2023 22:23:02 +0100 Subject: [PATCH 18/18] chore: bump to version v0.3.2 --- Cargo.lock | 2 +- src/ic-websocket-cdk/Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e4edbad..1745e84 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1001,7 +1001,7 @@ dependencies = [ [[package]] name = "ic-websocket-cdk" -version = "0.3.1" +version = "0.3.2" dependencies = [ "base64 0.21.5", "candid", diff --git a/src/ic-websocket-cdk/Cargo.toml b/src/ic-websocket-cdk/Cargo.toml index af582e2..99173b5 100644 --- a/src/ic-websocket-cdk/Cargo.toml +++ b/src/ic-websocket-cdk/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "ic-websocket-cdk" -version = "0.3.1" +version = "0.3.2" edition.workspace = true rust-version.workspace = true repository.workspace = true