use chia_protocol::{Bytes32, CoinSpend};
use chia_sdk_driver::{Nft, Puzzle, SpendContext};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct NftOwnershipProof {
pub launcher_id: Bytes32,
pub owner_puzzle_hash: Bytes32,
pub attributed_did: Option<Bytes32>,
pub nft_coin_id: Bytes32,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum GatingError {
NotAnNft,
WrongOwner { expected: Bytes32, got: Bytes32 },
WrongCollection {
required: Bytes32,
got: Option<Bytes32>,
},
WrongNft { required: Bytes32, got: Bytes32 },
}
impl GatingError {
pub fn code(&self) -> &'static str {
match self {
GatingError::NotAnNft => "NOT_AN_NFT",
GatingError::WrongOwner { .. } => "WRONG_OWNER",
GatingError::WrongCollection { .. } => "WRONG_COLLECTION",
GatingError::WrongNft { .. } => "WRONG_NFT",
}
}
}
impl core::fmt::Display for GatingError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
GatingError::NotAnNft => write!(f, "coin spend is not an NFT spend"),
GatingError::WrongOwner { expected, got } => {
write!(f, "NFT owner {got} is not the claimed owner {expected}")
}
GatingError::WrongCollection { required, got } => {
write!(f, "NFT collection {got:?} is not the required {required}")
}
GatingError::WrongNft { required, got } => {
write!(f, "NFT {got} is not the required {required}")
}
}
}
}
pub fn read_nft_ownership(parent_spend: &CoinSpend) -> Result<NftOwnershipProof, GatingError> {
let mut ctx = SpendContext::new();
let puzzle_ptr = ctx
.alloc(&parent_spend.puzzle_reveal)
.map_err(|_| GatingError::NotAnNft)?;
let solution_ptr = ctx
.alloc(&parent_spend.solution)
.map_err(|_| GatingError::NotAnNft)?;
let puzzle = Puzzle::parse(&ctx, puzzle_ptr);
let nft = Nft::parse_child(&mut ctx, parent_spend.coin, puzzle, solution_ptr)
.map_err(|_| GatingError::NotAnNft)?
.ok_or(GatingError::NotAnNft)?;
Ok(NftOwnershipProof {
launcher_id: nft.info.launcher_id,
owner_puzzle_hash: nft.info.p2_puzzle_hash,
attributed_did: nft.info.current_owner,
nft_coin_id: nft.coin.coin_id(),
})
}
pub fn prove_nft_ownership(
parent_spend: &CoinSpend,
claimed_owner_puzzle_hash: Bytes32,
required_nft: Option<Bytes32>,
) -> Result<NftOwnershipProof, GatingError> {
let proof = read_nft_ownership(parent_spend)?;
if proof.owner_puzzle_hash != claimed_owner_puzzle_hash {
return Err(GatingError::WrongOwner {
expected: claimed_owner_puzzle_hash,
got: proof.owner_puzzle_hash,
});
}
if let Some(required) = required_nft {
if proof.launcher_id != required {
return Err(GatingError::WrongNft {
required,
got: proof.launcher_id,
});
}
}
Ok(proof)
}
pub fn prove_collection_membership(
parent_spend: &CoinSpend,
claimed_owner_puzzle_hash: Bytes32,
required_did: Bytes32,
) -> Result<NftOwnershipProof, GatingError> {
let proof = prove_nft_ownership(parent_spend, claimed_owner_puzzle_hash, None)?;
if proof.attributed_did != Some(required_did) {
return Err(GatingError::WrongCollection {
required: required_did,
got: proof.attributed_did,
});
}
Ok(proof)
}