wip(client-cli): new aggregator discovery command

Author: jpraynaud

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

@@ -8,8 +8,10 @@ use mithril_common::{
 use crate::{AggregatorDiscoverer, AggregatorEndpoint, MithrilNetwork};
 
 /// Required capabilities for an aggregator.
-#[derive(Clone)]
+#[derive(Clone, PartialEq, Eq, Debug)]
 pub enum RequiredAggregatorCapabilities {
+    /// All
+    All,
     /// Signed entity type.
     SignedEntityType(SignedEntityTypeDiscriminants),
     /// Aggregate signature type.
@@ -24,6 +26,7 @@ impl RequiredAggregatorCapabilities {
     /// Check if the available capabilities match the required capabilities.
     fn matches(&self, available: &AggregatorCapabilities) -> bool {
         match self {
+            RequiredAggregatorCapabilities::All => true,
             RequiredAggregatorCapabilities::SignedEntityType(required_signed_entity_type) => {
                 available
                     .signed_entity_types
@@ -126,6 +129,18 @@ mod tests {
     mod required_capabilities {
         use super::*;
 
+        #[test]
+        fn required_capabilities_match_all_success() {
+            let required = RequiredAggregatorCapabilities::All;
+            let available = AggregatorCapabilities {
+                aggregate_signature_type: Concatenation,
+                signed_entity_types: BTreeSet::from([]),
+                cardano_transactions_prover: None,
+            };
+
+            assert!(required.matches(&available));
+        }
+
         #[test]
         fn required_capabilities_match_signed_entity_types_success() {
             let required =

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

@@ -63,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

@@ -1,27 +1,192 @@
 use clap::Parser;
 
-use mithril_client::{AggregatorDiscoveryType, ClientBuilder, MithrilNetwork, MithrilResult};
+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, default_value = "[]")]
+    signed_entity_types: Vec<SignedEntityTypeDiscriminants>,
+
+    /// Aggregate signature types to consider for the discovery
+    ///
+    /// If not provided, all aggregate signature types are considered.
+    #[clap(long, default_value = "[]")]
+    aggregate_signature_types: Vec<AggregateSignatureType>,
 }
 
 impl AggregatorSelectCommand {
     /// Main command execution
     pub async fn execute(&self) -> MithrilResult<()> {
-        let aggregator_endpoints =
-            ClientBuilder::new(AggregatorDiscoveryType::Automatic(self.network.to_owned()))
-                .with_default_aggregator_discoverer()
-                .discover_aggregator(&self.network)?;
+        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);
 
-        for endpoint in aggregator_endpoints.iter() {
-            println!("Discovered aggregator endpoint: {:?}", endpoint);
+        println!(
+            "Discovering at most {} aggregator endpoints:",
+            self.max_entries,
+        );
+        for endpoint in aggregator_endpoints {
+            println!("- Found: {endpoint}");
         }
 
         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-stm/src/aggregate_signature/signature.rs

@@ -59,6 +59,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 {