Feat/upgrade sns canisters for entire network

Cargo.toml

@@ -63,6 +63,7 @@ google-cloud-bigquery = { version = "0.13.1", default-features = false, features
     "auth",
     "rustls-tls",
 ] }
+hex = "0.4.3"
 
 [build-dependencies]
 tonic-build = "0.12"

src/canister/upgrade_user_token_sns_canister.rs

@@ -3,24 +3,42 @@ use axum::{
     Json,
 };
 use candid::Principal;
-use futures::{stream::FuturesUnordered, TryStreamExt};
+use futures::{stream::FuturesUnordered, StreamExt, TryStreamExt};
+use google_cloud_bigquery::storage::array::Array;
+use hex::ToHex;
 use ic_agent::Agent;
+use k256::elliptic_curve::rand_core::le;
 use serde::{Deserialize, Serialize};
 use std::{error::Error, sync::Arc, vec};
 use yral_canisters_client::{
     individual_user_template::{DeployedCdaoCanisters, IndividualUserTemplate},
-    platform_orchestrator::PlatformOrchestrator,
+    platform_orchestrator::{self, PlatformOrchestrator},
     sns_governance::{
-        self, Action, Command1, Configure, Follow, GetProposal, IncreaseDissolveDelay, ListNeurons,
-        ManageNeuron, NeuronId, Operation, Proposal, ProposalId, SnsGovernance,
+        self, Action, Command1, Configure, Follow, GetProposal, GetRunningSnsVersionArg,
+        IncreaseDissolveDelay, ListNeurons, ManageNeuron, NeuronId, Operation, Proposal,
+        ProposalId, SnsGovernance,
     },
+    user_index::UserIndex,
 };
 
 use crate::{consts::PLATFORM_ORCHESTRATOR_ID, qstash::client::QStashClient};
 
 use crate::app_state::AppState;
 use crate::utils::api_response::ApiResponse;
 
+pub const SNS_TOKEN_GOVERNANCE_MODULE_HASH: &'static str =
+    "bc91fd7bc4d6c01ea814b12510a1ff8f4f74fcac9ab16248ad4af7cb98d9c69d";
+pub const SNS_TOKEN_LEDGER_MODULE_HASH: &'static str =
+    "3d808fa63a3d8ebd4510c0400aa078e99a31afaa0515f0b68778f929ce4b2a46";
+pub const SNS_TOKEN_ROOT_MODULE_HASH: &'static str =
+    "431cb333feb3f762f742b0dea58745633a2a2ca41075e9933183d850b4ddb259";
+pub const SNS_TOKEN_SWAP_MODULE_HASH: &'static str =
+    "8313ac22d2ef0a0c1290a85b47f235cfa24ca2c96d095b8dbed5502483b9cd18";
+pub const SNS_TOKEN_INDEX_MODULE_HASH: &'static str =
+    "67b5f0bf128e801adf4a959ea26c3c9ca0cd399940e169a26a2eb237899a94dd";
+pub const SNS_TOKEN_ARCHIVE_MODULE_HASH: &'static str =
+    "317771544f0e828a60ad6efc97694c425c169c4d75d911ba592546912dba3116";
+
 #[derive(Serialize, Deserialize, Clone, Copy, Debug)]
 pub struct VerifyUpgradeProposalRequest {
     pub sns_canisters: SnsCanisters,
@@ -48,6 +66,53 @@ impl From<DeployedCdaoCanisters> for SnsCanisters {
     }
 }
 
+pub async fn upgrade_user_token_sns_canister_for_entire_network(
+    State(state): State<Arc<AppState>>,
+) -> Json<ApiResponse<()>> {
+    let result = state
+        .qstash_client
+        .upgrade_user_token_sns_canister_for_entire_network()
+        .await
+        .map_err(|e| e.into());
+
+    Json(ApiResponse::from(result))
+}
+
+pub async fn upgrade_user_token_sns_canister_for_entire_network_impl(
+    agent: &Agent,
+    qstash_client: &QStashClient,
+) -> Result<(), Box<dyn Error>> {
+    let platform_orchestrator = Principal::from_text(PLATFORM_ORCHESTRATOR_ID).unwrap();
+    let mut individual_canister_ids: Vec<Principal> = vec![];
+
+    let platform_orchestrator = PlatformOrchestrator(platform_orchestrator, agent);
+
+    let subnet_orchestrators = platform_orchestrator.get_all_subnet_orchestrators().await?;
+
+    for subnet_orchestrator in subnet_orchestrators {
+        let subnet_orchestrator = UserIndex(subnet_orchestrator, agent);
+        individual_canister_ids.extend(subnet_orchestrator.get_user_canister_list().await?);
+    }
+
+    let upgrade_governance_canister_tasks =
+        individual_canister_ids
+            .into_iter()
+            .map(|individual_canister| async move {
+                qstash_client
+                    .upgrade_all_sns_canisters_for_a_user_canister(individual_canister.to_text())
+                    .await
+            });
+
+    let stream = futures::stream::iter(upgrade_governance_canister_tasks)
+        .boxed()
+        .buffer_unordered(100);
+
+    let _upgrade_creator_dao_governance_canister_tasks =
+        stream.collect::<Vec<Result<(), anyhow::Error>>>().await;
+
+    Ok(())
+}
+
 pub async fn upgrade_user_token_sns_canister_handler(
     Path(user_canister_id): Path<String>,
     State(state): State<Arc<AppState>>,
@@ -62,7 +127,7 @@ pub async fn upgrade_user_token_sns_canister_handler(
     Json(ApiResponse::from(setup_for_upgrade_result))
 }
 
-async fn setup_sns_canisters_of_a_user_canister_for_upgrade(
+pub async fn setup_sns_canisters_of_a_user_canister_for_upgrade(
     agent: &Agent,
     qstash_client: &QStashClient,
     individual_canister_id: String,
@@ -197,6 +262,70 @@ async fn recharge_canister_using_platform_orchestrator(
     Ok(())
 }
 
+pub async fn is_upgrade_required(
+    sns_governance: &SnsGovernance<'_>,
+) -> Result<bool, Box<dyn Error + Send + Sync>> {
+    let deployed_version = sns_governance
+        .get_running_sns_version(GetRunningSnsVersionArg {})
+        .await?;
+
+    let deployed_version = deployed_version
+        .deployed_version
+        .ok_or("deployed version not found")?;
+
+    let governance_hash = deployed_version
+        .governance_wasm_hash
+        .to_vec()
+        .encode_hex::<String>();
+
+    let index_hash = deployed_version
+        .index_wasm_hash
+        .to_vec()
+        .encode_hex::<String>();
+
+    let swap_hash = deployed_version
+        .swap_wasm_hash
+        .to_vec()
+        .encode_hex::<String>();
+
+    let ledger_hash = deployed_version
+        .ledger_wasm_hash
+        .to_vec()
+        .encode_hex::<String>();
+
+    let root_hash = deployed_version
+        .root_wasm_hash
+        .to_vec()
+        .encode_hex::<String>();
+
+    let archive_hash = deployed_version
+        .archive_wasm_hash
+        .to_vec()
+        .encode_hex::<String>();
+
+    let hashes = vec![
+        governance_hash,
+        index_hash,
+        swap_hash,
+        ledger_hash,
+        root_hash,
+        archive_hash,
+    ];
+
+    let final_hashes = vec![
+        SNS_TOKEN_ARCHIVE_MODULE_HASH.to_owned(),
+        SNS_TOKEN_GOVERNANCE_MODULE_HASH.to_owned(),
+        SNS_TOKEN_INDEX_MODULE_HASH.to_owned(),
+        SNS_TOKEN_LEDGER_MODULE_HASH.to_owned(),
+        SNS_TOKEN_ROOT_MODULE_HASH.to_owned(),
+        SNS_TOKEN_SWAP_MODULE_HASH.to_owned(),
+    ];
+
+    let result = hashes.iter().all(|val| final_hashes.contains(val));
+
+    Ok(result)
+}
+
 pub async fn check_if_the_proposal_executed_successfully(
     sns_governance: &SnsGovernance<'_>,
     proposal_id: u64,

src/qstash/client.rs

@@ -154,4 +154,50 @@ impl QStashClient {
 
         Ok(())
     }
+
+    pub async fn upgrade_all_sns_canisters_for_a_user_canister(
+        &self,
+        user_canister_id: String,
+    ) -> Result<(), anyhow::Error> {
+        let off_chain_ep = OFF_CHAIN_AGENT_URL
+            .join(&format!(
+                "qstash/upgrade_all_sns_canisters_for_a_user_canister/{}",
+                user_canister_id
+            ))
+            .unwrap();
+
+        let url = self.base_url.join(&format!("publish/{}", off_chain_ep))?;
+
+        self.client
+            .post(url)
+            .header(CONTENT_TYPE, "application/json")
+            .header("upstash-method", "POST")
+            .header("upstash-retries", "0")
+            .send()
+            .await?;
+
+        Ok(())
+    }
+
+    pub async fn upgrade_user_token_sns_canister_for_entire_network(
+        &self,
+    ) -> Result<(), anyhow::Error> {
+        let off_chain_ep = OFF_CHAIN_AGENT_URL
+            .join(&format!(
+                "qstash/upgrade_user_token_sns_canister_for_entire_network/",
+            ))
+            .unwrap();
+
+        let url = self.base_url.join(&format!("publish/{}", off_chain_ep))?;
+
+        self.client
+            .post(url)
+            .header(CONTENT_TYPE, "application/json")
+            .header("upstash-method", "POST")
+            .header("upstash-retries", "0")
+            .send()
+            .await?;
+
+        Ok(())
+    }
 }

src/qstash/mod.rs

@@ -3,7 +3,7 @@ mod verify;
 use std::{str::FromStr, sync::Arc, time::Duration};
 
 use axum::{
-    extract::State,
+    extract::{Path, State},
     middleware::{self},
     response::Response,
     routing::post,
@@ -30,8 +30,10 @@ use yral_qstash_types::{ClaimTokensRequest, ParticipateInSwapRequest};
 use crate::{
     app_state::AppState,
     canister::upgrade_user_token_sns_canister::{
-        check_if_the_proposal_executed_successfully, upgrade_user_token_sns_canister_impl,
-        SnsCanisters, VerifyUpgradeProposalRequest,
+        check_if_the_proposal_executed_successfully, is_upgrade_required,
+        setup_sns_canisters_of_a_user_canister_for_upgrade,
+        upgrade_user_token_sns_canister_for_entire_network_impl,
+        upgrade_user_token_sns_canister_impl, SnsCanisters, VerifyUpgradeProposalRequest,
     },
     consts::ICP_LEDGER_CANISTER_ID,
     events::{
@@ -360,23 +362,75 @@ async fn verify_sns_canister_upgrade_proposal(
     )
     .await;
 
+    let is_upgrade_required = is_upgrade_required(&sns_governance).await.unwrap_or(true);
+
     match result {
         Ok(executed) if executed => {
-            state
-                .qstash_client
-                .upgrade_sns_creator_dao_canister(
-                    verify_sns_canister_proposal_request.sns_canisters,
-                )
-                .await
-                .map_err(|_e| StatusCode::INTERNAL_SERVER_ERROR)?;
+            if is_upgrade_required {
+                state
+                    .qstash_client
+                    .upgrade_sns_creator_dao_canister(
+                        verify_sns_canister_proposal_request.sns_canisters,
+                    )
+                    .await
+                    .map_err(|_e| StatusCode::INTERNAL_SERVER_ERROR)?;
+            }
 
             Ok(Response::builder()
                 .status(StatusCode::OK)
                 .body("Proposal executed successfully".into())
                 .unwrap())
         }
+        Ok(_) => Err(StatusCode::BAD_REQUEST),
+
+        Err(e) => Ok(Response::builder()
+            .status(StatusCode::BAD_REQUEST)
+            .body(e.to_string().into())
+            .unwrap()),
+    }
+}
+
+async fn upgrade_all_sns_canisters_for_a_user_canister(
+    Path(user_canister_id): Path<String>,
+    State(state): State<Arc<AppState>>,
+) -> Result<Response, StatusCode> {
+    let result = setup_sns_canisters_of_a_user_canister_for_upgrade(
+        &state.agent,
+        &state.qstash_client,
+        user_canister_id,
+    )
+    .await;
+
+    let res = match result {
+        Ok(_) => Response::builder()
+            .status(StatusCode::OK)
+            .body("setup for upgrade complete".into())
+            .unwrap(),
+        Err(e) => Response::builder()
+            .status(StatusCode::BAD_REQUEST)
+            .body(e.to_string().into())
+            .unwrap(),
+    };
+
+    Ok(res)
+}
+
+async fn upgrade_user_token_sns_canister_for_entire_network(
+    State(state): State<Arc<AppState>>,
+) -> Response {
+    let result =
+        upgrade_user_token_sns_canister_for_entire_network_impl(&state.agent, &state.qstash_client)
+            .await;
 
-        _ => Err(StatusCode::BAD_REQUEST),
+    match result {
+        Ok(()) => Response::builder()
+            .status(StatusCode::OK)
+            .body("Upgrade Started ".into())
+            .unwrap(),
+        Err(e) => Response::builder()
+            .status(StatusCode::BAD_REQUEST)
+            .body(e.to_string().into())
+            .unwrap(),
     }
 }
 
@@ -395,6 +449,14 @@ pub fn qstash_router<S>(app_state: Arc<AppState>) -> Router<S> {
             "/verify_sns_canister_upgrade_proposal",
             post(verify_sns_canister_upgrade_proposal),
         )
+        .route(
+            "/upgrade_all_sns_canisters_for_a_user_canister",
+            post(upgrade_all_sns_canisters_for_a_user_canister),
+        )
+        .route(
+            "upgrade_user_token_sns_canister_for_entire_network",
+            post(upgrade_user_token_sns_canister_for_entire_network),
+        )
         .layer(ServiceBuilder::new().layer(middleware::from_fn_with_state(
             app_state.qstash.clone(),
             verify_qstash_message,