Skip to content

Commit

Permalink
feat: Add duplicate version check in SNS-W upgrade steps response
Browse files Browse the repository at this point in the history
  • Loading branch information
anchpop committed Nov 27, 2024
1 parent 0976476 commit 73693a5
Show file tree
Hide file tree
Showing 3 changed files with 94 additions and 6 deletions.
20 changes: 16 additions & 4 deletions rs/sns/governance/src/cached_upgrade_steps.rs
Original file line number Diff line number Diff line change
Expand Up @@ -336,7 +336,22 @@ impl Governance {
.await;

let upgrade_steps = match upgrade_steps {
Ok(upgrade_steps) => upgrade_steps,
Ok(upgrade_steps) => {
// Check for duplicate versions in the response
let mut seen = std::collections::HashSet::new();
for version in &upgrade_steps {
if !seen.insert(version) {
log!(
ERROR,
"Cannot refresh cached_upgrade_steps: SNS-W response contains duplicate versions"
);
return;
}
}
Versions {
versions: upgrade_steps,
}
}
Err(err) => {
log!(
ERROR,
Expand All @@ -346,9 +361,6 @@ impl Governance {
return;
}
};
let upgrade_steps = Versions {
versions: upgrade_steps,
};

// Ensure `cached_upgrade_steps` is initialized
let cached_upgrade_steps = self
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -659,6 +659,72 @@ fn test_upgrade_periodic_task_lock_times_out() {
assert!(!gov.acquire_upgrade_periodic_task_lock());
}

#[tokio::test]
async fn test_refresh_cached_upgrade_steps_rejects_duplicate_versions() {
// Step 1: Prepare the world.
let root_canister_id = *TEST_ROOT_CANISTER_ID;
let governance_canister_id = *TEST_GOVERNANCE_CANISTER_ID;

let mut env = NativeEnvironment::new(Some(governance_canister_id));

let version = SnsVersion {
root_wasm_hash: vec![1, 2, 3],
governance_wasm_hash: vec![2, 3, 4],
ledger_wasm_hash: vec![3, 4, 5],
swap_wasm_hash: vec![4, 5, 6],
archive_wasm_hash: vec![5, 6, 7],
index_wasm_hash: vec![6, 7, 8],
};

// Set up environment to return upgrade steps with a duplicate version
env.set_call_canister_response(
SNS_WASM_CANISTER_ID,
"list_upgrade_steps",
Encode!(&ListUpgradeStepsRequest {
starting_at: Some(version.clone()),
sns_governance_canister_id: Some(governance_canister_id.into()),
limit: 0,
})
.unwrap(),
Ok(Encode!(&ListUpgradeStepsResponse {
steps: vec![
ListUpgradeStep {
version: Some(version.clone())
},
ListUpgradeStep {
version: Some(version.clone()) // Duplicate version
},
]
})
.unwrap()),
);

let mut governance = Governance::new(
GovernanceProto {
root_canister_id: Some(root_canister_id.get()),
deployed_version: Some(Version::from(version.clone())),
..basic_governance_proto()
}
.try_into()
.unwrap(),
Box::new(env),
Box::new(DoNothingLedger {}),
Box::new(DoNothingLedger {}),
Box::new(FakeCmc::new()),
);

// Step 2: Run code under test.
assert_eq!(governance.proto.cached_upgrade_steps, None);
governance.temporarily_lock_refresh_cached_upgrade_steps();
governance.refresh_cached_upgrade_steps().await;

// Step 3: Verify that the cached_upgrade_steps was not updated due to duplicate versions
assert_eq!(
governance.proto.cached_upgrade_steps.unwrap().upgrade_steps,
None
);
}

#[test]
fn test_upgrade_periodic_task_lock_doesnt_get_stuck_during_overflow() {
let env = NativeEnvironment::new(Some(*TEST_GOVERNANCE_CANISTER_ID));
Expand Down
14 changes: 12 additions & 2 deletions rs/sns/governance/tests/governance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2909,7 +2909,17 @@ async fn test_refresh_cached_upgrade_steps_noop_if_deployed_version_none() {
async fn test_refresh_cached_upgrade_steps() {
let mut canister_fixture = GovernanceCanisterFixtureBuilder::new().create();

let expected_upgrade_steps = vec![Version::default(), Version::default(), Version::default()];
let v1 = Version::default();
let v2 = Version {
governance_wasm_hash: vec![1],
..v1.clone()
};
let v3 = Version {
governance_wasm_hash: vec![1, 2],
..v2.clone()
};

let expected_upgrade_steps = vec![v1.clone(), v2.clone(), v3.clone()];

// Set up the fixture state
{
Expand All @@ -2927,7 +2937,7 @@ async fn test_refresh_cached_upgrade_steps() {
canister_fixture
.environment_fixture
.push_mocked_canister_reply(ListUpgradeStepsResponse { steps });
canister_fixture.governance.proto.deployed_version = Some(Version::default());
canister_fixture.governance.proto.deployed_version = Some(v1);
}

// Check that the initial state is None
Expand Down

0 comments on commit 73693a5

Please sign in to comment.