| 
 | 1 | +use {  | 
 | 2 | +    crate::{get_elgamal_registry_address, id},  | 
 | 3 | +    solana_program::{  | 
 | 4 | +        instruction::{AccountMeta, Instruction},  | 
 | 5 | +        program_error::ProgramError,  | 
 | 6 | +        pubkey::Pubkey,  | 
 | 7 | +        system_program, sysvar,  | 
 | 8 | +    },  | 
 | 9 | +    solana_zk_sdk::zk_elgamal_proof_program::{  | 
 | 10 | +        instruction::ProofInstruction, proof_data::PubkeyValidityProofData,  | 
 | 11 | +    },  | 
 | 12 | +    spl_token_confidential_transfer_proof_extraction::instruction::{ProofData, ProofLocation},  | 
 | 13 | +};  | 
 | 14 | + | 
 | 15 | +#[derive(Clone, Debug, PartialEq)]  | 
 | 16 | +#[repr(u8)]  | 
 | 17 | +pub enum RegistryInstruction {  | 
 | 18 | +    /// Initialize an ElGamal public key registry.  | 
 | 19 | +    ///  | 
 | 20 | +    /// 0. `[writable]` The account to be created  | 
 | 21 | +    /// 1. `[signer]` The wallet address (will also be the owner address for the  | 
 | 22 | +    ///    registry account)  | 
 | 23 | +    /// 2. `[]` System program  | 
 | 24 | +    /// 3. `[]` Instructions sysvar if `VerifyPubkeyValidity` is included in the  | 
 | 25 | +    ///    same transaction or context state account if `VerifyPubkeyValidity`  | 
 | 26 | +    ///    is pre-verified into a context state account.  | 
 | 27 | +    /// 4. `[]` (Optional) Record account if the accompanying proof is to be  | 
 | 28 | +    ///    read from a record account.  | 
 | 29 | +    CreateRegistry {  | 
 | 30 | +        /// Relative location of the `ProofInstruction::PubkeyValidityProof`  | 
 | 31 | +        /// instruction to the `CreateElGamalRegistry` instruction in the  | 
 | 32 | +        /// transaction. If the offset is `0`, then use a context state account  | 
 | 33 | +        /// for the proof.  | 
 | 34 | +        proof_instruction_offset: i8,  | 
 | 35 | +    },  | 
 | 36 | +    /// Update an ElGamal public key registry with a new ElGamal public key.  | 
 | 37 | +    ///  | 
 | 38 | +    /// 0. `[writable]` The ElGamal registry account  | 
 | 39 | +    /// 1. `[]` Instructions sysvar if `VerifyPubkeyValidity` is included in the  | 
 | 40 | +    ///    same transaction or context state account if `VerifyPubkeyValidity`  | 
 | 41 | +    ///    is pre-verified into a context state account.  | 
 | 42 | +    /// 2. `[]` (Optional) Record account if the accompanying proof is to be  | 
 | 43 | +    ///    read from a record account.  | 
 | 44 | +    /// 3. `[signer]` The owner of the ElGamal public key registry  | 
 | 45 | +    UpdateRegistry {  | 
 | 46 | +        /// Relative location of the `ProofInstruction::PubkeyValidityProof`  | 
 | 47 | +        /// instruction to the `UpdateElGamalRegistry` instruction in the  | 
 | 48 | +        /// transaction. If the offset is `0`, then use a context state account  | 
 | 49 | +        /// for the proof.  | 
 | 50 | +        proof_instruction_offset: i8,  | 
 | 51 | +    },  | 
 | 52 | +}  | 
 | 53 | + | 
 | 54 | +impl RegistryInstruction {  | 
 | 55 | +    /// Unpacks a byte buffer into a `RegistryInstruction`  | 
 | 56 | +    pub fn unpack(input: &[u8]) -> Result<Self, ProgramError> {  | 
 | 57 | +        let (&tag, rest) = input  | 
 | 58 | +            .split_first()  | 
 | 59 | +            .ok_or(ProgramError::InvalidInstructionData)?;  | 
 | 60 | + | 
 | 61 | +        Ok(match tag {  | 
 | 62 | +            0 => {  | 
 | 63 | +                let proof_instruction_offset =  | 
 | 64 | +                    *rest.first().ok_or(ProgramError::InvalidInstructionData)?;  | 
 | 65 | +                Self::CreateRegistry {  | 
 | 66 | +                    proof_instruction_offset: proof_instruction_offset as i8,  | 
 | 67 | +                }  | 
 | 68 | +            }  | 
 | 69 | +            1 => {  | 
 | 70 | +                let proof_instruction_offset =  | 
 | 71 | +                    *rest.first().ok_or(ProgramError::InvalidInstructionData)?;  | 
 | 72 | +                Self::UpdateRegistry {  | 
 | 73 | +                    proof_instruction_offset: proof_instruction_offset as i8,  | 
 | 74 | +                }  | 
 | 75 | +            }  | 
 | 76 | +            _ => return Err(ProgramError::InvalidInstructionData),  | 
 | 77 | +        })  | 
 | 78 | +    }  | 
 | 79 | + | 
 | 80 | +    /// Packs a `RegistryInstruction` into a byte buffer.  | 
 | 81 | +    pub fn pack(&self) -> Vec<u8> {  | 
 | 82 | +        let mut buf = vec![];  | 
 | 83 | +        match self {  | 
 | 84 | +            Self::CreateRegistry {  | 
 | 85 | +                proof_instruction_offset,  | 
 | 86 | +            } => {  | 
 | 87 | +                buf.push(0);  | 
 | 88 | +                buf.extend_from_slice(&proof_instruction_offset.to_le_bytes());  | 
 | 89 | +            }  | 
 | 90 | +            Self::UpdateRegistry {  | 
 | 91 | +                proof_instruction_offset,  | 
 | 92 | +            } => {  | 
 | 93 | +                buf.push(1);  | 
 | 94 | +                buf.extend_from_slice(&proof_instruction_offset.to_le_bytes());  | 
 | 95 | +            }  | 
 | 96 | +        };  | 
 | 97 | +        buf  | 
 | 98 | +    }  | 
 | 99 | +}  | 
 | 100 | + | 
 | 101 | +/// Create a `RegistryInstruction::CreateRegistry` instruction  | 
 | 102 | +pub fn create_registry(  | 
 | 103 | +    owner_address: &Pubkey,  | 
 | 104 | +    proof_location: ProofLocation<PubkeyValidityProofData>,  | 
 | 105 | +) -> Result<Vec<Instruction>, ProgramError> {  | 
 | 106 | +    let elgamal_registry_address = get_elgamal_registry_address(owner_address, &id());  | 
 | 107 | + | 
 | 108 | +    let mut accounts = vec![  | 
 | 109 | +        AccountMeta::new(elgamal_registry_address, false),  | 
 | 110 | +        AccountMeta::new_readonly(*owner_address, true),  | 
 | 111 | +        AccountMeta::new_readonly(system_program::id(), false),  | 
 | 112 | +    ];  | 
 | 113 | +    let proof_instruction_offset = proof_instruction_offset(&mut accounts, proof_location);  | 
 | 114 | + | 
 | 115 | +    let mut instructions = vec![Instruction {  | 
 | 116 | +        program_id: id(),  | 
 | 117 | +        accounts,  | 
 | 118 | +        data: RegistryInstruction::CreateRegistry {  | 
 | 119 | +            proof_instruction_offset,  | 
 | 120 | +        }  | 
 | 121 | +        .pack(),  | 
 | 122 | +    }];  | 
 | 123 | +    append_zk_elgamal_proof(&mut instructions, proof_location)?;  | 
 | 124 | +    Ok(instructions)  | 
 | 125 | +}  | 
 | 126 | + | 
 | 127 | +/// Create a `RegistryInstruction::UpdateRegistry` instruction  | 
 | 128 | +pub fn update_registry(  | 
 | 129 | +    owner_address: &Pubkey,  | 
 | 130 | +    proof_location: ProofLocation<PubkeyValidityProofData>,  | 
 | 131 | +) -> Result<Vec<Instruction>, ProgramError> {  | 
 | 132 | +    let elgamal_registry_address = get_elgamal_registry_address(owner_address, &id());  | 
 | 133 | + | 
 | 134 | +    let mut accounts = vec![AccountMeta::new(elgamal_registry_address, false)];  | 
 | 135 | +    let proof_instruction_offset = proof_instruction_offset(&mut accounts, proof_location);  | 
 | 136 | +    accounts.push(AccountMeta::new_readonly(*owner_address, true));  | 
 | 137 | + | 
 | 138 | +    let mut instructions = vec![Instruction {  | 
 | 139 | +        program_id: id(),  | 
 | 140 | +        accounts,  | 
 | 141 | +        data: RegistryInstruction::UpdateRegistry {  | 
 | 142 | +            proof_instruction_offset,  | 
 | 143 | +        }  | 
 | 144 | +        .pack(),  | 
 | 145 | +    }];  | 
 | 146 | +    append_zk_elgamal_proof(&mut instructions, proof_location)?;  | 
 | 147 | +    Ok(instructions)  | 
 | 148 | +}  | 
 | 149 | + | 
 | 150 | +/// Takes a `ProofLocation`, updates the list of accounts, and returns a  | 
 | 151 | +/// suitable proof location  | 
 | 152 | +fn proof_instruction_offset(  | 
 | 153 | +    accounts: &mut Vec<AccountMeta>,  | 
 | 154 | +    proof_location: ProofLocation<PubkeyValidityProofData>,  | 
 | 155 | +) -> i8 {  | 
 | 156 | +    match proof_location {  | 
 | 157 | +        ProofLocation::InstructionOffset(proof_instruction_offset, proof_data) => {  | 
 | 158 | +            accounts.push(AccountMeta::new_readonly(sysvar::instructions::id(), false));  | 
 | 159 | +            if let ProofData::RecordAccount(record_address, _) = proof_data {  | 
 | 160 | +                accounts.push(AccountMeta::new_readonly(*record_address, false));  | 
 | 161 | +            }  | 
 | 162 | +            proof_instruction_offset.into()  | 
 | 163 | +        }  | 
 | 164 | +        ProofLocation::ContextStateAccount(context_state_account) => {  | 
 | 165 | +            accounts.push(AccountMeta::new_readonly(*context_state_account, false));  | 
 | 166 | +            0  | 
 | 167 | +        }  | 
 | 168 | +    }  | 
 | 169 | +}  | 
 | 170 | + | 
 | 171 | +/// Takes a `RegistryInstruction` and appends the pubkey validity proof  | 
 | 172 | +/// instruction  | 
 | 173 | +fn append_zk_elgamal_proof(  | 
 | 174 | +    instructions: &mut Vec<Instruction>,  | 
 | 175 | +    proof_data_location: ProofLocation<PubkeyValidityProofData>,  | 
 | 176 | +) -> Result<(), ProgramError> {  | 
 | 177 | +    if let ProofLocation::InstructionOffset(proof_instruction_offset, proof_data) =  | 
 | 178 | +        proof_data_location  | 
 | 179 | +    {  | 
 | 180 | +        let proof_instruction_offset: i8 = proof_instruction_offset.into();  | 
 | 181 | +        if proof_instruction_offset != 1 {  | 
 | 182 | +            return Err(ProgramError::InvalidArgument);  | 
 | 183 | +        }  | 
 | 184 | +        match proof_data {  | 
 | 185 | +            ProofData::InstructionData(data) => instructions  | 
 | 186 | +                .push(ProofInstruction::VerifyPubkeyValidity.encode_verify_proof(None, data)),  | 
 | 187 | +            ProofData::RecordAccount(address, offset) => instructions.push(  | 
 | 188 | +                ProofInstruction::VerifyPubkeyValidity  | 
 | 189 | +                    .encode_verify_proof_from_account(None, address, offset),  | 
 | 190 | +            ),  | 
 | 191 | +        }  | 
 | 192 | +    }  | 
 | 193 | +    Ok(())  | 
 | 194 | +}  | 
0 commit comments