use cosmwasm_std::{
to_binary, to_vec, Addr, Binary, Deps, DepsMut, MessageInfo, Response, StdError, StdResult,
Storage,
};
use cosmwasm_storage::{Bucket, ReadonlyBucket};
use schemars::JsonSchema;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
const PREFIX_FEATURES: &[u8] = b"features";
const PREFIX_PAUSERS: &[u8] = b"pausers";
pub struct FeatureToggle;
impl FeatureToggleTrait for FeatureToggle {
const STORAGE_KEY: &'static [u8] = b"feature_toggle";
}
pub trait FeatureToggleTrait {
const STORAGE_KEY: &'static [u8];
fn init_features<T: Serialize>(
storage: &mut dyn Storage,
feature_statuses: Vec<FeatureStatus<T>>,
pausers: Vec<Addr>,
) -> StdResult<()> {
for fs in feature_statuses {
Self::set_feature_status(storage, &fs.feature, fs.status)?;
}
for p in pausers {
Self::set_pauser(storage, &p)?;
}
Ok(())
}
fn require_not_paused<T: Serialize>(storage: &dyn Storage, features: Vec<T>) -> StdResult<()> {
for feature in features {
let status = Self::get_feature_status(storage, &feature)?;
match status {
None => {
return Err(StdError::generic_err(format!(
"feature toggle: unknown feature '{}'",
String::from_utf8_lossy(&to_vec(&feature)?)
)))
}
Some(s) => match s {
Status::NotPaused => {}
Status::Paused => {
return Err(StdError::generic_err(format!(
"feature toggle: feature '{}' is paused",
String::from_utf8_lossy(&to_vec(&feature)?)
)));
}
},
}
}
Ok(())
}
fn pause<T: Serialize>(storage: &mut dyn Storage, features: Vec<T>) -> StdResult<()> {
for f in features {
Self::set_feature_status(storage, &f, Status::Paused)?;
}
Ok(())
}
fn unpause<T: Serialize>(storage: &mut dyn Storage, features: Vec<T>) -> StdResult<()> {
for f in features {
Self::set_feature_status(storage, &f, Status::NotPaused)?;
}
Ok(())
}
fn is_pauser(storage: &dyn Storage, key: &Addr) -> StdResult<bool> {
let feature_store: ReadonlyBucket<bool> =
ReadonlyBucket::multilevel(storage, &[Self::STORAGE_KEY, PREFIX_PAUSERS]);
feature_store.may_load(key.as_bytes()).map(|p| p.is_some())
}
fn set_pauser(storage: &mut dyn Storage, key: &Addr) -> StdResult<()> {
let mut feature_store = Bucket::multilevel(storage, &[Self::STORAGE_KEY, PREFIX_PAUSERS]);
feature_store.save(key.as_bytes(), &true )
}
fn remove_pauser(storage: &mut dyn Storage, key: &Addr) {
let mut feature_store: Bucket<bool> =
Bucket::multilevel(storage, &[Self::STORAGE_KEY, PREFIX_PAUSERS]);
feature_store.remove(key.as_bytes())
}
fn get_feature_status<T: Serialize>(
storage: &dyn Storage,
key: &T,
) -> StdResult<Option<Status>> {
let feature_store =
ReadonlyBucket::multilevel(storage, &[Self::STORAGE_KEY, PREFIX_FEATURES]);
feature_store.may_load(&cosmwasm_std::to_vec(&key)?)
}
fn set_feature_status<T: Serialize>(
storage: &mut dyn Storage,
key: &T,
item: Status,
) -> StdResult<()> {
let mut feature_store = Bucket::multilevel(storage, &[Self::STORAGE_KEY, PREFIX_FEATURES]);
feature_store.save(&cosmwasm_std::to_vec(&key)?, &item)
}
fn handle_pause<T: Serialize>(
deps: DepsMut,
info: &MessageInfo,
features: Vec<T>,
) -> StdResult<Response> {
if !Self::is_pauser(deps.storage, &info.sender)? {
return Err(StdError::generic_err("unauthorized"));
}
Self::pause(deps.storage, features)?;
Ok(Response::new().set_data(to_binary(&HandleAnswer::Pause {
status: ResponseStatus::Success,
})?))
}
fn handle_unpause<T: Serialize>(
deps: DepsMut,
info: &MessageInfo,
features: Vec<T>,
) -> StdResult<Response> {
if !Self::is_pauser(deps.storage, &info.sender)? {
return Err(StdError::generic_err("unauthorized"));
}
Self::unpause(deps.storage, features)?;
Ok(Response::new().set_data(to_binary(&HandleAnswer::Unpause {
status: ResponseStatus::Success,
})?))
}
fn handle_set_pauser(deps: DepsMut, address: Addr) -> StdResult<Response> {
Self::set_pauser(deps.storage, &address)?;
Ok(
Response::new().set_data(to_binary(&HandleAnswer::SetPauser {
status: ResponseStatus::Success,
})?),
)
}
fn handle_remove_pauser(deps: DepsMut, address: Addr) -> StdResult<Response> {
Self::remove_pauser(deps.storage, &address);
Ok(
Response::new().set_data(to_binary(&HandleAnswer::RemovePauser {
status: ResponseStatus::Success,
})?),
)
}
fn query_status<T: Serialize>(deps: Deps, features: Vec<T>) -> StdResult<Binary> {
let mut status = Vec::with_capacity(features.len());
for feature in features {
match Self::get_feature_status(deps.storage, &feature)? {
None => {
return Err(StdError::generic_err(format!(
"invalid feature: {} does not exist",
String::from_utf8_lossy(&to_vec(&feature)?)
)))
}
Some(s) => status.push(FeatureStatus { feature, status: s }),
}
}
to_binary(&FeatureToggleQueryAnswer::Status { features: status })
}
fn query_is_pauser(deps: Deps, address: Addr) -> StdResult<Binary> {
let is_pauser = Self::is_pauser(deps.storage, &address)?;
to_binary(&FeatureToggleQueryAnswer::<()>::IsPauser { is_pauser })
}
}
#[derive(Serialize, Debug, Deserialize, Clone, JsonSchema, PartialEq, Eq, Default)]
pub enum Status {
#[default]
NotPaused,
Paused,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum FeatureToggleHandleMsg<T: Serialize + DeserializeOwned> {
#[serde(bound = "")]
Pause {
features: Vec<T>,
},
#[serde(bound = "")]
Unpause {
features: Vec<T>,
},
SetPauser {
address: String,
},
RemovePauser {
address: String,
},
}
#[derive(Serialize, Deserialize, Clone, PartialEq, JsonSchema, Debug)]
#[serde(rename_all = "snake_case")]
enum ResponseStatus {
Success,
Failure,
}
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
#[serde(rename_all = "snake_case")]
enum HandleAnswer {
Pause { status: ResponseStatus },
Unpause { status: ResponseStatus },
SetPauser { status: ResponseStatus },
RemovePauser { status: ResponseStatus },
}
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum FeatureToggleQueryMsg<T: Serialize + DeserializeOwned> {
#[serde(bound = "")]
Status {
features: Vec<T>,
},
IsPauser {
address: String,
},
}
#[derive(Serialize, Deserialize, JsonSchema, Debug)]
#[serde(rename_all = "snake_case")]
enum FeatureToggleQueryAnswer<T: Serialize> {
Status { features: Vec<FeatureStatus<T>> },
IsPauser { is_pauser: bool },
}
#[derive(Serialize, Deserialize, JsonSchema, Debug)]
pub struct FeatureStatus<T: Serialize> {
pub feature: T,
pub status: Status,
}
#[cfg(test)]
mod tests {
use crate::feature_toggle::{
FeatureStatus, FeatureToggle, FeatureToggleHandleMsg, FeatureToggleQueryMsg,
FeatureToggleTrait, HandleAnswer, ResponseStatus, Status,
};
use cosmwasm_std::testing::{mock_dependencies, mock_info, MockStorage};
use cosmwasm_std::{from_binary, Addr, MemoryStorage, StdError, StdResult};
fn init_features(storage: &mut MemoryStorage) -> StdResult<()> {
FeatureToggle::init_features(
storage,
vec![
FeatureStatus {
feature: "Feature1".to_string(),
status: Status::NotPaused,
},
FeatureStatus {
feature: "Feature2".to_string(),
status: Status::NotPaused,
},
FeatureStatus {
feature: "Feature3".to_string(),
status: Status::Paused,
},
],
vec![Addr::unchecked("alice".to_string())],
)
}
#[test]
fn test_init_works() -> StdResult<()> {
let mut storage = MockStorage::new();
init_features(&mut storage)?;
assert_eq!(
FeatureToggle::get_feature_status(&storage, &"Feature1".to_string())?,
Some(Status::NotPaused)
);
assert_eq!(
FeatureToggle::get_feature_status(&storage, &"Feature2".to_string())?,
Some(Status::NotPaused)
);
assert_eq!(
FeatureToggle::get_feature_status(&storage, &"Feature3".to_string())?,
Some(Status::Paused)
);
assert_eq!(
FeatureToggle::get_feature_status(&storage, &"Feature4".to_string())?,
None
);
assert!(FeatureToggle::is_pauser(
&storage,
&Addr::unchecked("alice".to_string())
)?,);
assert!(!FeatureToggle::is_pauser(
&storage,
&Addr::unchecked("bob".to_string())
)?);
Ok(())
}
#[test]
fn test_unpause() -> StdResult<()> {
let mut storage = MockStorage::new();
init_features(&mut storage)?;
FeatureToggle::unpause(&mut storage, vec!["Feature3".to_string()])?;
assert_eq!(
FeatureToggle::get_feature_status(&storage, &"Feature3".to_string())?,
Some(Status::NotPaused)
);
Ok(())
}
#[test]
fn test_handle_unpause() -> StdResult<()> {
let mut deps = mock_dependencies();
init_features(&mut deps.storage)?;
let info = mock_info("non-pauser", &[]);
let error =
FeatureToggle::handle_unpause(deps.as_mut(), &info, vec!["Feature3".to_string()]);
assert_eq!(error, Err(StdError::generic_err("unauthorized")));
let info = mock_info("alice", &[]);
let response =
FeatureToggle::handle_unpause(deps.as_mut(), &info, vec!["Feature3".to_string()])?;
let answer: HandleAnswer = from_binary(&response.data.unwrap())?;
assert_eq!(
answer,
HandleAnswer::Unpause {
status: ResponseStatus::Success,
}
);
Ok(())
}
#[test]
fn test_pause() -> StdResult<()> {
let mut storage = MockStorage::new();
init_features(&mut storage)?;
FeatureToggle::pause(&mut storage, vec!["Feature1".to_string()])?;
assert_eq!(
FeatureToggle::get_feature_status(&storage, &"Feature1".to_string())?,
Some(Status::Paused)
);
Ok(())
}
#[test]
fn test_handle_pause() -> StdResult<()> {
let mut deps = mock_dependencies();
init_features(&mut deps.storage)?;
let info = mock_info("non-pauser", &[]);
let error = FeatureToggle::handle_pause(deps.as_mut(), &info, vec!["Feature2".to_string()]);
assert_eq!(error, Err(StdError::generic_err("unauthorized")));
let info = mock_info("alice", &[]);
let response =
FeatureToggle::handle_pause(deps.as_mut(), &info, vec!["Feature2".to_string()])?;
let answer: HandleAnswer = from_binary(&response.data.unwrap())?;
assert_eq!(
answer,
HandleAnswer::Pause {
status: ResponseStatus::Success,
}
);
Ok(())
}
#[test]
fn test_require_not_paused() -> StdResult<()> {
let mut storage = MockStorage::new();
init_features(&mut storage)?;
assert!(
FeatureToggle::require_not_paused(&storage, vec!["Feature1".to_string()]).is_ok(),
"{:?}",
FeatureToggle::require_not_paused(&storage, vec!["Feature1".to_string()])
);
assert!(
FeatureToggle::require_not_paused(&storage, vec!["Feature3".to_string()]).is_err(),
"{:?}",
FeatureToggle::require_not_paused(&storage, vec!["Feature3".to_string()])
);
Ok(())
}
#[test]
fn test_add_remove_pausers() -> StdResult<()> {
let mut storage = MockStorage::new();
init_features(&mut storage)?;
let bob = Addr::unchecked("bob".to_string());
FeatureToggle::set_pauser(&mut storage, &bob)?;
assert!(
FeatureToggle::is_pauser(&storage, &bob)?,
"{:?}",
FeatureToggle::is_pauser(&storage, &bob)
);
FeatureToggle::remove_pauser(&mut storage, &bob);
assert!(
!FeatureToggle::is_pauser(&storage, &bob)?,
"{:?}",
FeatureToggle::is_pauser(&storage, &bob)
);
Ok(())
}
#[test]
fn test_deserialize_messages() {
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Debug, PartialEq)]
#[serde(rename_all = "snake_case")]
enum Features {
Var1,
Var2,
}
let handle_msg = b"{\"pause\":{\"features\":[\"var1\",\"var2\"]}}";
let query_msg = b"{\"status\":{\"features\": [\"var1\"]}}";
let query_msg_invalid = b"{\"status\":{\"features\": [\"var3\"]}}";
let parsed: FeatureToggleHandleMsg<Features> =
cosmwasm_std::from_slice(handle_msg).unwrap();
assert_eq!(
parsed,
FeatureToggleHandleMsg::Pause {
features: vec![Features::Var1, Features::Var2]
}
);
let parsed: FeatureToggleQueryMsg<Features> = cosmwasm_std::from_slice(query_msg).unwrap();
assert_eq!(
parsed,
FeatureToggleQueryMsg::Status {
features: vec![Features::Var1]
}
);
let parsed: StdResult<FeatureToggleQueryMsg<Features>> =
cosmwasm_std::from_slice(query_msg_invalid);
assert!(parsed.is_err());
}
}