Skip to content

Commit

Permalink
refactor: deriveAddressSeed and deriveAddress utilities
Browse files Browse the repository at this point in the history
**Problems**

- The distinction between `derive_address_seed` and `derive_address` was
  unclear and we were inconsistent in it:
  - We ended up applying address Merkle tree public key in both
    functions, which is confusing.
- Before this change, there was no TypeScript function for deriving
  **address seed**. There was only `deriveAddress`, but deriving the
  unified seed was a mystery for developers.
- We have two utilities for hashing and truncating to BN254:
  - `hash_to_bn254_field_size_be` - the older one, which:
    - Searches for a bump in a loop, adds it to the hash inputs and then
      truncates the hash. That doesn't make sense, because truncating
      the hash should be sufficient, adding a bump is unnecessary.
    - Another limitation is that it takes only one sequence of bytes,
      making it difficult to provide multiple inputs without
      concatenating them.
  - `hashv_to_bn254_field_size` - the newer one, which:
    - Just truncates the hash result, without the bump mechanism.
    - Takes 2D byte slice as input, making it possible to pass multiple
      inputs.

**Changes**

- Don't add MT pubkey in `derive_address_seed`. It's not a correct place
  for it to be applied. The distinction between `derive_address_seed`
  and `derive_address` should be:
  - `derive_address_seed` takes provided seeds (defined by the
    developer) and hashes them together with the program ID. This
    operation is done only in the third-party program.
  - `derive_address` takes the address seed (result of
    `address_address_seed`) and hashes it together with the address
    Merkle tree public key. This is done both in the third-party program
    and in light-system-program. light-system-program does that as a
    check whether the correct Merkle tree is used.
- Adjust the stateless.js API:
  - Provide `deriveAddressSeed` function.
  - Add unit tests, make sure that `deriveAddressSeed` and
    `deriveAddress` provide the same results as the equivalent functions
    in Rust SDK.
  • Loading branch information
vadorovsky committed Sep 30, 2024
1 parent e9befae commit 74b20a2
Show file tree
Hide file tree
Showing 9 changed files with 179 additions and 11 deletions.
64 changes: 61 additions & 3 deletions js/stateless.js/src/utils/address.ts
Original file line number Diff line number Diff line change
@@ -1,8 +1,17 @@
import { AccountMeta, PublicKey } from '@solana/web3.js';
import { hashToBn254FieldSizeBe } from './conversion';
import { hashToBn254FieldSizeBe, hashvToBn254FieldSizeBe } from './conversion';
import { defaultTestStateTreeAccounts } from '../constants';
import { getIndexOrAdd } from '../instruction';

export function deriveAddressSeed(
seeds: Uint8Array[],
programId: PublicKey,
): Uint8Array {
const combinedSeeds: Uint8Array[] = [programId.toBytes(), ...seeds];
const hash = hashvToBn254FieldSizeBe(combinedSeeds);
return hash;
}

/**
* Derive an address for a compressed account from a seed and a merkle tree
* public key.
Expand All @@ -18,7 +27,7 @@ export async function deriveAddress(
): Promise<PublicKey> {
const bytes = merkleTreePubkey.toBytes();
const combined = Buffer.from([...bytes, ...seed]);
const hash = await hashToBn254FieldSizeBe(combined);
const hash = hashToBn254FieldSizeBe(combined);

if (hash === null) {
throw new Error('DeriveAddressError');
Expand Down Expand Up @@ -115,14 +124,63 @@ if (import.meta.vitest) {
//@ts-ignore
const { it, expect, describe } = import.meta.vitest;

const programId = new PublicKey('7yucc7fL3JGbyMwg4neUaenNSdySS39hbAk89Ao3t1Hz');

describe('derive address seed', () => {
it('should derive a valid address seed', () => {
const seeds: Uint8Array[] = [
new TextEncoder().encode("foo"),
new TextEncoder().encode("bar")
];
expect(deriveAddressSeed(seeds, programId)).toStrictEqual(new Uint8Array([
0, 246, 150, 3, 192, 95, 53, 123, 56, 139, 206, 179, 253, 133, 115, 103, 120, 155,
251, 72, 250, 47, 117, 217, 118, 59, 174, 207, 49, 101, 201, 110
]));
});

it('should derive a valid address seed', () => {
const seeds: Uint8Array[] = [
new TextEncoder().encode("ayy"),
new TextEncoder().encode("lmao")
];
expect(deriveAddressSeed(seeds, programId)).toStrictEqual(new Uint8Array([
0, 202, 44, 25, 221, 74, 144, 92, 69, 168, 38, 19, 206, 208, 29, 162, 53, 27, 120,
214, 152, 116, 15, 107, 212, 168, 33, 121, 187, 10, 76, 233
]));
});
});

describe('deriveAddress function', () => {
it('should derive a valid address from a seed and a merkle tree public key', async () => {
const seed = new Uint8Array([1, 2, 3, 4]);
const seeds: Uint8Array[] = [
new TextEncoder().encode("foo"),
new TextEncoder().encode("bar")
];
const seed = deriveAddressSeed(seeds, programId);
const merkleTreePubkey = new PublicKey(
'11111111111111111111111111111111',
);
const derivedAddress = await deriveAddress(seed, merkleTreePubkey);
expect(derivedAddress).toBeInstanceOf(PublicKey);
expect(derivedAddress).toStrictEqual(new PublicKey(
'139uhyyBtEh4e1CBDJ68ooK5nCeWoncZf9HPyAfRrukA'
));
});

it('should derive a valid address from a seed and a merkle tree public key', async () => {
const seeds: Uint8Array[] = [
new TextEncoder().encode("ayy"),
new TextEncoder().encode("lmao")
];
const seed = deriveAddressSeed(seeds, programId);
const merkleTreePubkey = new PublicKey(
'11111111111111111111111111111111',
);
const derivedAddress = await deriveAddress(seed, merkleTreePubkey);
expect(derivedAddress).toBeInstanceOf(PublicKey);
expect(derivedAddress).toStrictEqual(new PublicKey(
'12bhHm6PQjbNmEn3Yu1Gq9k7XwVn2rZpzYokmLwbFazN'
));
});
});

Expand Down
40 changes: 38 additions & 2 deletions js/stateless.js/src/utils/conversion.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,25 @@ function isSmallerThanBn254FieldSizeBe(bytes: Buffer): boolean {
return bigint.lt(FIELD_SIZE);
}

export async function hashToBn254FieldSizeBe(
/**
* Hash the provided `bytes` with Keccak256 and ensure the result fits in the
* BN254 prime field by repeatedly hashing the inputs with various "bump seeds"
* and truncating the resulting hash to 31 bytes.
*
* @deprecated Use `hashvToBn254FieldSizeBe` instead.
*/
export function hashToBn254FieldSizeBe(
bytes: Buffer,
): Promise<[Buffer, number] | null> {
): [Buffer, number] | null {
// TODO(vadorovsky, affects-onchain): Get rid of the bump mechanism, it
// makes no sense. Doing the same as in the `hashvToBn254FieldSizeBe` below
// - overwriting the most significant byte with zero - is sufficient for
// truncation, it's also faster, doesn't force us to return `Option` and
// care about handling an error which is practically never returned.
//
// The reason we can't do it now is that it would affect on-chain programs.
// Once we can update programs, we can get rid of the seed bump (or even of
// this function all together in favor of the `hashv` variant).
let bumpSeed = 255;
while (bumpSeed >= 0) {
const inputWithBumpSeed = Buffer.concat([
Expand All @@ -51,6 +67,26 @@ export async function hashToBn254FieldSizeBe(
return null;
}

/**
* Hash the provided `bytes` with Keccak256 and ensure that the result fits in
* the BN254 prime field by truncating the resulting hash to 31 bytes.
*
* @param bytes Input bytes
*
* @returns Hash digest
*/
export function hashvToBn254FieldSizeBe(
bytes: Uint8Array[],
): Uint8Array {
let hasher = keccak_256.create();
for (var input of bytes) {
hasher.update(input);
}
let hash = hasher.digest();
hash[0] = 0;
return hash;
}

/** Mutates array in place */
export function pushUniqueItems<T>(items: T[], map: T[]): void {
items.forEach(item => {
Expand Down
6 changes: 5 additions & 1 deletion programs/system/src/invoke/append_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use anchor_lang::{
use light_hasher::Poseidon;
use light_heap::{bench_sbf_end, bench_sbf_start};
use light_macros::heap_neutral;
#[allow(deprecated)]
use light_utils::hash_to_bn254_field_size_be;

use crate::{
Expand Down Expand Up @@ -151,7 +152,9 @@ pub fn create_cpi_accounts_and_instruction_data<'a>(
});
hashed_merkle_tree = match hashed_pubkeys.iter().find(|x| x.0 == account_info.key()) {
Some(hashed_merkle_tree) => hashed_merkle_tree.1,
None => {
None =>
{
#[allow(deprecated)]
hash_to_bn254_field_size_be(&account_info.key().to_bytes())
.unwrap()
.0
Expand Down Expand Up @@ -212,6 +215,7 @@ pub fn create_cpi_accounts_and_instruction_data<'a>(
{
Some(hashed_owner) => hashed_owner.1,
None => {
#[allow(deprecated)]
let hashed_owner =
hash_to_bn254_field_size_be(&account.compressed_account.owner.to_bytes())
.unwrap()
Expand Down
2 changes: 2 additions & 0 deletions programs/system/src/invoke/verify_state_proof.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use light_concurrent_merkle_tree::zero_copy::ConcurrentMerkleTreeZeroCopy;
use light_hasher::Poseidon;
use light_indexed_merkle_tree::zero_copy::IndexedMerkleTreeZeroCopy;
use light_macros::heap_neutral;
#[allow(deprecated)]
use light_utils::hash_to_bn254_field_size_be;
use light_verifier::{
verify_create_addresses_and_merkle_proof_zkp, verify_create_addresses_zkp,
Expand Down Expand Up @@ -87,6 +88,7 @@ pub fn fetch_roots_address_merkle_tree<
/// Merkle tree pubkeys should be ordered for efficiency.
#[inline(never)]
#[heap_neutral]
#[allow(deprecated)]
#[allow(unused_mut)]
pub fn hash_input_compressed_accounts<'a, 'b, 'c: 'info, 'info>(
remaining_accounts: &'a [AccountInfo<'info>],
Expand Down
3 changes: 3 additions & 0 deletions programs/system/src/sdk/address.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
use std::collections::HashMap;

use anchor_lang::{err, solana_program::pubkey::Pubkey, Result};
#[allow(deprecated)]
use light_utils::hash_to_bn254_field_size_be;

use crate::{errors::SystemProgramError, NewAddressParams, NewAddressParamsPacked};

pub fn derive_address(merkle_tree_pubkey: &Pubkey, seed: &[u8; 32]) -> Result<[u8; 32]> {
#[allow(deprecated)]
let hash = match hash_to_bn254_field_size_be(
[merkle_tree_pubkey.to_bytes(), *seed].concat().as_slice(),
) {
Expand Down
3 changes: 3 additions & 0 deletions programs/system/src/sdk/compressed_account.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::collections::HashMap;

use anchor_lang::prelude::*;
use light_hasher::{Hasher, Poseidon};
#[allow(deprecated)]
use light_utils::hash_to_bn254_field_size_be;

#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
Expand Down Expand Up @@ -162,6 +163,7 @@ impl CompressedAccount {
&merkle_tree_pubkey: &Pubkey,
leaf_index: &u32,
) -> Result<[u8; 32]> {
#[allow(deprecated)]
self.hash_with_hashed_values::<H>(
&hash_to_bn254_field_size_be(&self.owner.to_bytes())
.unwrap()
Expand All @@ -188,6 +190,7 @@ mod tests {
/// 5. no address and no data
/// 6. no address, no data, no lamports
#[test]
#[allow(deprecated)]
fn test_compressed_account_hash() {
let owner = Keypair::new().pubkey();
let address = [1u8; 32];
Expand Down
52 changes: 51 additions & 1 deletion sdk/src/address.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use anchor_lang::{solana_program::pubkey::Pubkey, AnchorDeserialize, AnchorSerialize};
#[allow(deprecated)]
use light_utils::{hash_to_bn254_field_size_be, hashv_to_bn254_field_size_be};

use crate::merkle_context::{AddressMerkleContext, RemainingAccounts};
Expand Down Expand Up @@ -63,7 +64,6 @@ pub fn pack_new_address_params(
/// let address = derive_address(
/// &[b"my_compressed_account"],
/// &crate::ID,
/// &address_merkle_context,
/// );
/// ```
pub fn derive_address_seed(seeds: &[&[u8]], program_id: &Pubkey) -> [u8; 32] {
Expand All @@ -90,5 +90,55 @@ pub fn derive_address(
// PANICS: Not being able to find the bump for truncating the hash is
// practically impossible. Quite frankly, we should just remove that error
// inside.
#[allow(deprecated)]
hash_to_bn254_field_size_be(input.as_slice()).unwrap().0
}

#[cfg(test)]
mod test {
use light_macros::pubkey;

use super::*;

#[test]
fn test_derive_address_seed() {
let program_id = pubkey!("7yucc7fL3JGbyMwg4neUaenNSdySS39hbAk89Ao3t1Hz");

let address_seed = derive_address_seed(&[b"foo", b"bar"], &program_id);
assert_eq!(
address_seed,
[
0, 246, 150, 3, 192, 95, 53, 123, 56, 139, 206, 179, 253, 133, 115, 103, 120, 155,
251, 72, 250, 47, 117, 217, 118, 59, 174, 207, 49, 101, 201, 110
]
);

let address_seed = derive_address_seed(&[b"ayy", b"lmao"], &program_id);
assert_eq!(
address_seed,
[
0, 202, 44, 25, 221, 74, 144, 92, 69, 168, 38, 19, 206, 208, 29, 162, 53, 27, 120,
214, 152, 116, 15, 107, 212, 168, 33, 121, 187, 10, 76, 233
]
);
}

#[test]
fn test_derive_address() {
let address_merkle_context = AddressMerkleContext {
address_merkle_tree_pubkey: pubkey!("11111111111111111111111111111111"),
address_queue_pubkey: pubkey!("22222222222222222222222222222222222222222222"),
};
let program_id = pubkey!("7yucc7fL3JGbyMwg4neUaenNSdySS39hbAk89Ao3t1Hz");

let address_seed = derive_address_seed(&[b"foo", b"bar"], &program_id);
let address = derive_address(&address_seed, &address_merkle_context);
let expected_address = pubkey!("139uhyyBtEh4e1CBDJ68ooK5nCeWoncZf9HPyAfRrukA");
assert_eq!(address, expected_address.to_bytes());

let address_seed = derive_address_seed(&[b"ayy", b"lmao"], &program_id);
let address = derive_address(&address_seed, &address_merkle_context);
let expected_address = pubkey!("12bhHm6PQjbNmEn3Yu1Gq9k7XwVn2rZpzYokmLwbFazN");
assert_eq!(address, expected_address.to_bytes());
}
}
2 changes: 2 additions & 0 deletions sdk/src/compressed_account.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use anchor_lang::prelude::{
AccountInfo, AnchorDeserialize, AnchorSerialize, ProgramError, Pubkey, Result,
};
use light_hasher::{DataHasher, Discriminator, Hasher, Poseidon};
#[allow(deprecated)]
use light_utils::hash_to_bn254_field_size_be;

use crate::{
Expand Down Expand Up @@ -477,6 +478,7 @@ impl CompressedAccount {
&merkle_tree_pubkey: &Pubkey,
leaf_index: &u32,
) -> Result<[u8; 32]> {
#[allow(deprecated)]
self.hash_with_hashed_values::<H>(
&hash_to_bn254_field_size_be(&self.owner.to_bytes())
.unwrap()
Expand Down
18 changes: 14 additions & 4 deletions utils/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,20 @@ pub fn is_smaller_than_bn254_field_size_be(bytes: &[u8; 32]) -> bool {
bigint < ark_bn254::Fr::MODULUS.into()
}

/// Hashes the provided `bytes` with Keccak256 and ensures the result fits
/// in the BN254 prime field by repeatedly hashing the inputs with various
/// "bump seeds" and truncating the resulting hash to 31 bytes.
#[deprecated(note = "use `hashv_to_bn254_field_size_be` instead")]
pub fn hash_to_bn254_field_size_be(bytes: &[u8]) -> Option<([u8; 32], u8)> {
// TODO(vadorovsky, affects-onchain): Get rid of the bump mechanism, it
// makes no sense. Doing the same as in the `hashv_to_bn254_field_size_be`
// below - overwriting the most significant byte with zero - is sufficient
// for truncation, it's also faster, doesn't force us to return `Option`
// and care about handling an error which is practically never returned.
//
// The reason we can't do it now is that it would affect on-chain programs.
// Once we can update programs, we can get rid of the seed bump (or even of
// this function all together in favor of the `hashv` variant).
let mut bump_seed = [u8::MAX];
// Loops with decreasing bump seed to find a valid hash which is less than
// bn254 Fr modulo field size.
Expand All @@ -73,10 +86,7 @@ pub fn hash_to_bn254_field_size_be(bytes: &[u8]) -> Option<([u8; 32], u8)> {
}

/// Hashes the provided `bytes` with Keccak256 and ensures the result fits
/// in the BN254 prime field by repeatedly hashing the inputs with various
/// "bump seeds" and truncating the resulting hash to 31 bytes.
///
/// The attempted "bump seeds" are bytes from 255 to 0.
/// in the BN254 prime field by truncating the resulting hash to 31 bytes.
///
/// # Examples
///
Expand Down

0 comments on commit 74b20a2

Please sign in to comment.