Add Support for AirdropOp for Bulk Distribution of Tokens (#15)

* initial airdrop opcode

* added airdrop tests

* added chunked test

Author: zfedoran

api/src/cvm/messages/airdrop.rs

@@ -0,0 +1,60 @@
+use steel::*;
+
+use crate::utils;
+use crate::types::Hash;
+use crate::cvm::{
+    CodeVmAccount,
+    VirtualDurableNonce, 
+    VirtualTimelockAccount
+};
+
+pub fn compact_airdrop_message(
+    src_timelock_address: &Pubkey,
+    dst_timelock_addresses: &[Pubkey],
+    amount: u64,
+    vdn: &VirtualDurableNonce,
+) -> Hash {
+    let mut msg = Vec::new();
+
+    msg.push(b"airdrop" as &[u8]);
+    msg.push(src_timelock_address.as_ref());
+    msg.push(vdn.address.as_ref());
+    msg.push(vdn.value.as_ref());
+
+    // Store the little-endian bytes in a local variable so it won't go out of scope
+    let amount_bytes = amount.to_le_bytes();
+    msg.push(&amount_bytes);
+
+    // Push each destination pubkey
+    for dst_pubkey in dst_timelock_addresses {
+        msg.push(dst_pubkey.as_ref());
+    }
+
+    utils::hashv(&msg)
+}
+
+pub fn create_airdrop_message(
+    vm: &CodeVmAccount,
+    src_vta: &VirtualTimelockAccount,
+    destinations: &[Pubkey],
+    amount: u64,
+    vdn: &VirtualDurableNonce,
+) -> Hash {
+
+    let src_timelock_address = src_vta.get_timelock_address(
+        &vm.get_mint(),
+        &vm.get_authority(),
+        vm.get_lock_duration(),
+    );
+
+    let src_token_address = src_vta.get_token_address(
+        &src_timelock_address,
+    );
+
+    compact_airdrop_message(
+        &src_token_address,
+        destinations,
+        amount,
+        vdn,
+    )
+}

api/src/cvm/messages/mod.rs

@@ -1,5 +1,7 @@
+mod airdrop;
 mod transfer;
 mod withdraw;
 
+pub use airdrop::*;
 pub use transfer::*;
 pub use withdraw::*;

api/src/opcode.rs

@@ -15,6 +15,8 @@ pub enum Opcode {
   ExternalRelayOp = 20,
 
   ConditionalTransferOp = 12,
+
+  AirdropOp = 30,
 }
 
 instruction!(Opcode, TransferOp);
@@ -24,6 +26,7 @@ instruction!(Opcode, ExternalTransferOp);
 instruction!(Opcode, ExternalWithdrawOp);
 instruction!(Opcode, ExternalRelayOp);
 instruction!(Opcode, ConditionalTransferOp);
+instruction!(Opcode, AirdropOp);
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug, Pod, Zeroable)]
@@ -210,3 +213,37 @@ pub struct ParsedConditionalTransferOp {
     pub signature: [u8; 64],
     pub amount: u64,
 }
+
+#[repr(C)]
+#[derive(Clone, Copy, Debug, Pod, Zeroable)]
+pub struct AirdropOp {
+    pub signature: [u8; 64],
+    pub amount: [u8; 8], // Pack u64 as [u8; 8]
+    pub count: u8,       // Up to 255 airdrops in a single tx (but CU limit will be hit first)
+}
+
+impl AirdropOp {
+    /// Converts the byte array `amount` to `u64`.
+    pub fn to_struct(&self) -> Result<ParsedAirdropOp, std::io::Error> {
+        Ok(ParsedAirdropOp {
+            signature: self.signature,
+            amount: u64::from_le_bytes(self.amount),
+            count: self.count,
+        })
+    }
+
+    /// Creates `AirdropOp` from the parsed struct by converting `u64` back to byte array.
+    pub fn from_struct(parsed: ParsedAirdropOp) -> Self {
+        AirdropOp {
+            signature: parsed.signature,
+            amount: parsed.amount.to_le_bytes(),
+            count: parsed.count,
+        }
+    }
+}
+
+pub struct ParsedAirdropOp {
+    pub signature: [u8; 64],
+    pub amount: u64,
+    pub count: u8,
+}

program/src/instruction/exec.rs

@@ -69,6 +69,8 @@ pub fn process_exec(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult
 
         Opcode::ConditionalTransferOp  => process_conditional_transfer(&ctx, &args),
 
+        Opcode::AirdropOp              => process_airdrop(&ctx, &args),
+
         _ => Err(ProgramError::InvalidInstructionData),
     }?;
 

program/src/opcode/airdrop.rs

@@ -0,0 +1,162 @@
+use code_vm_api::prelude::*;
+use steel::*;
+
+use crate::ExecContext;
+
+/*
+    This instruction is used to transfer tokens from *one* virtual account to a
+    number of virtual accounts. The signature of the source account is required
+    to authorize the transfer.
+
+    Extra accounts required by this instruction:
+    
+    | # | R/W | Type         | Req | PDA | Name   | Description  |
+    |---|-----|------------- |-----|-----|--------|--------------|
+    |...| The same as the vm_exec instruction.                   |
+    |---|-----|------------- |-----|-----|--------|--------------|
+    | 6 |     | <None>       |     |     |        |              |
+    | 7 |     | <None>       |     |     |        |              |
+    | 8 |     | <None>       |     |     |        |              |
+    | 9 |     | <None>       |     |     |        |              |
+    |10 |     | <None>       |     |     |        |              |
+
+
+    Instruction data:
+
+    0. signature: [u8;64]  - The opcode to execute.
+    1. amount: [u64]       - The account_indicies of the virtual accounts to use.
+    2. count: [u8]         - The number of destinations.
+*/
+pub fn process_airdrop(
+    ctx: &ExecContext,
+    data: &ExecIxData,
+) -> ProgramResult {
+
+    let vm = load_vm(ctx.vm_info)?;
+    let args = AirdropOp::try_from_bytes(&data.data)?.to_struct()?;
+
+    let mem_indicies = &data.mem_indicies;
+    let mem_banks = &data.mem_banks;
+    let num_accounts = 2 + (args.count as usize);
+
+    check_condition(
+        mem_indicies.len() == num_accounts,
+        "invalid number of memory indicies",
+    )?;
+
+    check_condition(
+        mem_banks.len() == num_accounts,
+        "invalid number of memory banks",
+    )?;
+
+    let nonce_index = mem_indicies[0];
+    let nonce_mem = mem_banks[0];
+
+    let src_index = mem_indicies[1];
+    let src_mem = mem_banks[1];
+
+    let vm_mem = ctx.get_banks();
+
+    check_condition(
+        vm_mem[nonce_mem as usize].is_some(),
+        "the nonce memory account must be provided",
+    )?;
+
+    check_condition(
+        vm_mem[src_mem as usize].is_some(),
+        "the source memory account must be provided",
+    )?;
+
+    let nonce_mem_info = vm_mem[nonce_mem as usize].unwrap();
+    let src_mem_info = vm_mem[src_mem as usize].unwrap();
+
+    let va = try_read(&nonce_mem_info, nonce_index)?;
+    let mut vdn = va.into_inner_nonce().unwrap();
+
+    let va = try_read(&src_mem_info, src_index)?;
+    let mut src_vta = va.into_inner_timelock().unwrap();
+
+    let total_amount = args.amount
+        .checked_mul(args.count as u64)
+        .ok_or(ProgramError::ArithmeticOverflow)?;
+
+    if src_vta.balance < total_amount {
+        return Err(ProgramError::InsufficientFunds);
+    }
+
+    src_vta.balance = src_vta.balance
+        .checked_sub(total_amount)
+        .ok_or(ProgramError::ArithmeticOverflow)?;
+
+    let mut dst_pubkeys = Vec::new();
+    for i in 0..args.count as usize {
+        let dst_index = mem_indicies[2 + i];
+        let dst_mem = mem_banks[2 + i];
+
+        check_condition(
+            vm_mem[dst_mem as usize].is_some(),
+            "a destination memory account must be provided",
+        )?;
+
+        let dst_mem_info = vm_mem[dst_mem as usize].unwrap();
+
+        let va = try_read(&dst_mem_info, dst_index)?;
+        let mut dst_vta = va.into_inner_timelock().unwrap();
+
+        // Check if this destination is actually the source.
+        let is_same_account = (src_mem == dst_mem) && (src_index == dst_index);
+        if is_same_account {
+            // If the source is also in the destinations list, it receives the airdrop as well.
+            src_vta.balance = src_vta.balance
+                .checked_add(args.amount)
+                .ok_or(ProgramError::ArithmeticOverflow)?;
+
+        } else {
+            // Normal destination: add the airdrop to its balance
+            dst_vta.balance = dst_vta.balance
+                .checked_add(args.amount)
+                .ok_or(ProgramError::ArithmeticOverflow)?;
+
+            // Write the updated destination back
+            try_write(
+                dst_mem_info,
+                dst_index,
+                &VirtualAccount::Timelock(dst_vta)
+            )?;
+        }
+
+        dst_pubkeys.push(dst_vta.owner);
+    }
+
+    let hash = create_airdrop_message(
+        &vm,
+        &src_vta,
+        &dst_pubkeys,
+        args.amount,
+        &vdn,
+    );
+
+    sig_verify(
+        src_vta.owner.as_ref(),
+        args.signature.as_ref(),
+        hash.as_ref(),
+    )?;
+
+    vdn.value = vm.get_current_poh();
+
+    // Finally, write back the updated source (which now includes
+    // any airdrop shares if the source was also in the destination list).
+    try_write(
+        src_mem_info,
+        src_index,
+        &VirtualAccount::Timelock(src_vta)
+    )?;
+
+    try_write(
+        nonce_mem_info,
+        nonce_index,
+        &VirtualAccount::Nonce(vdn)
+    )?;
+
+    Ok(())
+}

program/src/opcode/mod.rs

@@ -1,3 +1,4 @@
+mod airdrop;
 mod conditional_transfer;
 mod external_relay;
 mod external_transfer;
@@ -6,6 +7,7 @@ mod relay;
 mod transfer;
 mod withdraw;
 
+pub use airdrop::*;
 pub use conditional_transfer::*;
 pub use external_relay::*;
 pub use external_transfer::*;