Partially working POC without tests

rs/nns/test_utils/src/registry.rs

@@ -15,14 +15,16 @@ use ic_nervous_system_common_test_keys::{
     TEST_USER1_PRINCIPAL, TEST_USER2_PRINCIPAL, TEST_USER3_PRINCIPAL, TEST_USER4_PRINCIPAL,
     TEST_USER5_PRINCIPAL, TEST_USER6_PRINCIPAL, TEST_USER7_PRINCIPAL,
 };
-use ic_protobuf::registry::subnet::v1::{ChainKeyConfig, InitialNiDkgTranscriptRecord};
 use ic_protobuf::registry::{
     crypto::v1::{PublicKey, X509PublicKeyCert},
     node::v1::{ConnectionEndpoint, NodeRecord},
     node_operator::v1::NodeOperatorRecord,
     replica_version::v1::{BlessedReplicaVersions, ReplicaVersionRecord},
     routing_table::v1::RoutingTable as RoutingTablePB,
-    subnet::v1::{CatchUpPackageContents, SubnetListRecord, SubnetRecord},
+    subnet::v1::{
+        CatchUpPackageContents, ChainKeyConfig, InitialNiDkgTranscriptRecord, SubnetListRecord,
+        SubnetRecord,
+    },
 };
 use ic_registry_canister_api::AddNodePayload;
 use ic_registry_keys::{
@@ -42,9 +44,11 @@ use ic_registry_transport::{
     serialize_get_value_request, Error,
 };
 use ic_test_utilities_types::ids::{subnet_test_id, user_test_id};
-use ic_types::crypto::threshold_sig::ni_dkg::{NiDkgTag, NiDkgTargetId, NiDkgTranscript};
 use ic_types::{
-    crypto::{CurrentNodePublicKeys, KeyPurpose},
+    crypto::{
+        threshold_sig::ni_dkg::{NiDkgTag, NiDkgTargetId, NiDkgTranscript},
+        CurrentNodePublicKeys, KeyPurpose,
+    },
     NodeId, ReplicaVersion,
 };
 use maplit::btreemap;
@@ -54,8 +58,10 @@ use rand::RngCore;
 use registry_canister::mutations::node_management::{
     common::make_add_node_registry_mutations, do_add_node::connection_endpoint_from_string,
 };
-use std::collections::{BTreeMap, BTreeSet};
-use std::convert::TryFrom;
+use std::{
+    collections::{BTreeMap, BTreeSet},
+    convert::TryFrom,
+};
 
 /// ID used in multiple tests.
 pub const TEST_ID: u64 = 999;
@@ -827,6 +833,7 @@ pub fn prepare_add_node_payload(mutation_id: u8) -> (AddNodePayload, ValidNodePu
         // Unused section follows
         p2p_flow_endpoints: Default::default(),
         prometheus_metrics_endpoint: Default::default(),
+        node_type: None,
     };
 
     (payload, node_public_keys)

rs/orchestrator/src/registration.rs

@@ -222,6 +222,7 @@ impl NodeRegistration {
             // Unused section follows
             p2p_flow_endpoints: Default::default(),
             prometheus_metrics_endpoint: Default::default(),
+            node_type: None,
         }
     }
 

rs/registry/canister/api/src/lib.rs

@@ -1,10 +1,7 @@
 use candid::{CandidType, Deserialize};
 use ic_base_types::NodeId;
 use serde::Serialize;
-use std::collections::HashSet;
-use std::fmt;
-use std::net::Ipv4Addr;
-use std::str::FromStr;
+use std::{collections::HashSet, fmt, net::Ipv4Addr, str::FromStr};
 use thiserror::Error;
 
 #[derive(Debug, Error)]
@@ -186,6 +183,10 @@ pub struct AddNodePayload {
     // TODO(NNS1-2444): The fields below are deprecated and they are not read anywhere.
     pub p2p_flow_endpoints: Vec<String>,
     pub prometheus_metrics_endpoint: String,
+
+    // String representation of the node type.  Must be a valid type.
+    // Currently valid types are type0, type1, type2, type3, and type3.1.
+    pub node_type: Option<String>,
 }
 
 /// The payload of a request to update keys of the existing node.

rs/registry/canister/canister/registry.did

@@ -42,6 +42,7 @@ type AddNodePayload = record {
   transport_tls_cert : blob;
   ni_dkg_dealing_encryption_pk : blob;
   p2p_flow_endpoints : vec text;
+  node_type : opt text;
 };
 
 type AddNodesToSubnetPayload = record {

rs/registry/canister/src/mutations/node_management/do_add_node.rs

@@ -10,7 +10,7 @@ use ic_crypto_node_key_validation::ValidNodePublicKeys;
 use ic_crypto_utils_basic_sig::conversions as crypto_basicsig_conversions;
 use ic_protobuf::registry::{
     crypto::v1::{PublicKey, X509PublicKeyCert},
-    node::v1::{ConnectionEndpoint, IPv4InterfaceConfig, NodeRecord},
+    node::v1::{ConnectionEndpoint, IPv4InterfaceConfig, NodeRecord, NodeType},
 };
 use idna::domain_to_ascii_strict;
 
@@ -75,11 +75,26 @@ impl Registry {
             ));
         }
 
-        // 3. Validate keys and get the node id
+        // 3. Get valid type if type is in request
+        let node_type = payload
+            .node_type
+            .as_ref()
+            .map(|t| {
+                try_str_to_node_type(t).map_err(|e| {
+                    format!(
+                        "{}do_add_node: Error parsing node type from payload: {}",
+                        LOG_PREFIX, e
+                    )
+                })
+            })
+            .transpose()?
+            .map(|node_type| node_type as i32);
+
+        // 4. Validate keys and get the node id
         let (node_id, valid_pks) = valid_keys_from_payload(&payload)
             .map_err(|err| format!("{}do_add_node: {}", LOG_PREFIX, err))?;
 
-        // 4. Validate the domain is valid
+        // 5. Validate the domain is valid
         let domain: Option<String> = payload
             .domain
             .as_ref()
@@ -93,7 +108,7 @@ impl Registry {
             })
             .transpose()?;
 
-        // 5. If there is an IPv4 config, make sure that the IPv4 is not used by anyone else
+        // 6. If there is an IPv4 config, make sure that the IPv4 is not used by anyone else
         let ipv4_intf_config = payload.public_ipv4_config.clone().map(|ipv4_config| {
             ipv4_config.panic_on_invalid();
             IPv4InterfaceConfig {
@@ -113,7 +128,7 @@ impl Registry {
 
         println!("{}do_add_node: The node id is {:?}", LOG_PREFIX, node_id);
 
-        // 6. Create the Node Record
+        // 7. Create the Node Record
         let node_record = NodeRecord {
             xnet: Some(connection_endpoint_from_string(&payload.xnet_endpoint)),
             http: Some(connection_endpoint_from_string(&payload.http_endpoint)),
@@ -122,21 +137,21 @@ impl Registry {
             chip_id: payload.chip_id.clone(),
             public_ipv4_config: ipv4_intf_config,
             domain,
-            node_type,
+            node_type: node_type,
         };
 
-        // 7. Insert node, public keys, and crypto keys
+        // 8. Insert node, public keys, and crypto keys
         let mut mutations = make_add_node_registry_mutations(node_id, node_record, valid_pks);
 
-        // 8. Update the Node Operator record
+        // 9. Update the Node Operator record
         node_operator_record.node_allowance -= 1;
 
         let update_node_operator_record =
             make_update_node_operator_mutation(caller_id, &node_operator_record);
 
         mutations.push(update_node_operator_record);
 
-        // 9. Check invariants before applying mutations
+        // 10. Check invariants before applying mutations
         self.maybe_apply_mutation_internal(mutations);
 
         println!("{}do_add_node finished: {:?}", LOG_PREFIX, payload);
@@ -145,6 +160,18 @@ impl Registry {
     }
 }
 
+// try to convert input string inot NodeType enum
+fn try_str_to_node_type(type_string: &str) -> Result<NodeType, String> {
+    Ok(match type_string {
+        "type0" => NodeType::Type0,
+        "type1" => NodeType::Type1,
+        "type2" => NodeType::Type2,
+        "type3" => NodeType::Type3,
+        "type3.1" => NodeType::Type3dot1,
+        _ => return Err(format!("Invalid node type: {}", type_string)),
+    })
+}
+
 /// Parses the ConnectionEndpoint string
 ///
 /// The string is written in form: `ipv4:port` or `[ipv6]:port`.
@@ -302,6 +329,7 @@ mod tests {
             // Unused section follows
             p2p_flow_endpoints: Default::default(),
             prometheus_metrics_endpoint: Default::default(),
+            node_type: None,
         };
 
         (payload, node_public_keys)
@@ -340,6 +368,7 @@ mod tests {
             // Unused section follows
             p2p_flow_endpoints: Default::default(),
             prometheus_metrics_endpoint: Default::default(),
+            node_type: None,
         };
     }