feat(sns): add Root.register_extension (#5946)

aterga · web-flow · commit 758347643cf7 · 2025-07-16T18:05:23.000Z
This PR adds a new function `Root.register_extension`, that allows registering SNS extensions. For context, the first SNS extension will be the [SNS Liquidity Pools extension](https://forum.dfinity.org/t/sns-liquidity-pools-design/50439).
diff --git a/rs/nervous_system/agent/src/sns/root.rs b/rs/nervous_system/agent/src/sns/root.rs
@@ -21,6 +21,7 @@ pub struct RootCanister {
 pub struct SnsCanisters {
     pub sns: Sns,
     pub dapps: Vec<PrincipalId>,
+    pub extensions: Vec<PrincipalId>,
 }
 
 impl TryFrom<ListSnsCanistersResponse> for SnsCanisters {
@@ -35,6 +36,7 @@ impl TryFrom<ListSnsCanistersResponse> for SnsCanisters {
             index: Some(index_canister_id),
             archives,
             dapps,
+            extensions,
         } = src
         else {
             return Err(format!("Some SNS canisters were missing: {:?}", src));
@@ -49,7 +51,14 @@ impl TryFrom<ListSnsCanistersResponse> for SnsCanisters {
             archive: archives.into_iter().map(ArchiveCanister::new).collect(),
         };
 
-        Ok(Self { sns, dapps })
+        let extensions =
+            extensions.map_or_else(Vec::new, |extensions| extensions.extension_canister_ids);
+
+        Ok(Self {
+            sns,
+            dapps,
+            extensions,
+        })
     }
 }
 
diff --git a/rs/sns/cli/src/upgrade_sns_controlled_canister.rs b/rs/sns/cli/src/upgrade_sns_controlled_canister.rs
@@ -323,11 +323,19 @@ pub async fn find_sns<C: CallCanisters>(
         let root_canister = sns::root::RootCanister { canister_id };
 
         let response = root_canister.list_sns_canisters(agent).await?;
-        let SnsCanisters { sns, dapps } =
-            SnsCanisters::try_from(response).map_err(UpgradeSnsControlledCanisterError::Client)?;
+        let SnsCanisters {
+            sns,
+            dapps,
+            extensions,
+        } = response
+            .try_into()
+            .map_err(UpgradeSnsControlledCanisterError::Client)?;
+
+        let mut sns_controlled_canisters = dapps;
+        sns_controlled_canisters.extend(extensions.iter());
 
         // Check that the target is indeed controlled by this SNS.
-        if !BTreeSet::from_iter(&dapps[..]).contains(&target_canister_id.get()) {
+        if !BTreeSet::from_iter(&sns_controlled_canisters[..]).contains(&target_canister_id.get()) {
             return Err(UpgradeSnsControlledCanisterError::Client(format!(
                 "{} is not one of the canisters controlled by the SNS with Root canister {}",
                 target_canister_id.get(),
diff --git a/rs/sns/init/src/lib.rs b/rs/sns/init/src/lib.rs
@@ -652,11 +652,15 @@ impl SnsInitPayload {
                 .collect(),
         };
 
+        // TODO: Support specifying extensions at SNS initialization.
+        let extensions = Some(vec![].into());
+
         SnsRootCanister {
             governance_canister_id: Some(sns_canister_ids.governance),
             ledger_canister_id: Some(sns_canister_ids.ledger),
             swap_canister_id: Some(sns_canister_ids.swap),
             dapp_canister_ids,
+            extensions,
             archive_canister_ids: vec![],
             index_canister_id: Some(sns_canister_ids.index),
             testflight,
diff --git a/rs/sns/integration_tests/src/root.rs b/rs/sns/integration_tests/src/root.rs
@@ -13,7 +13,8 @@ use ic_nervous_system_root::change_canister::ChangeCanisterRequest;
 use ic_nns_constants::GOVERNANCE_CANISTER_ID;
 use ic_nns_test_utils::state_test_helpers::{get_controllers, set_controllers, update_with_sender};
 use ic_sns_root::{
-    pb::v1::SnsRootCanister, GetSnsCanistersSummaryRequest, GetSnsCanistersSummaryResponse,
+    pb::v1::{Extensions, SnsRootCanister},
+    GetSnsCanistersSummaryRequest, GetSnsCanistersSummaryResponse,
 };
 use ic_sns_test_utils::{
     itest_helpers::{
@@ -40,6 +41,9 @@ fn test_get_status() {
                 ledger_canister_id: Some(PrincipalId::new_user_test_id(43)),
                 swap_canister_id: Some(PrincipalId::new_user_test_id(44)),
                 dapp_canister_ids: vec![],
+                extensions: Some(Extensions {
+                    extension_canister_ids: vec![],
+                }),
                 archive_canister_ids: vec![],
                 index_canister_id: Some(PrincipalId::new_user_test_id(45)),
                 testflight: false,
diff --git a/rs/sns/integration_tests/src/timers.rs b/rs/sns/integration_tests/src/timers.rs
@@ -6,7 +6,7 @@ use ic_nervous_system_proto::pb::v1::{
 use ic_nns_test_utils::sns_wasm::{build_governance_sns_wasm, build_root_sns_wasm};
 use ic_sns_governance::pb::v1::governance::GovernanceCachedMetrics;
 use ic_sns_governance::{init::GovernanceCanisterInitPayloadBuilder, pb::v1::Governance};
-use ic_sns_root::pb::v1::SnsRootCanister;
+use ic_sns_root::pb::v1::{Extensions, SnsRootCanister};
 use ic_sns_swap::pb::v1::{
     GetStateRequest, GetStateResponse, Init, Lifecycle, NeuronBasketConstructionParameters,
 };
@@ -80,6 +80,9 @@ fn root_init() -> SnsRootCanister {
         index_canister_id: Some(PrincipalId::new_anonymous()),
         archive_canister_ids: vec![],
         dapp_canister_ids: vec![],
+        extensions: Some(Extensions {
+            extension_canister_ids: vec![],
+        }),
         testflight: false,
         timers: None,
     }
diff --git a/rs/sns/root/canister/canister.rs b/rs/sns/root/canister/canister.rs
@@ -19,6 +19,7 @@ use ic_nervous_system_proto::pb::v1::{
 };
 use ic_nervous_system_root::change_canister::ChangeCanisterRequest;
 use ic_nervous_system_runtime::{CdkRuntime, Runtime};
+use ic_sns_root::pb::v1::{RegisterExtensionRequest, RegisterExtensionResponse};
 use ic_sns_root::{
     logs::{ERROR, INFO},
     pb::v1::{
@@ -203,7 +204,7 @@ fn list_sns_canisters(_request: ListSnsCanistersRequest) -> ListSnsCanistersResp
     STATE.with(|sns_root_canister| {
         sns_root_canister
             .borrow()
-            .list_sns_canisters(ic_cdk::api::id())
+            .list_sns_canisters(PrincipalId(ic_cdk::api::id()))
     })
 }
 
@@ -248,6 +249,32 @@ fn change_canister(request: ChangeCanisterRequest) {
     });
 }
 
+#[candid_method(update)]
+#[update]
+async fn register_extension(request: RegisterExtensionRequest) -> RegisterExtensionResponse {
+    log!(INFO, "register_extension");
+    assert_eq_governance_canister_id(PrincipalId(ic_cdk::api::caller()));
+
+    let canister_id = match PrincipalId::try_from(request) {
+        Ok(canister_id) => canister_id,
+        Err(err) => {
+            return RegisterExtensionResponse::from(Err(err));
+        }
+    };
+
+    let root_canister_id = PrincipalId(ic_cdk::api::id());
+
+    let result = SnsRootCanister::register_extension(
+        &STATE,
+        &ManagementCanisterClientImpl::<CanisterRuntime>::new(None),
+        root_canister_id,
+        canister_id,
+    )
+    .await;
+
+    RegisterExtensionResponse::from(result)
+}
+
 /// This function is deprecated, and `register_dapp_canisters` should be used
 /// instead. (NNS1-1991)
 ///
@@ -277,7 +304,7 @@ async fn register_dapp_canister(
     let RegisterDappCanistersResponse {} = SnsRootCanister::register_dapp_canisters(
         &STATE,
         &ManagementCanisterClientImpl::<CanisterRuntime>::new(None),
-        ic_cdk::api::id(),
+        PrincipalId(ic_cdk::api::id()),
         request,
     )
     .await;
@@ -304,7 +331,7 @@ async fn register_dapp_canisters(
     SnsRootCanister::register_dapp_canisters(
         &STATE,
         &ManagementCanisterClientImpl::<CanisterRuntime>::new(None),
-        ic_cdk::api::id(),
+        PrincipalId(ic_cdk::api::id()),
         request,
     )
     .await
@@ -333,7 +360,7 @@ async fn set_dapp_controllers(request: SetDappControllersRequest) -> SetDappCont
     SnsRootCanister::set_dapp_controllers(
         &STATE,
         &ManagementCanisterClientImpl::<CanisterRuntime>::new(None),
-        ic_cdk::api::id(),
+        PrincipalId(ic_cdk::api::id()),
         PrincipalId(ic_cdk::api::caller()),
         &request,
     )
diff --git a/rs/sns/root/canister/root.did b/rs/sns/root/canister/root.did
@@ -113,6 +113,7 @@ type ListSnsCanistersResponse = record {
   index : opt principal;
   governance : opt principal;
   dapps : vec principal;
+  extensions : opt Extensions;
   archives : vec principal;
 };
 
@@ -137,6 +138,19 @@ type ManageDappCanisterSettingsResponse = record {
   failure_reason : opt text;
 };
 
+type RegisterExtensionRequest = record {
+  canister_id : opt principal;
+};
+
+type RegisterExtensionResult = variant {
+  Ok : record {};
+  Err : CanisterCallError;
+};
+
+type RegisterExtensionResponse = record {
+  result : opt RegisterExtensionResult;
+};
+
 type RegisterDappCanisterRequest = record {
   canister_id : opt principal;
 };
@@ -154,8 +168,13 @@ type SetDappControllersResponse = record {
   failed_updates : vec FailedUpdate;
 };
 
+type Extensions = record {
+  extension_canister_ids : vec principal;
+};
+
 type SnsRootCanister = record {
   dapp_canister_ids : vec principal;
+  extensions : opt Extensions;
   testflight : bool;
   archive_canister_ids : vec principal;
   governance_canister_id : opt principal;
@@ -186,6 +205,7 @@ service : (SnsRootCanister) -> {
   manage_dapp_canister_settings : (ManageDappCanisterSettingsRequest) -> (
       ManageDappCanisterSettingsResponse,
     );
+  register_extension : (RegisterExtensionRequest) -> (RegisterExtensionResponse);
   register_dapp_canister : (RegisterDappCanisterRequest) -> (record {});
   register_dapp_canisters : (RegisterDappCanistersRequest) -> (record {});
   set_dapp_controllers : (SetDappControllersRequest) -> (
diff --git a/rs/sns/root/proto/ic_sns_root/pb/v1/root.proto b/rs/sns/root/proto/ic_sns_root/pb/v1/root.proto
@@ -26,6 +26,9 @@ message SnsRootCanister {
   // Dapp canister IDs.
   repeated ic_base_types.pb.v1.PrincipalId dapp_canister_ids = 3;
 
+  // Extension canister IDs.
+  optional Extensions extensions = 11;
+
   // Required.
   //
   // The swap canister ID.
@@ -52,6 +55,19 @@ message SnsRootCanister {
   optional ic_nervous_system.pb.v1.Timers timers = 10;
 }
 
+message RegisterExtensionRequest {
+  ic_base_types.pb.v1.PrincipalId canister_id = 1;
+}
+
+message RegisterExtensionResponse {
+  message Ok {}
+
+  oneof result {
+    Ok ok = 1;
+    CanisterCallError err = 2;
+  }
+}
+
 message RegisterDappCanisterRequest {
   ic_base_types.pb.v1.PrincipalId canister_id = 1;
 }
@@ -97,6 +113,10 @@ message ListSnsCanistersRequest {
   // This struct intentionally left blank (for now).
 }
 
+message Extensions {
+  repeated ic_base_types.pb.v1.PrincipalId extension_canister_ids = 1;
+}
+
 // Response struct for the ListSnsCanisters API on the
 // SNS Root canister. ListSnsCanisters will return Principals
 // of all the associated canisters in an SNS.
@@ -108,6 +128,7 @@ message ListSnsCanistersResponse {
   repeated ic_base_types.pb.v1.PrincipalId dapps = 5;
   repeated ic_base_types.pb.v1.PrincipalId archives = 6;
   ic_base_types.pb.v1.PrincipalId index = 7;
+  optional Extensions extensions = 8;
 }
 
 enum LogVisibility {
diff --git a/rs/sns/root/src/gen/ic_sns_root.pb.v1.rs b/rs/sns/root/src/gen/ic_sns_root.pb.v1.rs
@@ -28,6 +28,9 @@ pub struct SnsRootCanister {
     /// Dapp canister IDs.
     #[prost(message, repeated, tag = "3")]
     pub dapp_canister_ids: ::prost::alloc::vec::Vec<::ic_base_types::PrincipalId>,
+    /// Extension canister IDs.
+    #[prost(message, optional, tag = "11")]
+    pub extensions: ::core::option::Option<Extensions>,
     /// Required.
     ///
     /// The swap canister ID.
@@ -57,6 +60,57 @@ pub struct SnsRootCanister {
     PartialEq,
     ::prost::Message,
 )]
+pub struct RegisterExtensionRequest {
+    #[prost(message, optional, tag = "1")]
+    pub canister_id: ::core::option::Option<::ic_base_types::PrincipalId>,
+}
+#[derive(
+    candid::CandidType,
+    candid::Deserialize,
+    comparable::Comparable,
+    Clone,
+    PartialEq,
+    ::prost::Message,
+)]
+pub struct RegisterExtensionResponse {
+    #[prost(oneof = "register_extension_response::Result", tags = "1, 2")]
+    pub result: ::core::option::Option<register_extension_response::Result>,
+}
+/// Nested message and enum types in `RegisterExtensionResponse`.
+pub mod register_extension_response {
+    #[derive(
+        candid::CandidType,
+        candid::Deserialize,
+        comparable::Comparable,
+        Clone,
+        Copy,
+        PartialEq,
+        ::prost::Message,
+    )]
+    pub struct Ok {}
+    #[derive(
+        candid::CandidType,
+        candid::Deserialize,
+        comparable::Comparable,
+        Clone,
+        PartialEq,
+        ::prost::Oneof,
+    )]
+    pub enum Result {
+        #[prost(message, tag = "1")]
+        Ok(Ok),
+        #[prost(message, tag = "2")]
+        Err(super::CanisterCallError),
+    }
+}
+#[derive(
+    candid::CandidType,
+    candid::Deserialize,
+    comparable::Comparable,
+    Clone,
+    PartialEq,
+    ::prost::Message,
+)]
 pub struct RegisterDappCanisterRequest {
     #[prost(message, optional, tag = "1")]
     pub canister_id: ::core::option::Option<::ic_base_types::PrincipalId>,
@@ -184,6 +238,18 @@ pub struct CanisterCallError {
     ::prost::Message,
 )]
 pub struct ListSnsCanistersRequest {}
+#[derive(
+    candid::CandidType,
+    candid::Deserialize,
+    comparable::Comparable,
+    Clone,
+    PartialEq,
+    ::prost::Message,
+)]
+pub struct Extensions {
+    #[prost(message, repeated, tag = "1")]
+    pub extension_canister_ids: ::prost::alloc::vec::Vec<::ic_base_types::PrincipalId>,
+}
 /// Response struct for the ListSnsCanisters API on the
 /// SNS Root canister. ListSnsCanisters will return Principals
 /// of all the associated canisters in an SNS.
@@ -210,6 +276,8 @@ pub struct ListSnsCanistersResponse {
     pub archives: ::prost::alloc::vec::Vec<::ic_base_types::PrincipalId>,
     #[prost(message, optional, tag = "7")]
     pub index: ::core::option::Option<::ic_base_types::PrincipalId>,
+    #[prost(message, optional, tag = "8")]
+    pub extensions: ::core::option::Option<Extensions>,
 }
 #[derive(
     candid::CandidType,
diff --git a/rs/sns/root/src/lib.rs b/rs/sns/root/src/lib.rs
diff --git a/rs/sns/root/src/types.rs b/rs/sns/root/src/types.rs
diff --git a/rs/sns/root/unreleased_changelog.md b/rs/sns/root/unreleased_changelog.md
