seperate pick and receive randomness

OmniFlix · Apr 12, 2024 · 28f1eb2 · 28f1eb2
1 parent 4ed59cb
commit 28f1eb2
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 21 deletions.
diff --git 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::<Vec<u32>>();
-
-    let winners = pick(randomness, 100, participant_arr);
-    let winners_string = winners
-        .iter()
-        .map(|&x| x.to_string())
-        .collect::<Vec<String>>()
-        .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<Response, ContractError> {
+    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::<Vec<u32>>();
+
+    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::<Vec<String>>()
+                .join(","),
+        );
+    Ok(res)
+}
+
+pub fn pick_final_winners(
+    deps: DepsMut,
+    _env: Env,
+    info: MessageInfo,
+) -> Result<Response, ContractError> {
+    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::<Vec<u32>>();
+    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::<Vec<String>>()
+                .join(","),
+        );
     Ok(res)
 }
 

diff --git 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
@@ -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
@@ -6,3 +6,5 @@ pub const NOIS_PROXY: Item<Addr> = Item::new("nois_proxy");
 pub const ADMIN: Item<Addr> = Item::new("admin");
 pub const WINNERS: Item<Vec<u32>> = Item::new("winners");
 pub const TEST_WINNERS: Item<Vec<u32>> = 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");