mpl_token_auth_rules/state/v2/constraint/
pubkey_tree_match.rs

1use solana_program::{msg, pubkey::PUBKEY_BYTES};
2
3use crate::{
4    error::RuleSetError,
5    state::RuleResult,
6    state::{
7        try_from_bytes,
8        v2::{Constraint, ConstraintType, Str32, HEADER_SECTION},
9        Header,
10    },
11    utils::compute_merkle_root,
12};
13
14/// Constraing representing a test where a `Pubkey` must be a member of the Merkle tree in the rule.
15///
16/// This constraint requires `PayloadType` values of `PayloadType::Pubkey` and `PayloadType::MerkleProof`.
17/// The `field` values in the Rule are used to locate them in the `Payload`. The `Pubkey` and the proof
18/// are used to calculate a Merkle root which is compared against the root stored in the rule.
19pub struct PubkeyTreeMatch<'a> {
20    /// The field in the `Payload` to be compared when looking for the `Pubkey`.
21    pub pubkey_field: &'a Str32,
22    /// The field in the `Payload` to be compared when looking for the Merkle proof.
23    pub proof_field: &'a Str32,
24    /// The root of the Merkle tree.
25    pub root: &'a [u8; PUBKEY_BYTES],
26}
27
28impl<'a> PubkeyTreeMatch<'a> {
29    /// Deserialize a constraint from a byte array.
30    pub fn from_bytes(bytes: &'a [u8]) -> Result<Self, RuleSetError> {
31        let pubkey_field = try_from_bytes::<Str32>(0, Str32::SIZE, bytes)?;
32        let mut cursor = Str32::SIZE;
33
34        let proof_field = try_from_bytes::<Str32>(cursor, Str32::SIZE, bytes)?;
35        cursor += Str32::SIZE;
36
37        let root = try_from_bytes::<[u8; 32]>(cursor, PUBKEY_BYTES, bytes)?;
38
39        Ok(Self {
40            pubkey_field,
41            proof_field,
42            root,
43        })
44    }
45
46    /// Serialize a constraint into a byte array.
47    pub fn serialize(
48        pubkey_field: String,
49        proof_field: String,
50        root: &[u8; PUBKEY_BYTES],
51    ) -> Result<Vec<u8>, RuleSetError> {
52        let length = (Str32::SIZE + Str32::SIZE + PUBKEY_BYTES) as u32;
53        let mut data = Vec::with_capacity(HEADER_SECTION + length as usize);
54
55        // Header
56        Header::serialize(ConstraintType::PubkeyTreeMatch, length, &mut data);
57
58        // Constraint
59        // - pubkey_field
60        let mut field_bytes = [0u8; Str32::SIZE];
61        field_bytes[..pubkey_field.len()].copy_from_slice(pubkey_field.as_bytes());
62        data.extend(field_bytes);
63        // - proof_field
64        let mut field_bytes = [0u8; Str32::SIZE];
65        field_bytes[..proof_field.len()].copy_from_slice(proof_field.as_bytes());
66        data.extend(field_bytes);
67        // - root
68        data.extend_from_slice(root);
69
70        Ok(data)
71    }
72}
73
74impl<'a> Constraint<'a> for PubkeyTreeMatch<'a> {
75    fn constraint_type(&self) -> ConstraintType {
76        ConstraintType::PubkeyTreeMatch
77    }
78
79    fn validate(
80        &self,
81        _accounts: &std::collections::HashMap<
82            solana_program::pubkey::Pubkey,
83            &solana_program::account_info::AccountInfo,
84        >,
85        payload: &crate::payload::Payload,
86        _update_rule_state: bool,
87        _rule_set_state_pda: &Option<&solana_program::account_info::AccountInfo>,
88        _rule_authority: &Option<&solana_program::account_info::AccountInfo>,
89    ) -> RuleResult {
90        msg!("Validating PubkeyTreeMatch");
91
92        // Get the `Pubkey` we are checking from the payload.
93        let leaf = match payload.get_pubkey(&self.pubkey_field.to_string()) {
94            Some(pubkey) => pubkey,
95            _ => return RuleResult::Error(RuleSetError::MissingPayloadValue.into()),
96        };
97
98        // Get the Merkle proof from the payload.
99        let merkle_proof = match payload.get_merkle_proof(&self.proof_field.to_string()) {
100            Some(merkle_proof) => merkle_proof,
101            _ => return RuleResult::Error(RuleSetError::MissingPayloadValue.into()),
102        };
103
104        // Check if the computed hash (root) is equal to the root in the rule.
105        let computed_root = compute_merkle_root(leaf, merkle_proof);
106
107        if computed_root == *self.root {
108            RuleResult::Success(self.constraint_type().to_error())
109        } else {
110            RuleResult::Failure(self.constraint_type().to_error())
111        }
112    }
113}