Audit revision v2

packages/ics721/src/error.rs

@@ -55,14 +55,17 @@ pub enum ContractError {
     #[error("Transfer Doesn't contain any action, no redemption or creation")]
     InvalidTransferNoAction,
 
-    #[error("Couldn't find nft contract for this class id: {0}")]
+    #[error("Couldn't find nft contract for class id: {0}")]
     NoNftContractForClassId(String),
 
     #[error("Unknown nft contract: {child_collection}, Class Id: {class_id}, Token ID: {token_id} => NFT contract: {cw721_addr}")]
-    NoClassIdForNftContract {
+    NoNftContractMatch {
         child_collection: String,
         class_id: String,
         token_id: String,
         cw721_addr: String,
     },
+
+    #[error("Couldn't find class id for nft contract: {0}")]
+    NoClassIdForNftContract(String),
 }

packages/ics721/src/ibc.rs

@@ -12,9 +12,10 @@ use crate::{
     helpers::ack_callback_msg,
     ibc_helpers::{ack_fail, ack_success, try_get_ack_error, validate_order_and_version},
     ibc_packet_receive::receive_ibc_packet,
+    query::{load_class_id_for_nft_contract, load_nft_contract_for_class_id},
     state::{
-        CLASS_ID_TO_NFT_CONTRACT, INCOMING_CLASS_TOKEN_TO_CHANNEL, INCOMING_PROXY,
-        NFT_CONTRACT_TO_CLASS_ID, OUTGOING_CLASS_TOKEN_TO_CHANNEL, OUTGOING_PROXY, TOKEN_METADATA,
+        INCOMING_CLASS_TOKEN_TO_CHANNEL, INCOMING_PROXY, OUTGOING_CLASS_TOKEN_TO_CHANNEL,
+        OUTGOING_PROXY, TOKEN_METADATA,
     },
     ContractError,
 };
@@ -121,7 +122,8 @@ where
         } else {
             let msg: NonFungibleTokenPacketData = from_json(&ack.original_packet.data)?;
 
-            let nft_contract = CLASS_ID_TO_NFT_CONTRACT.load(deps.storage, msg.class_id.clone())?;
+            let nft_contract =
+                load_nft_contract_for_class_id(deps.storage, msg.class_id.to_string())?;
             // Burn all of the tokens being transfered out that were
             // previously transfered in on this channel.
             let burn_notices = msg.token_ids.iter().cloned().try_fold(
@@ -188,7 +190,8 @@ where
         error: &str,
     ) -> Result<IbcBasicResponse, ContractError> {
         let message: NonFungibleTokenPacketData = from_json(&packet.data)?;
-        let nft_contract = CLASS_ID_TO_NFT_CONTRACT.load(deps.storage, message.class_id.clone())?;
+        let nft_contract =
+            load_nft_contract_for_class_id(deps.storage, message.class_id.to_string())?;
         let sender = deps.api.addr_validate(&message.sender)?;
 
         let messages = message
@@ -241,15 +244,18 @@ where
                 // `ACK_AND_DO_NOTHING`.
 
                 let res = parse_reply_instantiate_data(reply)?;
-                let cw721_addr = deps.api.addr_validate(&res.contract_address)?;
+                let nft_contract = deps.api.addr_validate(&res.contract_address)?;
 
                 // cw721 addr has already been stored, here just check whether it exists, otherwise a NotFound error is thrown
-                let class_id = NFT_CONTRACT_TO_CLASS_ID.load(deps.storage, cw721_addr.clone())?;
-
-                Ok(Response::default()
-                    .add_attribute("method", "instantiate_cw721_reply")
-                    .add_attribute("class_id", class_id)
-                    .add_attribute("cw721_addr", cw721_addr))
+                match load_class_id_for_nft_contract(deps.storage, &nft_contract)? {
+                    Some(class_id) => Ok(Response::default()
+                        .add_attribute("method", "instantiate_cw721_reply")
+                        .add_attribute("class_id", class_id)
+                        .add_attribute("cw721_addr", nft_contract)),
+                    None => Err(ContractError::NoClassIdForNftContract(
+                        nft_contract.to_string(),
+                    )),
+                }
             }
             INSTANTIATE_OUTGOING_PROXY_REPLY_ID => {
                 let res = parse_reply_instantiate_data(reply)?;

packages/ics721/src/ibc_packet_receive.rs

@@ -13,7 +13,8 @@ use crate::{
     ibc::ACK_AND_DO_NOTHING_REPLY_ID,
     ibc_helpers::{get_endpoint_prefix, try_pop_source_prefix},
     msg::{CallbackMsg, ExecuteMsg},
-    state::{CLASS_ID_TO_NFT_CONTRACT, CW721_CODE_ID, OUTGOING_CLASS_TOKEN_TO_CHANNEL, PO},
+    query::load_nft_contract_for_class_id,
+    state::{CW721_CODE_ID, OUTGOING_CLASS_TOKEN_TO_CHANNEL, PO},
     token_types::{VoucherCreation, VoucherRedemption},
     ContractError,
 };
@@ -228,8 +229,7 @@ fn create_callback_msg(
         let nft_contract = if is_redemption {
             // If its a redemption, it means we already have the contract address in storage
 
-            CLASS_ID_TO_NFT_CONTRACT
-                .load(deps.storage, local_class_id.clone())
+            load_nft_contract_for_class_id(deps.storage, local_class_id.to_string())
                 .map_err(|_| ContractError::NoNftContractForClassId(local_class_id.to_string()))
         } else {
             // If its a creation action, we can use the instantiate2 function to get the nft contract

packages/ics721/src/query.rs

@@ -1,13 +1,12 @@
-use cosmwasm_std::{to_json_binary, Addr, Binary, Deps, Env, Order, StdResult};
-use cw_storage_plus::Map;
+use cosmwasm_std::{to_json_binary, Addr, Binary, Deps, Env, Order, StdError, StdResult, Storage};
+use cw_storage_plus::{Bound, Map};
 
 use crate::{
     msg::QueryMsg,
     state::{
-        UniversalAllNftInfoResponse, ADMIN_USED_FOR_CW721, CLASS_ID_TO_CLASS,
-        CLASS_ID_TO_NFT_CONTRACT, CW721_CODE_ID, INCOMING_CLASS_TOKEN_TO_CHANNEL, INCOMING_PROXY,
-        NFT_CONTRACT_TO_CLASS_ID, OUTGOING_CLASS_TOKEN_TO_CHANNEL, OUTGOING_PROXY, PO,
-        TOKEN_METADATA,
+        UniversalAllNftInfoResponse, ADMIN_USED_FOR_CW721, CLASS_ID_AND_NFT_CONTRACT_INFO,
+        CLASS_ID_TO_CLASS, CW721_CODE_ID, INCOMING_CLASS_TOKEN_TO_CHANNEL, INCOMING_PROXY,
+        OUTGOING_CLASS_TOKEN_TO_CHANNEL, OUTGOING_PROXY, PO, TOKEN_METADATA,
     },
 };
 use ics721_types::token_types::{Class, ClassId, ClassToken, Token, TokenId};
@@ -16,28 +15,28 @@ pub trait Ics721Query {
     fn query(&self, deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
         match msg {
             QueryMsg::ClassId { contract } => {
-                to_json_binary(&self.query_class_id_for_nft_contract(deps, contract)?)
+                to_json_binary(&query_class_id_for_nft_contract(deps, contract)?)
             }
             QueryMsg::NftContract { class_id } => {
-                to_json_binary(&self.query_nft_contract_for_class_id(deps, class_id)?)
+                to_json_binary(&query_nft_contract_for_class_id(deps.storage, class_id)?)
             }
             QueryMsg::ClassMetadata { class_id } => {
-                to_json_binary(&self.query_class_metadata(deps, class_id)?)
+                to_json_binary(&query_class_metadata(deps, class_id)?)
             }
             QueryMsg::TokenMetadata { class_id, token_id } => {
-                to_json_binary(&self.query_token_metadata(deps, class_id, token_id)?)
+                to_json_binary(&query_token_metadata(deps, class_id, token_id)?)
             }
             QueryMsg::Owner { class_id, token_id } => {
-                to_json_binary(&self.query_owner(deps, class_id, token_id)?)
+                to_json_binary(&query_owner(deps, class_id, token_id)?)
             }
             QueryMsg::Pauser {} => to_json_binary(&PO.query_pauser(deps.storage)?),
             QueryMsg::Paused {} => to_json_binary(&PO.query_paused(deps.storage)?),
             QueryMsg::OutgoingProxy {} => to_json_binary(&OUTGOING_PROXY.load(deps.storage)?),
             QueryMsg::IncomingProxy {} => to_json_binary(&INCOMING_PROXY.load(deps.storage)?),
-            QueryMsg::Cw721CodeId {} => to_json_binary(&self.query_cw721_code_id(deps)?),
+            QueryMsg::Cw721CodeId {} => to_json_binary(&query_cw721_code_id(deps)?),
             QueryMsg::Cw721Admin {} => to_json_binary(&ADMIN_USED_FOR_CW721.load(deps.storage)?),
             QueryMsg::NftContracts { start_after, limit } => {
-                to_json_binary(&self.query_nft_contracts(deps, start_after, limit)?)
+                to_json_binary(&query_nft_contracts(deps, start_after, limit)?)
             }
             QueryMsg::OutgoingChannels { start_after, limit } => to_json_binary(&query_channels(
                 deps,
@@ -53,97 +52,120 @@ pub trait Ics721Query {
             )?),
         }
     }
+}
 
-    fn query_class_id_for_nft_contract(
-        &self,
-        deps: Deps,
-        contract: String,
-    ) -> StdResult<Option<ClassId>> {
-        let contract = deps.api.addr_validate(&contract)?;
-        NFT_CONTRACT_TO_CLASS_ID.may_load(deps.storage, contract)
-    }
+pub fn query_class_id_for_nft_contract(deps: Deps, contract: String) -> StdResult<Option<ClassId>> {
+    let contract = deps.api.addr_validate(&contract)?;
+    load_class_id_for_nft_contract(deps.storage, &contract)
+}
 
-    fn query_nft_contract_for_class_id(
-        &self,
-        deps: Deps,
-        class_id: String,
-    ) -> StdResult<Option<Addr>> {
-        CLASS_ID_TO_NFT_CONTRACT.may_load(deps.storage, ClassId::new(class_id))
-    }
+pub fn load_class_id_for_nft_contract(
+    storage: &dyn Storage,
+    contract: &Addr,
+) -> StdResult<Option<ClassId>> {
+    CLASS_ID_AND_NFT_CONTRACT_INFO
+        .idx
+        .address
+        .item(storage, contract.clone())
+        .map(|e| e.map(|(_, c)| c.class_id))
+}
 
-    fn query_class_metadata(&self, deps: Deps, class_id: String) -> StdResult<Option<Class>> {
-        CLASS_ID_TO_CLASS.may_load(deps.storage, ClassId::new(class_id))
-    }
+pub fn query_nft_contract_for_class_id(
+    storage: &dyn Storage,
+    class_id: String,
+) -> StdResult<Option<Addr>> {
+    // Convert the class_id string to ClassId type if necessary
+    let class_id_key = ClassId::new(class_id);
 
-    fn query_token_metadata(
-        &self,
-        deps: Deps,
-        class_id: String,
-        token_id: String,
-    ) -> StdResult<Option<Token>> {
-        let token_id = TokenId::new(token_id);
-        let class_id = ClassId::new(class_id);
+    // Query the IndexedMap using the class_id index
+    CLASS_ID_AND_NFT_CONTRACT_INFO
+        .idx
+        .class_id
+        .item(storage, class_id_key)
+        .map(|e| e.map(|(_, v)| v.address))
+}
 
-        let Some(token_metadata) =
-            TOKEN_METADATA.may_load(deps.storage, (class_id.clone(), token_id.clone()))?
-        else {
-            // Token metadata is set unconditionaly on mint. If we have no
-            // metadata entry, we have no entry for this token at all.
-            return Ok(None);
-        };
-        let Some(token_contract) = CLASS_ID_TO_NFT_CONTRACT.may_load(deps.storage, class_id)?
-        else {
-            debug_assert!(false, "token_metadata != None => token_contract != None");
-            return Ok(None);
-        };
-        let UniversalAllNftInfoResponse { info, .. } = deps.querier.query_wasm_smart(
-            token_contract,
-            &cw721::Cw721QueryMsg::AllNftInfo {
-                token_id: token_id.clone().into(),
-                include_expired: None,
-            },
-        )?;
-        Ok(Some(Token {
-            id: token_id,
-            uri: info.token_uri,
-            data: token_metadata,
-        }))
-    }
+pub fn load_nft_contract_for_class_id(storage: &dyn Storage, class_id: String) -> StdResult<Addr> {
+    query_nft_contract_for_class_id(storage, class_id.clone())?.map_or_else(
+        || {
+            Err(StdError::NotFound {
+                kind: format!("NFT contract not found for class id {}", class_id),
+            })
+        },
+        Ok,
+    )
+}
 
-    fn query_owner(
-        &self,
-        deps: Deps,
-        class_id: String,
-        token_id: String,
-    ) -> StdResult<cw721::OwnerOfResponse> {
-        let class_uri = CLASS_ID_TO_NFT_CONTRACT.load(deps.storage, ClassId::new(class_id))?;
-        let resp: cw721::OwnerOfResponse = deps.querier.query_wasm_smart(
-            class_uri,
-            &cw721::Cw721QueryMsg::OwnerOf {
-                token_id,
-                include_expired: None,
-            },
-        )?;
-        Ok(resp)
-    }
+pub fn query_class_metadata(deps: Deps, class_id: String) -> StdResult<Option<Class>> {
+    CLASS_ID_TO_CLASS.may_load(deps.storage, ClassId::new(class_id))
+}
 
-    fn query_cw721_code_id(&self, deps: Deps) -> StdResult<u64> {
-        CW721_CODE_ID.load(deps.storage)
-    }
+pub fn query_token_metadata(
+    deps: Deps,
+    class_id: String,
+    token_id: String,
+) -> StdResult<Option<Token>> {
+    let token_id = TokenId::new(token_id);
+    let class_id = ClassId::new(class_id);
 
-    fn query_nft_contracts(
-        &self,
-        deps: Deps,
-        start_after: Option<ClassId>,
-        limit: Option<u32>,
-    ) -> StdResult<Vec<(String, Addr)>> {
-        cw_paginate_storage::paginate_map(
-            deps,
-            &CLASS_ID_TO_NFT_CONTRACT,
-            start_after,
-            limit,
-            Order::Ascending,
-        )
+    let Some(token_metadata) =
+        TOKEN_METADATA.may_load(deps.storage, (class_id.clone(), token_id.clone()))?
+    else {
+        // Token metadata is set unconditionaly on mint. If we have no
+        // metadata entry, we have no entry for this token at all.
+        return Ok(None);
+    };
+    let Some(nft_contract) = query_nft_contract_for_class_id(deps.storage, class_id.to_string())?
+    else {
+        debug_assert!(false, "token_metadata != None => token_contract != None");
+        return Ok(None);
+    };
+    let UniversalAllNftInfoResponse { info, .. } = deps.querier.query_wasm_smart(
+        nft_contract,
+        &cw721::Cw721QueryMsg::AllNftInfo {
+            token_id: token_id.clone().into(),
+            include_expired: None,
+        },
+    )?;
+    Ok(Some(Token {
+        id: token_id,
+        uri: info.token_uri,
+        data: token_metadata,
+    }))
+}
+
+pub fn query_owner(
+    deps: Deps,
+    class_id: String,
+    token_id: String,
+) -> StdResult<cw721::OwnerOfResponse> {
+    let nft_contract = load_nft_contract_for_class_id(deps.storage, class_id)?;
+    let resp: cw721::OwnerOfResponse = deps.querier.query_wasm_smart(
+        nft_contract,
+        &cw721::Cw721QueryMsg::OwnerOf {
+            token_id,
+            include_expired: None,
+        },
+    )?;
+    Ok(resp)
+}
+
+pub fn query_cw721_code_id(deps: Deps) -> StdResult<u64> {
+    CW721_CODE_ID.load(deps.storage)
+}
+
+pub fn query_nft_contracts(
+    deps: Deps,
+    start_after: Option<ClassId>,
+    limit: Option<u32>,
+) -> StdResult<Vec<(String, Addr)>> {
+    let start = start_after.map(|s| Bound::ExclusiveRaw(s.to_string().into()));
+    let all = CLASS_ID_AND_NFT_CONTRACT_INFO
+        .range(deps.storage, start, None, Order::Ascending)
+        .map(|item| item.map(|(k, v)| (k, v.address)));
+    match limit {
+        Some(limit) => all.take(limit as usize).collect(),
+        None => all.collect(),
     }
 }