use crate::hash::CryptoHash;
use crate::types::AccountId;
use borsh::{BorshDeserialize, BorshSerialize};
use near_crypto::PublicKey;
use std::mem::size_of;
pub(crate) const ACCOUNT_DATA_SEPARATOR: u8 = b',';
pub(crate) mod col {
pub const ACCOUNT: u8 = 0;
pub const CONTRACT_CODE: u8 = 1;
pub const ACCESS_KEY: u8 = 2;
pub const RECEIVED_DATA: u8 = 3;
pub const POSTPONED_RECEIPT_ID: u8 = 4;
pub const PENDING_DATA_COUNT: u8 = 5;
pub const POSTPONED_RECEIPT: u8 = 6;
pub const DELAYED_RECEIPT_INDICES: u8 = 7;
pub const DELAYED_RECEIPT: u8 = 8;
pub const CONTRACT_DATA: u8 = 9;
pub const NON_DELAYED_RECEIPT_COLUMNS: [(u8, &str); 8] = [
(ACCOUNT, "Account"),
(CONTRACT_CODE, "ContractCode"),
(ACCESS_KEY, "AccessKey"),
(RECEIVED_DATA, "ReceivedData"),
(POSTPONED_RECEIPT_ID, "PostponedReceiptId"),
(PENDING_DATA_COUNT, "PendingDataCount"),
(POSTPONED_RECEIPT, "PostponedReceipt"),
(CONTRACT_DATA, "ContractData"),
];
}
#[derive(Debug, Clone, PartialEq, Eq, BorshDeserialize, BorshSerialize)]
pub enum TrieKey {
Account { account_id: AccountId },
ContractCode { account_id: AccountId },
AccessKey { account_id: AccountId, public_key: PublicKey },
ReceivedData { receiver_id: AccountId, data_id: CryptoHash },
PostponedReceiptId { receiver_id: AccountId, data_id: CryptoHash },
PendingDataCount { receiver_id: AccountId, receipt_id: CryptoHash },
PostponedReceipt { receiver_id: AccountId, receipt_id: CryptoHash },
DelayedReceiptIndices,
DelayedReceipt { index: u64 },
ContractData { account_id: AccountId, key: Vec<u8> },
}
trait Byte {
fn len(self) -> usize;
}
impl Byte for u8 {
fn len(self) -> usize {
1
}
}
impl TrieKey {
pub fn len(&self) -> usize {
match self {
TrieKey::Account { account_id } => col::ACCOUNT.len() + account_id.len(),
TrieKey::ContractCode { account_id } => col::CONTRACT_CODE.len() + account_id.len(),
TrieKey::AccessKey { account_id, public_key } => {
col::ACCESS_KEY.len() * 2 + account_id.len() + public_key.len()
}
TrieKey::ReceivedData { receiver_id, data_id } => {
col::RECEIVED_DATA.len()
+ receiver_id.len()
+ ACCOUNT_DATA_SEPARATOR.len()
+ data_id.as_ref().len()
}
TrieKey::PostponedReceiptId { receiver_id, data_id } => {
col::POSTPONED_RECEIPT_ID.len()
+ receiver_id.len()
+ ACCOUNT_DATA_SEPARATOR.len()
+ data_id.as_ref().len()
}
TrieKey::PendingDataCount { receiver_id, receipt_id } => {
col::PENDING_DATA_COUNT.len()
+ receiver_id.len()
+ ACCOUNT_DATA_SEPARATOR.len()
+ receipt_id.as_ref().len()
}
TrieKey::PostponedReceipt { receiver_id, receipt_id } => {
col::POSTPONED_RECEIPT.len()
+ receiver_id.len()
+ ACCOUNT_DATA_SEPARATOR.len()
+ receipt_id.as_ref().len()
}
TrieKey::DelayedReceiptIndices => col::DELAYED_RECEIPT_INDICES.len(),
TrieKey::DelayedReceipt { .. } => col::DELAYED_RECEIPT.len() + size_of::<u64>(),
TrieKey::ContractData { account_id, key } => {
col::CONTRACT_DATA.len()
+ account_id.len()
+ ACCOUNT_DATA_SEPARATOR.len()
+ key.len()
}
}
}
pub fn append_into(&self, buf: &mut Vec<u8>) {
let expected_len = self.len();
let start_len = buf.len();
buf.reserve(self.len());
match self {
TrieKey::Account { account_id } => {
buf.push(col::ACCOUNT);
buf.extend(account_id.as_ref().as_bytes());
}
TrieKey::ContractCode { account_id } => {
buf.push(col::CONTRACT_CODE);
buf.extend(account_id.as_ref().as_bytes());
}
TrieKey::AccessKey { account_id, public_key } => {
buf.push(col::ACCESS_KEY);
buf.extend(account_id.as_ref().as_bytes());
buf.push(col::ACCESS_KEY);
buf.extend(public_key.try_to_vec().unwrap());
}
TrieKey::ReceivedData { receiver_id, data_id } => {
buf.push(col::RECEIVED_DATA);
buf.extend(receiver_id.as_ref().as_bytes());
buf.push(ACCOUNT_DATA_SEPARATOR);
buf.extend(data_id.as_ref());
}
TrieKey::PostponedReceiptId { receiver_id, data_id } => {
buf.push(col::POSTPONED_RECEIPT_ID);
buf.extend(receiver_id.as_ref().as_bytes());
buf.push(ACCOUNT_DATA_SEPARATOR);
buf.extend(data_id.as_ref());
}
TrieKey::PendingDataCount { receiver_id, receipt_id } => {
buf.push(col::PENDING_DATA_COUNT);
buf.extend(receiver_id.as_ref().as_bytes());
buf.push(ACCOUNT_DATA_SEPARATOR);
buf.extend(receipt_id.as_ref());
}
TrieKey::PostponedReceipt { receiver_id, receipt_id } => {
buf.push(col::POSTPONED_RECEIPT);
buf.extend(receiver_id.as_ref().as_bytes());
buf.push(ACCOUNT_DATA_SEPARATOR);
buf.extend(receipt_id.as_ref());
}
TrieKey::DelayedReceiptIndices => {
buf.push(col::DELAYED_RECEIPT_INDICES);
}
TrieKey::DelayedReceipt { index } => {
buf.push(col::DELAYED_RECEIPT_INDICES);
buf.extend(&index.to_le_bytes());
}
TrieKey::ContractData { account_id, key } => {
buf.push(col::CONTRACT_DATA);
buf.extend(account_id.as_ref().as_bytes());
buf.push(ACCOUNT_DATA_SEPARATOR);
buf.extend(key);
}
};
debug_assert_eq!(expected_len, buf.len() - start_len);
}
pub fn to_vec(&self) -> Vec<u8> {
let mut buf = Vec::with_capacity(self.len());
self.append_into(&mut buf);
buf
}
}
pub mod trie_key_parsers {
use super::*;
pub fn parse_public_key_from_access_key_key(
raw_key: &[u8],
account_id: &AccountId,
) -> Result<PublicKey, std::io::Error> {
let prefix_len = col::ACCESS_KEY.len() * 2 + account_id.len();
if raw_key.len() < prefix_len {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"raw key is too short for TrieKey::AccessKey",
));
}
PublicKey::try_from_slice(&raw_key[prefix_len..])
}
pub fn parse_data_key_from_contract_data_key<'a>(
raw_key: &'a [u8],
account_id: &AccountId,
) -> Result<&'a [u8], std::io::Error> {
let prefix_len = col::CONTRACT_DATA.len() + account_id.len() + ACCOUNT_DATA_SEPARATOR.len();
if raw_key.len() < prefix_len {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"raw key is too short for TrieKey::ContractData",
));
}
Ok(&raw_key[prefix_len..])
}
pub fn parse_account_id_prefix<'a>(
column: u8,
raw_key: &'a [u8],
) -> Result<&'a [u8], std::io::Error> {
let prefix = std::slice::from_ref(&column);
if let Some(tail) = raw_key.strip_prefix(prefix) {
Ok(tail)
} else {
Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"raw key is does not start with a proper column marker",
))
}
}
fn parse_account_id_from_slice(
data: &[u8],
trie_key: &str,
) -> Result<AccountId, std::io::Error> {
std::str::from_utf8(data)
.map_err(|_| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!(
"raw key AccountId has invalid UTF-8 format to be TrieKey::{}",
trie_key
),
)
})?
.parse()
.map_err(|_| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("raw key does not have a valid AccountId to be TrieKey::{}", trie_key),
)
})
}
fn next_token(data: &[u8], separator: u8) -> Option<&[u8]> {
data.iter().position(|&byte| byte == separator).map(|idx| &data[..idx])
}
pub fn parse_account_id_from_contract_data_key(
raw_key: &[u8],
) -> Result<AccountId, std::io::Error> {
let account_id_prefix = parse_account_id_prefix(col::CONTRACT_DATA, raw_key)?;
if let Some(account_id) = next_token(account_id_prefix, ACCOUNT_DATA_SEPARATOR) {
parse_account_id_from_slice(account_id, "ContractData")
} else {
Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"raw key does not have ACCOUNT_DATA_SEPARATOR to be TrieKey::ContractData",
))
}
}
pub fn parse_account_id_from_account_key(raw_key: &[u8]) -> Result<AccountId, std::io::Error> {
let account_id = parse_account_id_prefix(col::ACCOUNT, raw_key)?;
parse_account_id_from_slice(account_id, "Account")
}
pub fn parse_account_id_from_access_key_key(
raw_key: &[u8],
) -> Result<AccountId, std::io::Error> {
let account_id_prefix = parse_account_id_prefix(col::ACCESS_KEY, raw_key)?;
if let Some(account_id) = next_token(account_id_prefix, col::ACCESS_KEY) {
parse_account_id_from_slice(account_id, "AccessKey")
} else {
Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"raw key does not have public key to be TrieKey::AccessKey",
))
}
}
pub fn parse_account_id_from_contract_code_key(
raw_key: &[u8],
) -> Result<AccountId, std::io::Error> {
let account_id = parse_account_id_prefix(col::CONTRACT_CODE, raw_key)?;
parse_account_id_from_slice(account_id, "ContractCode")
}
pub fn parse_trie_key_access_key_from_raw_key(
raw_key: &[u8],
) -> Result<TrieKey, std::io::Error> {
let account_id = parse_account_id_from_access_key_key(raw_key)?;
let public_key = parse_public_key_from_access_key_key(raw_key, &account_id)?;
Ok(TrieKey::AccessKey { account_id, public_key })
}
#[allow(unused)]
pub fn parse_account_id_from_raw_key(
raw_key: &[u8],
) -> Result<Option<AccountId>, std::io::Error> {
for (col, col_name) in col::NON_DELAYED_RECEIPT_COLUMNS {
if parse_account_id_prefix(col, raw_key).is_err() {
continue;
}
let account_id = match col {
col::ACCOUNT => parse_account_id_from_account_key(raw_key)?,
col::CONTRACT_CODE => parse_account_id_from_contract_code_key(raw_key)?,
col::ACCESS_KEY => parse_account_id_from_access_key_key(raw_key)?,
_ => parse_account_id_from_trie_key_with_separator(col, raw_key, col_name)?,
};
return Ok(Some(account_id));
}
Ok(None)
}
fn parse_account_id_from_trie_key_with_separator(
col: u8,
raw_key: &[u8],
col_name: &str,
) -> Result<AccountId, std::io::Error> {
let account_id_prefix = parse_account_id_prefix(col, raw_key)?;
if let Some(account_id) = next_token(account_id_prefix, ACCOUNT_DATA_SEPARATOR) {
parse_account_id_from_slice(account_id, col_name)
} else {
Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("raw key does not have ACCOUNT_DATA_SEPARATOR to be TrieKey::{}", col_name),
))
}
}
pub fn parse_account_id_from_received_data_key(
raw_key: &[u8],
) -> Result<AccountId, std::io::Error> {
parse_account_id_from_trie_key_with_separator(col::RECEIVED_DATA, raw_key, "ReceivedData")
}
pub fn parse_data_id_from_received_data_key(
raw_key: &[u8],
account_id: &AccountId,
) -> Result<CryptoHash, std::io::Error> {
let prefix_len = col::ACCESS_KEY.len() * 2 + account_id.len();
if raw_key.len() < prefix_len {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"raw key is too short for TrieKey::ReceivedData",
));
}
CryptoHash::try_from(&raw_key[prefix_len..]).map_err(|_| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
"Can't parse CryptoHash for TrieKey::ReceivedData",
)
})
}
pub fn get_raw_prefix_for_access_keys(account_id: &AccountId) -> Vec<u8> {
let mut res = Vec::with_capacity(col::ACCESS_KEY.len() * 2 + account_id.len());
res.push(col::ACCESS_KEY);
res.extend(account_id.as_bytes());
res.push(col::ACCESS_KEY);
res
}
pub fn get_raw_prefix_for_contract_data(account_id: &AccountId, prefix: &[u8]) -> Vec<u8> {
let mut res = Vec::with_capacity(
col::CONTRACT_DATA.len()
+ account_id.len()
+ ACCOUNT_DATA_SEPARATOR.len()
+ prefix.len(),
);
res.push(col::CONTRACT_DATA);
res.extend(account_id.as_bytes());
res.push(ACCOUNT_DATA_SEPARATOR);
res.extend(prefix);
res
}
}
#[cfg(test)]
mod tests {
use near_crypto::KeyType;
use super::*;
const OK_ACCOUNT_IDS: &[&str] = &[
"aa",
"a-a",
"a-aa",
"100",
"0o",
"com",
"near",
"bowen",
"b-o_w_e-n",
"b.owen",
"bro.wen",
"a.ha",
"a.b-a.ra",
"system",
"over.9000",
"google.com",
"illia.cheapaccounts.near",
"0o0ooo00oo00o",
"alex-skidanov",
"10-4.8-2",
"b-o_w_e-n",
"no_lols",
"0123456789012345678901234567890123456789012345678901234567890123",
"near.a",
];
#[test]
fn test_key_for_account_consistency() {
for account_id in OK_ACCOUNT_IDS.iter().map(|x| x.parse::<AccountId>().unwrap()) {
let key = TrieKey::Account { account_id: account_id.clone() };
let raw_key = key.to_vec();
assert_eq!(raw_key.len(), key.len());
assert_eq!(
trie_key_parsers::parse_account_id_from_account_key(&raw_key).unwrap(),
account_id
);
assert_eq!(
trie_key_parsers::parse_account_id_from_raw_key(&raw_key).unwrap().unwrap(),
account_id
);
}
}
#[test]
fn test_key_for_access_key_consistency() {
let public_key = PublicKey::empty(KeyType::ED25519);
for account_id in OK_ACCOUNT_IDS.iter().map(|x| x.parse::<AccountId>().unwrap()) {
let key = TrieKey::AccessKey {
account_id: account_id.clone(),
public_key: public_key.clone(),
};
let raw_key = key.to_vec();
assert_eq!(raw_key.len(), key.len());
assert_eq!(
trie_key_parsers::parse_trie_key_access_key_from_raw_key(&raw_key).unwrap(),
key
);
assert_eq!(
trie_key_parsers::parse_account_id_from_access_key_key(&raw_key).unwrap(),
account_id
);
assert_eq!(
trie_key_parsers::parse_public_key_from_access_key_key(&raw_key, &account_id)
.unwrap(),
public_key
);
assert_eq!(
trie_key_parsers::parse_account_id_from_raw_key(&raw_key).unwrap().unwrap(),
account_id
);
}
}
#[test]
fn test_key_for_data_consistency() {
let data_key = b"0123456789" as &[u8];
for account_id in OK_ACCOUNT_IDS.iter().map(|x| x.parse::<AccountId>().unwrap()) {
let key =
TrieKey::ContractData { account_id: account_id.clone(), key: data_key.to_vec() };
let raw_key = key.to_vec();
assert_eq!(raw_key.len(), key.len());
assert_eq!(
trie_key_parsers::parse_account_id_from_contract_data_key(&raw_key).unwrap(),
account_id
);
assert_eq!(
trie_key_parsers::parse_data_key_from_contract_data_key(&raw_key, &account_id)
.unwrap(),
data_key
);
assert_eq!(
trie_key_parsers::parse_account_id_from_raw_key(&raw_key).unwrap().unwrap(),
account_id
);
}
}
#[test]
fn test_key_for_code_consistency() {
for account_id in OK_ACCOUNT_IDS.iter().map(|x| x.parse::<AccountId>().unwrap()) {
let key = TrieKey::ContractCode { account_id: account_id.clone() };
let raw_key = key.to_vec();
assert_eq!(raw_key.len(), key.len());
assert_eq!(
trie_key_parsers::parse_account_id_from_contract_code_key(&raw_key).unwrap(),
account_id
);
assert_eq!(
trie_key_parsers::parse_account_id_from_raw_key(&raw_key).unwrap().unwrap(),
account_id
);
}
}
#[test]
fn test_key_for_received_data_consistency() {
for account_id in OK_ACCOUNT_IDS.iter().map(|x| x.parse::<AccountId>().unwrap()) {
let key = TrieKey::ReceivedData {
receiver_id: account_id.clone(),
data_id: CryptoHash::default(),
};
let raw_key = key.to_vec();
assert_eq!(raw_key.len(), key.len());
assert_eq!(
trie_key_parsers::parse_account_id_from_received_data_key(&raw_key).unwrap(),
account_id
);
assert_eq!(
trie_key_parsers::parse_account_id_from_raw_key(&raw_key).unwrap().unwrap(),
account_id
);
assert_eq!(
trie_key_parsers::parse_data_id_from_received_data_key(&raw_key, &account_id)
.unwrap(),
CryptoHash::default(),
);
}
}
#[test]
fn test_key_for_postponed_receipt_consistency() {
for account_id in OK_ACCOUNT_IDS.iter().map(|x| x.parse::<AccountId>().unwrap()) {
let key = TrieKey::PostponedReceipt {
receiver_id: account_id.clone(),
receipt_id: CryptoHash::default(),
};
let raw_key = key.to_vec();
assert_eq!(raw_key.len(), key.len());
assert_eq!(
trie_key_parsers::parse_account_id_from_raw_key(&raw_key).unwrap().unwrap(),
account_id
);
}
}
#[test]
fn test_key_for_postponed_receipt_id_consistency() {
for account_id in OK_ACCOUNT_IDS.iter().map(|x| x.parse::<AccountId>().unwrap()) {
let key = TrieKey::PostponedReceiptId {
receiver_id: account_id.clone(),
data_id: CryptoHash::default(),
};
let raw_key = key.to_vec();
assert_eq!(raw_key.len(), key.len());
assert_eq!(
trie_key_parsers::parse_account_id_from_raw_key(&raw_key).unwrap().unwrap(),
account_id
);
}
}
#[test]
fn test_key_for_pending_data_count_consistency() {
for account_id in OK_ACCOUNT_IDS.iter().map(|x| x.parse::<AccountId>().unwrap()) {
let key = TrieKey::PendingDataCount {
receiver_id: account_id.clone(),
receipt_id: CryptoHash::default(),
};
let raw_key = key.to_vec();
assert_eq!(raw_key.len(), key.len());
assert_eq!(
trie_key_parsers::parse_account_id_from_raw_key(&raw_key).unwrap().unwrap(),
account_id
);
}
}
#[test]
fn test_key_for_delayed_receipts_consistency() {
let key = TrieKey::DelayedReceiptIndices;
let raw_key = key.to_vec();
assert!(trie_key_parsers::parse_account_id_from_raw_key(&raw_key).unwrap().is_none());
let key = TrieKey::DelayedReceipt { index: 0 };
let raw_key = key.to_vec();
assert!(trie_key_parsers::parse_account_id_from_raw_key(&raw_key).unwrap().is_none());
}
}