From 5269a3b16c6a3964836fde846883221907f9c349 Mon Sep 17 00:00:00 2001 From: mymiracle0118 Date: Thu, 18 Apr 2024 19:52:44 +0000 Subject: [PATCH] complete main structure and logic --- Cargo.lock | 14 ++++++++ Cargo.toml | 1 + src/contract.rs | 88 ++++++++++++++++++++++++++++++++++++++++++++++--- src/error.rs | 3 ++ src/msg.rs | 28 +++++++++++++++- src/state.rs | 2 ++ 6 files changed, 131 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3670994..4c83777 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -264,6 +264,19 @@ dependencies = [ "thiserror", ] +[[package]] +name = "cw721" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3c4d286625ccadc957fe480dd3bdc54ada19e0e6b5b9325379db3130569e914" +dependencies = [ + "cosmwasm-schema", + "cosmwasm-std", + "cw-utils", + "schemars", + "serde", +] + [[package]] name = "der" version = "0.7.9" @@ -472,6 +485,7 @@ dependencies = [ "cw-multi-test", "cw-storage-plus", "cw2", + "cw721", "schemars", "serde", "thiserror", diff --git a/Cargo.toml b/Cargo.toml index 5d403fe..4b95a86 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,6 +42,7 @@ cosmwasm-std = { version = "1.5.0", features = [ ] } cw-storage-plus = "1.1.0" cw2 = "1.1.1" +cw721 = "0.18.0" schemars = "0.8.15" serde = { version = "1.0.189", default-features = false, features = ["derive"] } thiserror = { version = "1.0.49" } diff --git a/src/contract.rs b/src/contract.rs index 32b123b..6ca2b8a 100644 --- a/src/contract.rs +++ b/src/contract.rs @@ -1,11 +1,12 @@ #[cfg(not(feature = "library"))] use cosmwasm_std::entry_point; -use cosmwasm_std::{to_json_binary, Binary, Deps, DepsMut, Env, MessageInfo, Response, StdResult}; +use cosmwasm_std::{to_json_binary, Binary, Deps, DepsMut, Env, MessageInfo, Response, StdResult, Addr, CosmosMsg, WasmMsg, SubMsg}; use cw2::set_contract_version; +use cw721::Cw721ExecuteMsg; use crate::error::ContractError; -use crate::msg::{ExecuteMsg, GetCountResponse, InstantiateMsg, QueryMsg}; -use crate::state::{State, STATE}; +use crate::msg::{ExecuteMsg, GetCountResponse, InstantiateMsg, QueryMsg, AllNftsResponse, NftContractAddrResponse}; +use crate::state::{State, STATE, NFTS, NFT_CONTRACT_ADDR}; // version info for migration info const CONTRACT_NAME: &str = "crates.io:jarvis-airdrop"; @@ -34,17 +35,22 @@ pub fn instantiate( #[cfg_attr(not(feature = "library"), entry_point)] pub fn execute( deps: DepsMut, - _env: Env, + env: Env, info: MessageInfo, msg: ExecuteMsg, ) -> Result { match msg { ExecuteMsg::Increment {} => execute::increment(deps), ExecuteMsg::Reset { count } => execute::reset(deps, info, count), + ExecuteMsg::SetNftContractAddr { addr } => execute::set_nft_contract_addr(deps, env, info, addr), + ExecuteMsg::ReceiveNft { sender, token_id, msg } => execute::receive_nft(deps, env, info, token_id), + ExecuteMsg::SendNfts { allocations } => execute::send_nfts(deps, env, info, allocations), } } pub mod execute { + use std::ops::Sub; + use super::*; pub fn increment(deps: DepsMut) -> Result { @@ -66,12 +72,76 @@ pub mod execute { })?; Ok(Response::new().add_attribute("action", "reset")) } + + pub fn set_nft_contract_addr( + deps: DepsMut, + _env: Env, + info: MessageInfo, + addr: String, + ) -> Result { + // Optionally, add authorization checks here to ensure only specific addresses can update this + let nft_contract_addr = deps.api.addr_validate(&addr)?; + NFT_CONTRACT_ADDR.save(deps.storage, &nft_contract_addr)?; + Ok(Response::new().add_attribute("action", "set_nft_contract_addr").add_attribute("address", addr)) + } + + pub fn receive_nft( + deps: DepsMut, + _env: Env, + _info: MessageInfo, + token_id: String, + ) -> Result { + let mut nfts = NFTS.load(deps.storage).unwrap_or_default(); + nfts.push(token_id); + NFTS.save(deps.storage, &nfts)?; + + Ok(Response::new().add_attribute("action", "receive_nft")) + } + + + pub fn send_nfts( + deps: DepsMut, + _env: Env, + info: MessageInfo, + allocations: Vec<(Addr, u32)>, + ) -> Result { + let nft_contract_addr = NFT_CONTRACT_ADDR.load(deps.storage)?; + let mut nfts = NFTS.load(deps.storage)?; + let mut response = Response::new().add_attribute("action", "send_nfts"); + + for (recipient, amount) in allocations { + for _ in 0..amount { + if let Some(token_id) = nfts.pop() { + // Create a transfer message for the cw721 NFT + let transfer_msg = Cw721ExecuteMsg::TransferNft { + recipient: recipient.to_string(), + token_id: token_id, + }; + + let msg = CosmosMsg::Wasm(WasmMsg::Execute { + contract_addr: nft_contract_addr.clone().to_string(), + msg: to_json_binary(&transfer_msg)?, + funds: vec![], + }); + + response.messages.push(SubMsg::new(msg)); + } else { + return Err(ContractError::InsufficientNFTs {}); + } + } + } + + NFTS.save(deps.storage, &nfts)?; + Ok(response) + } } #[cfg_attr(not(feature = "library"), entry_point)] pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult { match msg { QueryMsg::GetCount {} => to_json_binary(&query::count(deps)?), + QueryMsg::GetAllNfts { } => to_json_binary(&query::all_nfts(deps)?), + QueryMsg::GetNftContractAddr { } => to_json_binary(&query::nft_contract_addr(deps)?), } } @@ -82,6 +152,16 @@ pub mod query { let state = STATE.load(deps.storage)?; Ok(GetCountResponse { count: state.count }) } + + pub fn all_nfts(deps: Deps) -> StdResult { + let nfts = NFTS.load(deps.storage)?; + Ok(AllNftsResponse { nfts }) + } + + pub fn nft_contract_addr(deps: Deps) -> StdResult { + let nft_contract_addr = NFT_CONTRACT_ADDR.load(deps.storage)?; + Ok(NftContractAddrResponse { nft_contract_addr }) + } } #[cfg(test)] diff --git a/src/error.rs b/src/error.rs index 4a69d8f..a5e82fe 100644 --- a/src/error.rs +++ b/src/error.rs @@ -8,6 +8,9 @@ pub enum ContractError { #[error("Unauthorized")] Unauthorized {}, + + #[error("Insufficient NFTs")] + InsufficientNFTs {}, // Add any other custom errors you like here. // Look at https://docs.rs/thiserror/1.0.21/thiserror/ for details. } diff --git a/src/msg.rs b/src/msg.rs index 0edfa32..accfe3a 100644 --- a/src/msg.rs +++ b/src/msg.rs @@ -1,4 +1,5 @@ use cosmwasm_schema::{cw_serde, QueryResponses}; +use cosmwasm_std::{Addr, Binary}; #[cw_serde] pub struct InstantiateMsg { @@ -9,6 +10,15 @@ pub struct InstantiateMsg { pub enum ExecuteMsg { Increment {}, Reset { count: i32 }, + ReceiveNft { + sender: String, + token_id: String, + msg: Binary, + }, + SendNfts { + allocations: Vec<(Addr, u32)>, // Each tuple contains an address and the number of NFTs to send + }, + SetNftContractAddr { addr: String }, } #[cw_serde] @@ -17,10 +27,26 @@ pub enum QueryMsg { // GetCount returns the current count as a json-encoded number #[returns(GetCountResponse)] GetCount {}, + + #[returns(AllNftsResponse)] + GetAllNfts {}, + + #[returns(NftContractAddrResponse)] + GetNftContractAddr {}, } // We define a custom struct for each query response #[cw_serde] pub struct GetCountResponse { - pub count: i32, + pub count: i32, +} + +#[cw_serde] +pub struct AllNftsResponse { + pub nfts: Vec, } + +#[cw_serde] +pub struct NftContractAddrResponse { + pub nft_contract_addr: Addr, +} \ No newline at end of file diff --git a/src/state.rs b/src/state.rs index bad9202..b88cb87 100644 --- a/src/state.rs +++ b/src/state.rs @@ -11,3 +11,5 @@ pub struct State { } pub const STATE: Item = Item::new("state"); +pub const NFTS: Item> = Item::new("nfts"); +pub const NFT_CONTRACT_ADDR: Item = Item::new("nft_contract_addr");