feat(client-cli): add a 'aggregator-discovery' command

Author: jpraynaud

internal/mithril-aggregator-discovery/src/model.rs

@@ -25,6 +25,12 @@ impl MithrilNetwork {
     }
 }
 
+impl From<String> for MithrilNetwork {
+    fn from(name: String) -> Self {
+        MithrilNetwork::new(name)
+    }
+}
+
 /// Representation of an aggregator endpoint
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct AggregatorEndpoint {
@@ -57,3 +63,9 @@ impl From<AggregatorEndpoint> for String {
         endpoint.url
     }
 }
+
+impl std::fmt::Display for AggregatorEndpoint {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.url)
+    }
+}

mithril-client-cli/src/commands/tools/aggregator_discovery.rs

@@ -0,0 +1,197 @@
+use clap::Parser;
+
+use mithril_client::{
+    AggregatorDiscoveryType, ClientBuilder, MithrilNetwork, MithrilResult,
+    RequiredAggregatorCapabilities,
+    common::{AggregateSignatureType, SignedEntityTypeDiscriminants},
+};
+
+/// Clap command to select an aggregator from the available ones with automatic discovery.
+#[derive(Parser, Debug, Clone)]
+pub struct AggregatorSelectCommand {
+    /// Path to the Cardano node database directory.
+    #[clap(long)]
+    network: MithrilNetwork,
+
+    /// Maximum number of entries to retrieve
+    #[clap(long, default_value_t = 1)]
+    max_entries: usize,
+
+    /// Signed entity types to consider for the discovery
+    ///
+    /// If not provided, all signed entity types are considered.
+    #[clap(long, value_parser, num_args = 0.., value_delimiter = ',')]
+    signed_entity_types: Vec<SignedEntityTypeDiscriminants>,
+
+    /// Aggregate signature types to consider for the discovery
+    ///
+    /// If not provided, all aggregate signature types are considered.
+    #[clap(long, value_parser, num_args = 0.., value_delimiter = ',')]
+    aggregate_signature_types: Vec<AggregateSignatureType>,
+}
+
+impl AggregatorSelectCommand {
+    /// Main command execution
+    pub async fn execute(&self) -> MithrilResult<()> {
+        let required_capabilities = self.build_required_capabilities();
+        let client_builder =
+            ClientBuilder::new(AggregatorDiscoveryType::Automatic(self.network.clone()))
+                .with_capabilities(required_capabilities)
+                .with_default_aggregator_discoverer();
+        let aggregator_endpoints = client_builder
+            .discover_aggregator(&self.network)?
+            .take(self.max_entries);
+
+        println!(
+            "Discovering at most {} aggregator endpoints:",
+            self.max_entries,
+        );
+        let mut found_aggregators = 0;
+        for endpoint in aggregator_endpoints {
+            println!("- Found: {endpoint}");
+            found_aggregators += 1;
+        }
+        if found_aggregators == 0 {
+            println!("- No aggregator endpoint found matching the requirements.");
+        }
+
+        Ok(())
+    }
+
+    fn build_required_capabilities(&self) -> RequiredAggregatorCapabilities {
+        if self.signed_entity_types.is_empty() && self.aggregate_signature_types.is_empty() {
+            return RequiredAggregatorCapabilities::All;
+        }
+
+        let mut required_capabilities = vec![];
+        if !self.signed_entity_types.is_empty() {
+            let mut required_capabilities_signed_entity_types = vec![];
+            for signed_entity_type in &self.signed_entity_types {
+                required_capabilities_signed_entity_types.push(
+                    RequiredAggregatorCapabilities::SignedEntityType(*signed_entity_type),
+                );
+            }
+            required_capabilities.push(RequiredAggregatorCapabilities::Or(
+                required_capabilities_signed_entity_types,
+            ));
+        }
+
+        if !self.aggregate_signature_types.is_empty() {
+            let mut required_capabilities_aggregate_signature_types = vec![];
+            for aggregate_signature_type in &self.aggregate_signature_types {
+                required_capabilities_aggregate_signature_types.push(
+                    RequiredAggregatorCapabilities::AggregateSignatureType(
+                        *aggregate_signature_type,
+                    ),
+                );
+            }
+            required_capabilities.push(RequiredAggregatorCapabilities::Or(
+                required_capabilities_aggregate_signature_types,
+            ));
+        }
+        if required_capabilities.len() == 1 {
+            return required_capabilities.into_iter().next().unwrap();
+        } else {
+            return RequiredAggregatorCapabilities::And(required_capabilities);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use mithril_client::common::SignedEntityTypeDiscriminants;
+
+    #[test]
+    fn test_build_required_capabilities_all() {
+        let command = AggregatorSelectCommand {
+            network: MithrilNetwork::dummy(),
+            max_entries: 1,
+            signed_entity_types: vec![],
+            aggregate_signature_types: vec![],
+        };
+
+        let required_capabilities = command.build_required_capabilities();
+        assert_eq!(required_capabilities, RequiredAggregatorCapabilities::All);
+    }
+
+    #[test]
+    fn test_build_required_capabilities_signed_entity_types() {
+        let command = AggregatorSelectCommand {
+            network: MithrilNetwork::dummy(),
+            max_entries: 1,
+            signed_entity_types: vec![
+                SignedEntityTypeDiscriminants::CardanoTransactions,
+                SignedEntityTypeDiscriminants::CardanoStakeDistribution,
+            ],
+            aggregate_signature_types: vec![],
+        };
+
+        let required_capabilities = command.build_required_capabilities();
+
+        assert_eq!(
+            required_capabilities,
+            RequiredAggregatorCapabilities::Or(vec![
+                RequiredAggregatorCapabilities::SignedEntityType(
+                    SignedEntityTypeDiscriminants::CardanoTransactions
+                ),
+                RequiredAggregatorCapabilities::SignedEntityType(
+                    SignedEntityTypeDiscriminants::CardanoStakeDistribution
+                ),
+            ])
+        );
+    }
+
+    #[test]
+    fn test_build_required_capabilities_aggregate_signature_types() {
+        let command = AggregatorSelectCommand {
+            network: MithrilNetwork::dummy(),
+            max_entries: 1,
+            signed_entity_types: vec![],
+            aggregate_signature_types: vec![AggregateSignatureType::Concatenation],
+        };
+        let required_capabilities = command.build_required_capabilities();
+
+        assert_eq!(
+            required_capabilities,
+            RequiredAggregatorCapabilities::Or(vec![
+                RequiredAggregatorCapabilities::AggregateSignatureType(
+                    AggregateSignatureType::Concatenation
+                ),
+            ])
+        );
+    }
+
+    #[test]
+    fn test_build_required_capabilities_both() {
+        let command = AggregatorSelectCommand {
+            network: MithrilNetwork::dummy(),
+            max_entries: 1,
+            signed_entity_types: vec![
+                SignedEntityTypeDiscriminants::CardanoTransactions,
+                SignedEntityTypeDiscriminants::CardanoStakeDistribution,
+            ],
+            aggregate_signature_types: vec![AggregateSignatureType::Concatenation],
+        };
+        let required_capabilities = command.build_required_capabilities();
+
+        assert_eq!(
+            required_capabilities,
+            RequiredAggregatorCapabilities::And(vec![
+                RequiredAggregatorCapabilities::Or(vec![
+                    RequiredAggregatorCapabilities::SignedEntityType(
+                        SignedEntityTypeDiscriminants::CardanoTransactions
+                    ),
+                    RequiredAggregatorCapabilities::SignedEntityType(
+                        SignedEntityTypeDiscriminants::CardanoStakeDistribution
+                    ),
+                ]),
+                RequiredAggregatorCapabilities::Or(vec![
+                    RequiredAggregatorCapabilities::AggregateSignatureType(
+                        AggregateSignatureType::Concatenation
+                    ),
+                ]),
+            ])
+        );
+    }
+}

mithril-client-cli/src/commands/tools/mod.rs

@@ -3,8 +3,10 @@
 //! Provides utility subcommands such as converting restored InMemory UTxO-HD ledger snapshot
 //! to different flavors (Legacy, LMDB).
 
+mod aggregator_discovery;
 mod snapshot_converter;
 
+pub use aggregator_discovery::*;
 pub use snapshot_converter::*;
 
 use anyhow::anyhow;
@@ -18,13 +20,17 @@ pub enum ToolsCommands {
     /// UTxO-HD related commands
     #[clap(subcommand, name = "utxo-hd")]
     UTxOHD(UTxOHDCommands),
+    /// Aggregator discovery related commands
+    #[clap(subcommand, name = "aggregator-discovery")]
+    AggregatorDiscovery(AggregatorDiscoveryCommands),
 }
 
 impl ToolsCommands {
     /// Execute Tools command
     pub async fn execute(&self) -> MithrilResult<()> {
         match self {
             Self::UTxOHD(cmd) => cmd.execute().await,
+            Self::AggregatorDiscovery(cmd) => cmd.execute().await,
         }
     }
 }
@@ -52,3 +58,20 @@ impl UTxOHDCommands {
         }
     }
 }
+
+/// Aggregator discovery related commands
+#[derive(Subcommand, Debug, Clone)]
+pub enum AggregatorDiscoveryCommands {
+    /// Select an aggregator from the available ones with automatic discovery
+    #[clap(arg_required_else_help = false)]
+    Select(AggregatorSelectCommand),
+}
+
+impl AggregatorDiscoveryCommands {
+    /// Execute Aggregator discovery command
+    pub async fn execute(&self) -> MithrilResult<()> {
+        match self {
+            Self::Select(cmd) => cmd.execute().await,
+        }
+    }
+}

mithril-client/src/client.rs

@@ -9,8 +9,8 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use mithril_aggregator_discovery::{
-    AggregatorDiscoverer, CapableAggregatorDiscoverer, HttpConfigAggregatorDiscoverer,
-    MithrilNetwork, RequiredAggregatorCapabilities,
+    AggregatorDiscoverer, AggregatorEndpoint, CapableAggregatorDiscoverer,
+    HttpConfigAggregatorDiscoverer, MithrilNetwork, RequiredAggregatorCapabilities,
 };
 use mithril_common::api_version::APIVersionProvider;
 use mithril_common::{MITHRIL_CLIENT_TYPE_HEADER, MITHRIL_ORIGIN_TAG_HEADER};
@@ -294,8 +294,7 @@ impl ClientBuilder {
             None => {
                 return Err(anyhow!(
                     "The genesis verification key must be provided to build the client with the 'with_genesis_verification_key' function"
-                )
-                .into());
+                ));
             }
         };
 
@@ -405,40 +404,47 @@ impl ClientBuilder {
         })
     }
 
+    /// Discover available aggregator endpoints for the given Mithril network and required capabilities.
+    pub fn discover_aggregator(
+        &self,
+        network: &MithrilNetwork,
+    ) -> MithrilResult<impl Iterator<Item = AggregatorEndpoint>> {
+        match self.aggregator_discoverer.clone() {
+            Some(discoverer) => {
+                let discoverer = if let Some(capabilities) = &self.aggregator_capabilities {
+                    Arc::new(CapableAggregatorDiscoverer::new(
+                        capabilities.to_owned(),
+                        discoverer.clone(),
+                    )) as Arc<dyn AggregatorDiscoverer>
+                } else {
+                    discoverer as Arc<dyn AggregatorDiscoverer>
+                };
+                tokio::task::block_in_place(move || {
+                    tokio::runtime::Handle::current().block_on(async move {
+                        discoverer
+                            .get_available_aggregators(network.to_owned())
+                            .await
+                            .with_context(|| "Discovering aggregator endpoint failed")
+                    })
+                })
+            }
+            None => Err(anyhow!(
+                "The aggregator discoverer must be provided to build the client with automatic discovery using the 'with_aggregator_discoverer' function"
+            )),
+        }
+    }
+
     fn build_aggregator_client(
         &self,
         logger: Logger,
     ) -> Result<AggregatorHTTPClient, anyhow::Error> {
         let aggregator_endpoint = match self.aggregator_discovery {
             AggregatorDiscoveryType::Url(ref url) => url.clone(),
-            AggregatorDiscoveryType::Automatic(ref network) => {
-                match self.aggregator_discoverer.clone() {
-                    Some(discoverer) => {
-                        let discoverer = if let Some(capabilities) = &self.aggregator_capabilities {
-                            Arc::new(CapableAggregatorDiscoverer::new(
-                                capabilities.to_owned(),
-                                discoverer.clone(),
-                            )) as Arc<dyn AggregatorDiscoverer>
-                        } else {
-                            discoverer as Arc<dyn AggregatorDiscoverer>
-                        };
-                        tokio::task::block_in_place(move || {
-                            tokio::runtime::Handle::current().block_on(async move {
-                                discoverer
-                                    .get_available_aggregators(network.to_owned())
-                                    .await
-                                    .with_context(|| "Discovering aggregator endpoint failed")?
-                                    .next()
-                                    .ok_or(anyhow!("No aggregator was available through discovery"))
-                            })
-                        })?
-                        .into()
-                    }
-                    None => {
-                        return Err(anyhow!("The aggregator discoverer must be provided to build the client with automatic discovery using the 'with_aggregator_discoverer' function").into());
-                    }
-                }
-            }
+            AggregatorDiscoveryType::Automatic(ref network) => self
+                .discover_aggregator(network)?
+                .next()
+                .ok_or_else(|| anyhow!("No aggregator was available through discovery"))?
+                .into(),
         };
         let endpoint_url = Url::parse(&aggregator_endpoint).with_context(|| {
             format!("Invalid aggregator endpoint, it must be a correctly formed url: '{aggregator_endpoint}'")

mithril-client/src/type_alias.rs

@@ -89,3 +89,6 @@ pub mod common {
 
 /// Required capabilities for an aggregator.
 pub use mithril_aggregator_discovery::RequiredAggregatorCapabilities;
+
+/// Mithril network
+pub use mithril_aggregator_discovery::MithrilNetwork;

mithril-stm/src/aggregate_signature/signature.rs

@@ -60,6 +60,17 @@ impl<D: Clone + Digest + FixedOutput + Send + Sync> From<&AggregateSignature<D>>
     }
 }
 
+impl From<String> for AggregateSignatureType {
+    fn from(s: String) -> Self {
+        match s.as_str() {
+            "Concatenation" => AggregateSignatureType::Concatenation,
+            #[cfg(feature = "future_proof_system")]
+            "Future" => AggregateSignatureType::Future,
+            _ => panic!("Unknown aggregate signature type: {}", s),
+        }
+    }
+}
+
 impl Display for AggregateSignatureType {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {