Pass message policy to Client::new; add default impl for make_chain_client. (#4045)

## Motivation

`Client::new` uses a default value for the message policy, and we
explicitly set it each time `make_chain_client` is called.

## Proposal

Pass the correct policy into `Client::new` directly, so we can have a
default impl for `make_chain_client`.

## Test Plan

CI

## Release Plan

- Nothing to do / These changes follow the usual release cycle.

## Links

- [reviewer
checklist](https://github.com/linera-io/linera-protocol/blob/main/CONTRIBUTING.md#reviewer-checklist)

Author: afck

linera-client/src/chain_listener.rs

@@ -28,7 +28,10 @@ use linera_storage::{Clock as _, Storage as _};
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, info, instrument, warn, Instrument as _};
 
-use crate::{wallet::Wallet, Error};
+use crate::{
+    wallet::{UserChain, Wallet},
+    Error,
+};
 
 #[derive(Debug, Default, Clone, clap::Args)]
 pub struct ChainListenerConfig {
@@ -68,9 +71,22 @@ pub trait ClientContext {
 
     fn storage(&self) -> &<Self::Environment as linera_core::Environment>::Storage;
 
-    fn make_chain_client(&self, chain_id: ChainId) -> ContextChainClient<Self>;
-
-    fn client(&self) -> &linera_core::client::Client<Self::Environment>;
+    fn client(&self) -> &Arc<linera_core::client::Client<Self::Environment>>;
+
+    fn make_chain_client(&self, chain_id: ChainId) -> ChainClient<Self::Environment> {
+        let chain = self
+            .wallet()
+            .get(chain_id)
+            .cloned()
+            .unwrap_or_else(|| UserChain::make_other(chain_id, Timestamp::from(0)));
+        self.client().create_chain_client(
+            chain_id,
+            chain.block_hash,
+            chain.next_block_height,
+            chain.pending_proposal,
+            chain.owner,
+        )
+    }
 
     async fn update_wallet_for_new_chain(
         &mut self,

linera-client/src/client_context.rs

@@ -3,7 +3,7 @@
 
 #[cfg(with_testing)]
 use std::num::NonZeroUsize;
-use std::{collections::HashSet, sync::Arc};
+use std::sync::Arc;
 
 use futures::Future;
 use linera_base::{
@@ -15,7 +15,7 @@ use linera_base::{
 };
 use linera_chain::types::ConfirmedBlockCertificate;
 use linera_core::{
-    client::{BlanketMessagePolicy, ChainClient, Client, MessagePolicy, PendingProposal},
+    client::{ChainClient, Client, MessagePolicy},
     data_types::{ChainInfoQuery, ClientOutcome},
     join_set_ext::JoinSet,
     node::{CrossChainMessageDelivery, ValidatorNode},
@@ -71,8 +71,6 @@ pub struct ClientContext<Env: Environment, W> {
     pub retry_delay: Duration,
     pub max_retries: u32,
     pub chain_listeners: JoinSet,
-    pub blanket_message_policy: BlanketMessagePolicy,
-    pub restrict_chain_ids_to: Option<HashSet<ChainId>>,
 }
 
 impl<Env: Environment, W> chain_listener::ClientContext for ClientContext<Env, W>
@@ -89,24 +87,7 @@ where
         self.client.storage_client()
     }
 
-    fn make_chain_client(&self, chain_id: ChainId) -> ChainClient<Env> {
-        // We only create clients for chains we have in the wallet, or for the admin chain.
-        let chain = self
-            .wallet
-            .get(chain_id)
-            .cloned()
-            .unwrap_or_else(|| UserChain::make_other(chain_id, Timestamp::from(0)));
-
-        self.make_chain_client_internal(
-            chain_id,
-            chain.block_hash,
-            chain.next_block_height,
-            chain.pending_proposal,
-            chain.owner,
-        )
-    }
-
-    fn client(&self) -> &Client<Env> {
+    fn client(&self) -> &Arc<Client<Env>> {
         &self.client
     }
 
@@ -153,6 +134,10 @@ where
             },
             options.max_pending_message_bundles,
             wallet.genesis_admin_chain(),
+            MessagePolicy::new(
+                options.blanket_message_policy,
+                options.restrict_chain_ids_to,
+            ),
             delivery,
             options.long_lived_services,
             chain_ids,
@@ -170,8 +155,6 @@ where
             retry_delay: options.retry_delay,
             max_retries: options.max_retries,
             chain_listeners: JoinSet::default(),
-            blanket_message_policy: options.blanket_message_policy,
-            restrict_chain_ids_to: options.restrict_chain_ids_to,
         }
     }
 
@@ -204,6 +187,7 @@ where
             },
             10,
             wallet.genesis_admin_chain(),
+            MessagePolicy::new_accept_all(),
             delivery,
             false,
             chain_ids,
@@ -221,8 +205,6 @@ where
             retry_delay,
             max_retries,
             chain_listeners: JoinSet::default(),
-            blanket_message_policy: BlanketMessagePolicy::Accept,
-            restrict_chain_ids_to: None,
         }
     }
 }
@@ -267,28 +249,6 @@ impl<Env: Environment, W: Persist<Target = Wallet>> ClientContext<Env, W> {
             .expect("No non-admin chain specified in wallet with no non-admin chain")
     }
 
-    fn make_chain_client_internal(
-        &self,
-        chain_id: ChainId,
-        block_hash: Option<CryptoHash>,
-        next_block_height: BlockHeight,
-        pending_proposal: Option<PendingProposal>,
-        preferred_owner: Option<AccountOwner>,
-    ) -> ChainClient<Env> {
-        let mut chain_client = self.client.create_chain_client(
-            chain_id,
-            block_hash,
-            next_block_height,
-            pending_proposal,
-            preferred_owner,
-        );
-        chain_client.options_mut().message_policy = MessagePolicy::new(
-            self.blanket_message_policy,
-            self.restrict_chain_ids_to.clone(),
-        );
-        chain_client
-    }
-
     pub fn make_node_provider(&self) -> NodeProvider {
         NodeProvider::new(self.make_node_options())
     }
@@ -907,7 +867,7 @@ where
                 benchmark_chains.insert(chain_id, pub_key.into());
                 self.client.track_chain(chain_id);
 
-                let mut chain_client = self.make_chain_client_internal(
+                let mut chain_client = self.client.create_chain_client(
                     chain_id,
                     None,
                     BlockHeight::ZERO,

linera-client/src/unit_tests/chain_listener.rs

@@ -13,7 +13,7 @@ use linera_base::{
     ownership::{ChainOwnership, TimeoutConfig},
 };
 use linera_core::{
-    client::{ChainClient, Client},
+    client::{ChainClient, Client, MessagePolicy},
     environment,
     node::CrossChainMessageDelivery,
     test_utils::{MemoryStorageBuilder, StorageBuilder as _, TestBuilder},
@@ -46,25 +46,10 @@ impl chain_listener::ClientContext for ClientContext {
         self.client.storage_client()
     }
 
-    fn client(&self) -> &linera_core::client::Client<Self::Environment> {
+    fn client(&self) -> &Arc<linera_core::client::Client<Self::Environment>> {
         &self.client
     }
 
-    fn make_chain_client(&self, chain_id: ChainId) -> ChainClient<environment::Test> {
-        let chain = self
-            .wallet
-            .get(chain_id)
-            .cloned()
-            .unwrap_or_else(|| UserChain::make_other(chain_id, Timestamp::from(0)));
-        self.client.create_chain_client(
-            chain_id,
-            chain.block_hash,
-            chain.next_block_height,
-            chain.pending_proposal.clone(),
-            chain.owner,
-        )
-    }
-
     async fn update_wallet_for_new_chain(
         &mut self,
         chain_id: ChainId,
@@ -129,6 +114,7 @@ async fn test_chain_listener() -> anyhow::Result<()> {
             },
             10,
             admin_id,
+            MessagePolicy::new_accept_all(),
             delivery,
             false,
             [chain_id0],
@@ -217,6 +203,7 @@ async fn test_chain_listener_admin_chain() -> anyhow::Result<()> {
             },
             10,
             admin_id,
+            MessagePolicy::new_accept_all(),
             delivery,
             false,
             [],

linera-core/src/client/mod.rs

@@ -181,6 +181,7 @@ impl<Env: Environment> Client<Env> {
         environment: Env,
         max_pending_message_bundles: usize,
         admin_id: ChainId,
+        message_policy: MessagePolicy,
         cross_chain_message_delivery: CrossChainMessageDelivery,
         long_lived_services: bool,
         tracked_chains: impl IntoIterator<Item = ChainId>,
@@ -207,7 +208,7 @@ impl<Env: Environment> Client<Env> {
             chains: DashMap::new(),
             max_pending_message_bundles,
             admin_id,
-            message_policy: MessagePolicy::new(BlanketMessagePolicy::Accept, None),
+            message_policy,
             cross_chain_message_delivery,
             grace_period,
             tracked_chains,
@@ -1259,6 +1260,14 @@ impl MessagePolicy {
         }
     }
 
+    #[cfg(with_testing)]
+    pub fn new_accept_all() -> Self {
+        Self {
+            blanket: BlanketMessagePolicy::Accept,
+            restrict_chain_ids_to: None,
+        }
+    }
+
     #[instrument(level = "trace", skip(self))]
     fn must_handle(&self, bundle: &mut IncomingBundle) -> bool {
         if self.is_reject() {

linera-core/src/unit_tests/test_utils.rs

@@ -49,7 +49,7 @@ use {
 };
 
 use crate::{
-    client::Client,
+    client::{Client, MessagePolicy},
     data_types::*,
     node::{
         CrossChainMessageDelivery, NodeError, NotificationStream, ValidatorNode,
@@ -952,6 +952,7 @@ where
             },
             10,
             self.admin_id(),
+            MessagePolicy::new_accept_all(),
             CrossChainMessageDelivery::NonBlocking,
             false,
             [chain_id],

linera-faucet/server/src/tests.rs

@@ -36,7 +36,7 @@ impl chain_listener::ClientContext for ClientContext {
         self.client.storage_client()
     }
 
-    fn client(&self) -> &linera_core::client::Client<environment::Test> {
+    fn client(&self) -> &Arc<linera_core::client::Client<environment::Test>> {
         unimplemented!()
     }
 

linera-service/src/schema_export.rs

@@ -1,6 +1,8 @@
 // Copyright (c) Zefchain Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
+use std::sync::Arc;
+
 use linera_base::{
     crypto::CryptoHash,
     data_types::{BlobContent, NetworkDescription, Timestamp},
@@ -180,11 +182,7 @@ impl ClientContext for DummyContext {
         unimplemented!()
     }
 
-    fn client(&self) -> &linera_core::client::Client<Self::Environment> {
-        unimplemented!()
-    }
-
-    fn make_chain_client(&self, _: ChainId) -> ChainClient<Self::Environment> {
+    fn client(&self) -> &Arc<linera_core::client::Client<Self::Environment>> {
         unimplemented!()
     }