Skip to content

Commit

Permalink
Market-contract v2: Rework query authentication framework
Browse files Browse the repository at this point in the history
  • Loading branch information
@ueco-jb
ueco-jb committed Dec 13, 2023
1 parent 6eea7c7 commit b85a823
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 88 deletions.
96 changes: 44 additions & 52 deletions contracts/lending/market_v2/src/contract.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use shade_protocol::{
},
contract_interfaces::{
oracles::{band::ReferenceData, oracle::QueryMsg::Price},
query_auth::helpers::authenticate_vk,
query_auth::helpers::{authenticate_permit, authenticate_vk, PermitAuthentication},
snip20::Snip20ReceiveMsg,
},
query_authentication::viewing_keys,
Expand All @@ -16,7 +16,10 @@ use shade_protocol::{

use crate::{
error::ContractError,
msg::{AuthQueryMsg, ExecuteMsg, InstantiateMsg, QueryMsg, ReceiveMsg, TotalDebtResponse},
msg::{
AuthPermit, Authentication, ExecuteMsg, InstantiateMsg, QueryMsg, ReceiveMsg,
TotalDebtResponse,
},
state::{debt, Config, CONFIG, VIEWING_KEY},
};

Expand Down Expand Up @@ -735,59 +738,66 @@ mod execute {
}
}

pub fn authenticate(deps: Deps, auth: Authentication, account: &Addr) -> StdResult<()> {
let config = CONFIG.load(deps.storage)?;
match auth {
Authentication::ViewingKey { key, address } => {
let address = deps.api.addr_validate(&address)?;
if !authenticate_vk(address.clone(), key, &deps.querier, &config.query_auth)? {
return Err(StdError::generic_err("Invalid Viewing Key"));
}
if &address != account {
return Err(StdError::generic_err(
"Trying to query using viewing key not matching the account",
));
}
Ok(())
}
Authentication::Permit(permit) => {
let res: PermitAuthentication<AuthPermit> =
authenticate_permit(permit, &deps.querier, config.query_auth)?;
if res.revoked {
return Err(StdError::generic_err("Permit Revoked"));
}
if res.sender != config.governance_contract {
return Err(StdError::generic_err(
"Unauthorized: Only credit agency can query using permit",
));
}
Ok(())
}
}
}

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> Result<Binary, ContractError> {
let res = match msg {
QueryMsg::TokensBalance {
account,
viewing_key,
authentication,
} => {
let config = CONFIG.load(deps.storage)?;
authenticate_vk(
account.clone(),
viewing_key,
&deps.querier,
&config.query_auth,
)?;
authenticate(deps, authentication, &account)?;
to_binary(&query::tokens_balance(deps, env, account.to_string())?)?
}
QueryMsg::Withdrawable {
account,
viewing_key,
authentication,
} => {
let config = CONFIG.load(deps.storage)?;
authenticate_vk(
account.clone(),
viewing_key,
&deps.querier,
&config.query_auth,
)?;
authenticate(deps, authentication, &account)?;
to_binary(&query::withdrawable(deps, env, account.to_string())?)?
}
QueryMsg::Borrowable {
account,
viewing_key,
authentication,
} => {
let config = CONFIG.load(deps.storage)?;
authenticate_vk(
account.clone(),
viewing_key,
&deps.querier,
&config.query_auth,
)?;
authenticate(deps, authentication, &account)?;
to_binary(&query::borrowable(deps, env, account.to_string())?)?
}
QueryMsg::CreditLine {
account,
viewing_key,
authentication,
} => {
let config = CONFIG.load(deps.storage)?;
authenticate_vk(
account.clone(),
viewing_key,
&deps.querier,
&config.query_auth,
)?;
authenticate(deps, authentication, &account)?;
to_binary(&query::credit_line(deps, env, account)?)?
}
QueryMsg::Configuration {} => to_binary(&query::config(deps, env)?)?,
Expand All @@ -804,24 +814,6 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> Result<Binary, ContractErro
let (total, multiplier) = debt::total(deps.storage)?;
to_binary(&TotalDebtResponse { total, multiplier })?
}
QueryMsg::WithPermit { permit, query_msg } => {
// Handle AuthQueryMsg here
match query_msg {
AuthQueryMsg::TokensBalance { account } => {
to_binary(&query::tokens_balance(deps, env, account)?)?
}
AuthQueryMsg::Withdrawable { account } => {
to_binary(&query::withdrawable(deps, env, account)?)?
}
AuthQueryMsg::Borrowable { account } => {
to_binary(&query::borrowable(deps, env, account)?)?
}
AuthQueryMsg::CreditLine { account } => {
let account = deps.api.addr_validate(&account)?;
to_binary(&query::credit_line(deps, env, account)?)?
}
}
}
};
Ok(res)
}
Expand Down
60 changes: 24 additions & 36 deletions contracts/lending/market_v2/src/msg.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,15 @@ pub enum CreditAgencyExecuteMsg {
EnterMarket { account: String },
}

#[cw_serde]
pub struct AuthPermit {}

#[cw_serde]
pub enum Authentication {
ViewingKey { key: String, address: String },
Permit(QueryPermit),
}

#[cw_serde]
#[derive(QueryResponses)]
pub enum QueryMsg {
Expand Down Expand Up @@ -123,55 +132,34 @@ pub enum QueryMsg {
TotalDebt {},
/// Returns TokensBalanceResponse
#[returns(TokensBalanceResponse)]
TokensBalance { account: Addr, viewing_key: String },
TokensBalance {
account: Addr,
authentication: Authentication,
},
/// Returns the amount that the given account can withdraw
#[returns(Coin)]
Withdrawable { account: Addr, viewing_key: String },
Withdrawable {
account: Addr,
authentication: Authentication,
},
/// Returns the amount that the given account can borrow
#[returns(Coin)]
Borrowable { account: Addr, viewing_key: String },
Borrowable {
account: Addr,
authentication: Authentication,
},
/// Returns CreditLineResponse
#[returns(lending_utils::credit_line::CreditLineResponse)]
CreditLine { account: Addr, viewing_key: String },

/// Queries that requires viewing permit
#[returns(AuthQueryResponse)]
WithPermit {
permit: QueryPermit,
query_msg: AuthQueryMsg,
CreditLine {
account: Addr,
authentication: Authentication,
},
}

impl Query for QueryMsg {
const BLOCK_SIZE: usize = 256;
}

#[cw_serde]
#[derive(QueryResponses)]
pub enum AuthQueryMsg {
/// Returns TokensBalanceResponse
#[returns(TokensBalanceResponse)]
TokensBalance { account: String },
/// Returns the amount that the given account can withdraw
#[returns(Coin)]
Withdrawable { account: String },
/// Returns the amount that the given account can borrow
#[returns(Coin)]
Borrowable { account: String },
/// Returns CreditLineResponse
#[returns(lending_utils::credit_line::CreditLineResponse)]
CreditLine { account: String },
}

// Define an enum for all possible return types of AuthQueryMsg
#[cw_serde]
enum AuthQueryResponse {
TokensBalance(TokensBalanceResponse),
Withdrawable(Coin),
Borrowable(Coin),
CreditLine(lending_utils::credit_line::CreditLineResponse),
}

#[cw_serde]
pub enum QueryTotalCreditLine {
TotalCreditLine { account: String },
Expand Down

0 comments on commit b85a823

Please sign in to comment.