From 28f1eb28f2fa0251508fbea559d1888b78b4f9cc Mon Sep 17 00:00:00 2001 From: Ninjatosba Date: Fri, 12 Apr 2024 19:04:04 +0300 Subject: [PATCH] seperate pick and receive randomness --- src/contract.rs | 97 +++++++++++++++++++++++++++++++++++++++---------- src/error.rs | 10 ++++- src/msg.rs | 2 + src/state.rs | 2 + 4 files changed, 90 insertions(+), 21 deletions(-) diff --git a/src/contract.rs b/src/contract.rs index 36c3e6e..b8c46a2 100644 --- a/src/contract.rs +++ b/src/contract.rs @@ -1,6 +1,8 @@ use crate::error::ContractError; use crate::msg::{ExecuteMsg, InstantiateMsg, QueryMsg}; -use crate::state::{ADMIN, NOIS_PROXY, PARTICIPANT_COUNT, TEST_WINNERS, WINNERS}; +use crate::state::{ + ADMIN, FINAL_RANDOMNESS, NOIS_PROXY, PARTICIPANT_COUNT, TEST_RANDOMNESS, TEST_WINNERS, WINNERS, +}; #[cfg(not(feature = "library"))] use cosmwasm_std::entry_point; use cosmwasm_std::{ @@ -49,7 +51,11 @@ pub fn execute( ExecuteMsg::RequestRandomness { job_id } => { execute_request_randomness(deps, env, info, job_id) } - ExecuteMsg::NoisReceive { callback } => execute_pick_winners(deps, env, info, callback), + ExecuteMsg::NoisReceive { callback } => { + execute_receive_randomness(deps, env, info, callback) + } + ExecuteMsg::PickTestWinners {} => pick_test_winners(deps, env, info), + ExecuteMsg::PickFinalWinners {} => pick_final_winners(deps, env, info), } } @@ -73,7 +79,7 @@ pub fn execute_request_randomness( Ok(res) } -pub fn execute_pick_winners( +pub fn execute_receive_randomness( deps: DepsMut, _env: Env, info: MessageInfo, @@ -90,30 +96,83 @@ pub fn execute_pick_winners( .to_array() .map_err(|_| ContractError::InvalidRandomness {})?; - let participant_count = PARTICIPANT_COUNT.load(deps.storage)?; - let participant_arr = (0..participant_count).collect::>(); - - let winners = pick(randomness, 100, participant_arr); - let winners_string = winners - .iter() - .map(|&x| x.to_string()) - .collect::>() - .join(", "); - match job_id.as_str() { "test" => { - TEST_WINNERS.save(deps.storage, &winners)?; + TEST_RANDOMNESS.save(deps.storage, &randomness)?; } _ => { - let old_winners = WINNERS.may_load(deps.storage)?; - if old_winners.is_some() { - return Err(ContractError::WinnersAlreadyPicked {}); + let old_randomness = FINAL_RANDOMNESS.may_load(deps.storage)?; + if let Some(_) = old_randomness { + return Err(ContractError::FinalRandomnessAlreadySet {}); } - WINNERS.save(deps.storage, &winners)?; + FINAL_RANDOMNESS.save(deps.storage, &randomness)?; } } - let res = Response::new().add_attribute("winners", winners_string); + let res = Response::default() + .add_attribute("action", "receive_randomness") + .add_attribute("job_id", job_id); + Ok(res) +} + +pub fn pick_test_winners( + deps: DepsMut, + _env: Env, + _info: MessageInfo, +) -> Result { + let participant_count = PARTICIPANT_COUNT.load(deps.storage)?; + let randomness = TEST_RANDOMNESS + .load(deps.storage) + .map_err(|_| ContractError::TestRandomnessNotSet {})?; + let participants = (0..participant_count).collect::>(); + + let winners = pick(randomness, 100, participants); + TEST_WINNERS.save(deps.storage, &winners)?; + + let res = Response::new() + .add_attribute("action", "pick_test_winners") + .add_attribute( + "winners", + winners + .iter() + .map(|x| x.to_string()) + .collect::>() + .join(","), + ); + Ok(res) +} + +pub fn pick_final_winners( + deps: DepsMut, + _env: Env, + info: MessageInfo, +) -> Result { + if info.sender != ADMIN.load(deps.storage)? { + return Err(ContractError::Unauthorized {}); + } + let participant_count = PARTICIPANT_COUNT.load(deps.storage)?; + let randomness = FINAL_RANDOMNESS + .load(deps.storage) + .map_err(|_| ContractError::InvalidRandomness {})?; + + let participants = (0..participant_count).collect::>(); + let winners = pick(randomness, 100, participants); + + if WINNERS.may_load(deps.storage)?.is_some() { + return Err(ContractError::FinalWinnersAlreadyPicked {}); + } + WINNERS.save(deps.storage, &winners)?; + + let res = Response::new() + .add_attribute("action", "pick_final_winners") + .add_attribute( + "winners", + winners + .iter() + .map(|x| x.to_string()) + .collect::>() + .join(","), + ); Ok(res) } diff --git a/src/error.rs b/src/error.rs index 4317089..a7cd47d 100644 --- a/src/error.rs +++ b/src/error.rs @@ -18,6 +18,12 @@ pub enum ContractError { #[error("Invalid randomness")] InvalidRandomness {}, - #[error("Winners already picked")] - WinnersAlreadyPicked {}, + #[error("Final Randomness already set")] + FinalRandomnessAlreadySet {}, + + #[error("Test Randomness not set")] + TestRandomnessNotSet {}, + + #[error("Final Winners already picked")] + FinalWinnersAlreadyPicked {}, } diff --git a/src/msg.rs b/src/msg.rs index ad8b803..1b329f2 100644 --- a/src/msg.rs +++ b/src/msg.rs @@ -11,6 +11,8 @@ pub struct InstantiateMsg { pub enum ExecuteMsg { RequestRandomness { job_id: String }, NoisReceive { callback: NoisCallback }, + PickTestWinners {}, + PickFinalWinners {}, } #[cw_serde] diff --git a/src/state.rs b/src/state.rs index ef3a548..19829d1 100644 --- a/src/state.rs +++ b/src/state.rs @@ -6,3 +6,5 @@ pub const NOIS_PROXY: Item = Item::new("nois_proxy"); pub const ADMIN: Item = Item::new("admin"); pub const WINNERS: Item> = Item::new("winners"); pub const TEST_WINNERS: Item> = Item::new("test_winners"); +pub const TEST_RANDOMNESS: Item<[u8; 32]> = Item::new("test_randomness"); +pub const FINAL_RANDOMNESS: Item<[u8; 32]> = Item::new("final_randomness");