diff --git a/Cargo.lock b/Cargo.lock index d8e7775..70027da 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -266,6 +266,7 @@ dependencies = [ "atoma-demo", "linera-sdk", "proptest", + "rand", "serde", "serde_json", "test-log", diff --git a/Cargo.toml b/Cargo.toml index d52ebee..caaad83 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,7 @@ test-strategy = { version = "0.4.0", optional = true } [dev-dependencies] atoma-demo = { path = ".", features = ["test"] } linera-sdk = { git = "https://github.com/jvff/linera-protocol", rev = "26a5299", features = ["test"] } +rand = "0.8.5" [target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies] linera-sdk = { git = "https://github.com/jvff/linera-protocol", rev = "26a5299", features = ["test", "wasmer", "unstable-oracles"] } diff --git a/src/contract.rs b/src/contract.rs index 4235275..e865a1b 100644 --- a/src/contract.rs +++ b/src/contract.rs @@ -8,7 +8,7 @@ mod state; #[path = "./contract_unit_tests.rs"] mod tests; -use atoma_demo::{ChatInteraction, Operation}; +use atoma_demo::{ChatInteraction, Operation, PublicKey}; use linera_sdk::{ base::WithContractAbi, views::{RootView, View}, @@ -43,9 +43,10 @@ impl Contract for ApplicationContract { async fn instantiate(&mut self, _argument: Self::InstantiationArgument) {} async fn execute_operation(&mut self, operation: Self::Operation) -> Self::Response { - let Operation::LogChatInteraction { interaction } = operation; - - self.log_chat_interaction(interaction); + match operation { + Operation::UpdateNodes { add, remove } => self.update_nodes(add, remove), + Operation::LogChatInteraction { interaction } => self.log_chat_interaction(interaction), + } } async fn execute_message(&mut self, _message: Self::Message) {} @@ -56,6 +57,47 @@ impl Contract for ApplicationContract { } impl ApplicationContract { + /// Handles an [`Operation::UpdateNodes`] by adding the `nodes_to_add` and removing the + /// `nodes_to_remove`. + fn update_nodes(&mut self, nodes_to_add: Vec, nodes_to_remove: Vec) { + assert!( + self.runtime.chain_id() == self.runtime.application_id().creation.chain_id, + "Only the chain that created the application can manage the set of active nodes" + ); + + Self::assert_key_sets_are_disjoint(&nodes_to_add, &nodes_to_remove); + + for node in nodes_to_remove { + self.state + .active_atoma_nodes + .remove(&node) + .expect("Failed to remove a node from the set of active Atoma nodes"); + } + + for node in nodes_to_add { + self.state + .active_atoma_nodes + .insert(&node) + .expect("Failed to add a node to the set of active Atoma nodes"); + } + } + + /// Checks if two sets of [`PublicKey`]s are disjoint. + fn assert_key_sets_are_disjoint(left: &[PublicKey], right: &[PublicKey]) { + let (smallest_set, largest_set) = if left.len() < right.len() { + (left, right) + } else { + (right, left) + }; + + let disjoint = largest_set.iter().all(|key| !smallest_set.contains(key)); + + assert!( + disjoint, + "Conflicting request to add and remove the same node" + ); + } + /// Handles an [`Operation::LogChatInteraction`] by adding a [`ChatInteraction`] to the chat /// log. fn log_chat_interaction(&mut self, interaction: ChatInteraction) { diff --git a/src/contract_unit_tests.rs b/src/contract_unit_tests.rs index f8605a0..c0ce543 100644 --- a/src/contract_unit_tests.rs +++ b/src/contract_unit_tests.rs @@ -1,12 +1,101 @@ // Copyright (c) Zefchain Labs, Inc. // SPDX-License-Identifier: Apache-2.0 -use atoma_demo::{ChatInteraction, Operation}; -use linera_sdk::{util::BlockingWait, Contract, ContractRuntime}; +use std::{ + collections::{BTreeSet, HashSet}, + iter, panic, +}; + +use atoma_demo::{ChatInteraction, Operation, PublicKey}; +use linera_sdk::{ + base::{ApplicationId, ChainId}, + util::BlockingWait, + Contract, ContractRuntime, +}; +use proptest::{ + prelude::{Arbitrary, BoxedStrategy}, + sample::size_range, + strategy::Strategy, +}; +use rand::Rng; use test_strategy::proptest; use super::ApplicationContract; +/// Tests if nodes can be added to and removed from the set of active Atoma nodes. +#[proptest] +fn updating_nodes( + application_id: ApplicationId, + test_operations: TestUpdateNodesOperations, +) { + let mut test = NodeSetTest::new(application_id); + + for test_operation in test_operations.0 { + let operation = test.prepare_operation(test_operation); + + test.contract.execute_operation(operation).blocking_wait(); + + test.check_active_atoma_nodes(); + } +} + +/// Tests if the set of active Atoma nodes can only be changed in the chain where the application +/// was created. +#[proptest] +fn only_creation_chain_can_track_nodes( + application_id: ApplicationId, + chain_id: ChainId, + test_operation: TestUpdateNodesOperation, +) { + let result = panic::catch_unwind(move || { + let mut test = NodeSetTest::new(application_id).with_chain_id(chain_id); + let operation = test.prepare_operation(test_operation); + + test.contract.execute_operation(operation).blocking_wait(); + + test + }); + + match result { + Ok(test) => { + assert_eq!( + chain_id, application_id.creation.chain_id, + "Contract executed `Operation::UpdateNodes` \ + outside of the application's creation chain" + ); + test.check_active_atoma_nodes(); + } + Err(_panic_cause) => { + assert_ne!( + chain_id, application_id.creation.chain_id, + "Contract failed to execute `Operation::UpdateNodes` \ + on the application's creation chain" + ); + } + } +} + +/// Tests if the contract rejects adding a node twice. +#[proptest] +fn cant_add_and_remove_node_in_the_same_operation( + application_id: ApplicationId, + #[any(size_range(1..5).lift())] conflicting_nodes: HashSet, + mut test_operation: TestUpdateNodesOperation, +) { + let result = panic::catch_unwind(move || { + let mut test = NodeSetTest::new(application_id); + + test_operation.add.extend(conflicting_nodes.iter().copied()); + test_operation.remove.extend(conflicting_nodes); + + let operation = test.prepare_operation(test_operation); + + test.contract.execute_operation(operation).blocking_wait(); + }); + + assert!(result.is_err()); +} + /// Tests if chat interactions are logged on chain. #[proptest] fn chat_interactions_are_logged_on_chain(interactions: Vec) { @@ -34,3 +123,174 @@ fn setup_contract() -> ApplicationContract { ApplicationContract::load(runtime).blocking_wait() } + +/// Helper type with shared code for active Atoma node set tests. +pub struct NodeSetTest { + contract: ApplicationContract, + expected_nodes: HashSet, +} + +impl NodeSetTest { + /// Creates a new [`NodeSetTest`], setting up the contract and the runtime. + /// + /// The test configures the contract to run on the specified `chain_id`, or on the + /// application's creation chain if it's [`None`]. + pub fn new(application_id: ApplicationId) -> Self { + let mut contract = setup_contract(); + let chain_id = application_id.creation.chain_id; + + contract.runtime.set_application_id(application_id); + contract.runtime.set_chain_id(chain_id); + + NodeSetTest { + contract, + expected_nodes: HashSet::new(), + } + } + + /// Changes the [`ChainId`] for the chain that executes the contract. + pub fn with_chain_id(mut self, chain_id: ChainId) -> Self { + self.contract.runtime.set_chain_id(chain_id); + self + } + + /// Prepares an [`Operation::UpdateNodes`] based on the configured + /// [`TestUpdateNodesOperation`]. + /// + /// Updates the expected active Atoma nodes state to reflect the execution of the operation. + pub fn prepare_operation(&mut self, test_operation: TestUpdateNodesOperation) -> Operation { + let nodes_to_add = test_operation.add; + let nodes_to_remove = test_operation.remove; + + for node_to_remove in &nodes_to_remove { + self.expected_nodes.remove(node_to_remove); + } + + self.expected_nodes.extend(nodes_to_add.iter().copied()); + + Operation::UpdateNodes { + add: nodes_to_add, + remove: nodes_to_remove, + } + } + + /// Asserts that the contract's state has exactly the same nodes as the expected nodes. + pub fn check_active_atoma_nodes(&self) { + let node_count = self + .contract + .state + .active_atoma_nodes + .count() + .blocking_wait() + .expect("Failed to read active Atoma node set size"); + + let mut active_nodes = HashSet::with_capacity(node_count); + self.contract + .state + .active_atoma_nodes + .for_each_index(|node| { + assert!( + active_nodes.insert(node), + "`SetView` should not have duplicate elements" + ); + Ok(()) + }) + .blocking_wait() + .expect("Failed to read active Atoma nodes from state"); + + assert_eq!(node_count, self.expected_nodes.len()); + assert_eq!(active_nodes, self.expected_nodes); + } +} + +/// A list of test configurations for a sequence of [`Operation::UpdateNodes`]. +#[derive(Clone, Debug)] +pub struct TestUpdateNodesOperations(Vec); + +impl Arbitrary for TestUpdateNodesOperations { + type Parameters = (); + type Strategy = BoxedStrategy; + + /// Creates an arbitrary [`TestUpdateNodesOperations`]. + /// + /// This is done by creating a random set of nodes, and partitioning it into an arbitrary + /// number of operations. After an operation X adds a node for the first time, that node can be + /// removed at an operation Y > X, and then re-added at an operation Z > Y, and so on. + fn arbitrary_with((): Self::Parameters) -> Self::Strategy { + (1_usize..10) + .prop_flat_map(|operation_count| { + let node_keys = BTreeSet::::arbitrary_with(size_range(1..100).lift()) + .prop_map(Vec::from_iter) + .prop_shuffle(); + + node_keys.prop_perturb(move |node_keys, mut random| { + let mut add_operations = iter::repeat(vec![]) + .take(operation_count) + .collect::>(); + let mut remove_operations = add_operations.clone(); + + for node_key in node_keys { + let mut is_active = false; + let mut index = 0; + + random.gen_range(0..operation_count); + + while index < operation_count { + if is_active { + remove_operations[index].push(node_key); + } else { + add_operations[index].push(node_key); + } + + is_active = !is_active; + index = random.gen_range(index..operation_count) + 1; + } + } + + TestUpdateNodesOperations( + add_operations + .into_iter() + .zip(remove_operations) + .map(|(add, remove)| TestUpdateNodesOperation { add, remove }) + .collect(), + ) + }) + }) + .boxed() + } +} + +/// The test configuration for an [`Operation::UpdateNodes`]. +#[derive(Clone, Debug)] +pub struct TestUpdateNodesOperation { + add: Vec, + remove: Vec, +} + +impl Arbitrary for TestUpdateNodesOperation { + type Parameters = (); + type Strategy = BoxedStrategy; + + /// Creates an arbitrary [`TestUpdateNodesOperation`]. + /// + /// This is done by creating a random set of nodes, and splitting it in two, one with the nodes + /// to add and one with the nodes to remove. + fn arbitrary_with((): Self::Parameters) -> Self::Strategy { + (..20_usize, ..20_usize) + .prop_flat_map(|(add_count, remove_count)| { + BTreeSet::::arbitrary_with(size_range(add_count + remove_count).lift()) + .prop_map(Vec::from_iter) + .prop_shuffle() + .prop_map(move |mut node_keys| { + let nodes_to_remove = node_keys.split_off(add_count); + let nodes_to_add = node_keys; + + TestUpdateNodesOperation { + add: nodes_to_add, + remove: nodes_to_remove, + } + }) + }) + .boxed() + } +} diff --git a/src/lib.rs b/src/lib.rs index 4a1d52d..30654fd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,6 +19,12 @@ impl ServiceAbi for ApplicationAbi { /// Operations that the contract can execute. #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] pub enum Operation { + /// Update the set of active Atoma nodes. + UpdateNodes { + add: Vec, + remove: Vec, + }, + /// Log an interaction with the AI. LogChatInteraction { interaction: ChatInteraction }, } @@ -32,3 +38,9 @@ pub struct ChatInteraction { #[cfg_attr(feature = "test", strategy("[A-Za-z0-9., ]*"))] pub response: String, } + +/// Representation of an Atoma node's public key. +#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Deserialize, Serialize)] +#[cfg_attr(feature = "test", derive(test_strategy::Arbitrary))] +pub struct PublicKey([u8; 32]); +async_graphql::scalar!(PublicKey); diff --git a/src/state.rs b/src/state.rs index 4aae513..c0757f1 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,11 +1,12 @@ // Copyright (c) Zefchain Labs, Inc. // SPDX-License-Identifier: Apache-2.0 -use atoma_demo::ChatInteraction; -use linera_sdk::views::{linera_views, LogView, RootView, ViewStorageContext}; +use atoma_demo::{ChatInteraction, PublicKey}; +use linera_sdk::views::{linera_views, LogView, RootView, SetView, ViewStorageContext}; #[derive(RootView, async_graphql::SimpleObject)] #[view(context = "ViewStorageContext")] pub struct Application { + pub active_atoma_nodes: SetView, pub chat_log: LogView, } diff --git a/tests/chat_transcript.rs b/tests/chat_transcript.rs index 95425be..343278c 100644 --- a/tests/chat_transcript.rs +++ b/tests/chat_transcript.rs @@ -57,7 +57,10 @@ async fn service_queries_atoma() { let Operation::LogChatInteraction { interaction: ChatInteraction { response, .. }, - } = operation; + } = operation + else { + panic!("Unexpected operation returned from service"); + }; assert!(response.contains("Rio de Janeiro")); }