def set_balance(self, address, balance):
validate_canonical_address(address, title="Storage Address")
validate_uint256(balance, title="Account Balance")
key = get_balance_key(address)
self._check_accessibility(key)
self._trie[key] = int_to_big_endian(balance)
def get_balance(self, address):
validate_canonical_address(address, title="Storage Address")
key = get_balance_key(address)
self._check_accessibility(key)
if key in self._trie:
return big_endian_to_int(self._trie[key])
else:
return 0
def test_access_restriction():
# populate db
state = ShardingAccountStateDB(MemoryDB())
state.set_balance(ADDRESS, 2)
state.set_code(ADDRESS, b"code")
state.set_storage(ADDRESS, 123, 4)
original_db = state.db
original_root_hash = state.root_hash
def make_state(access_list):
kv_store = copy.deepcopy(original_db.wrapped_db.kv_store)
db = MemoryDB(kv_store)
return ShardingAccountStateDB(db,
original_root_hash,
access_list=access_list)
# access lists to use
CODE_ACCESS_LIST = [get_code_key(ADDRESS)]
BALANCE_ACCESS_LIST = [get_balance_key(ADDRESS)]
STORAGE_ACCESS_LIST = [get_storage_key(ADDRESS, 123)]
# test with access list
state = make_state(BALANCE_ACCESS_LIST)
state.get_balance(ADDRESS)
state.set_balance(ADDRESS, 3)
state.delta_balance(ADDRESS, 1)
state = make_state(CODE_ACCESS_LIST)
state.get_code(ADDRESS)
state.set_code(ADDRESS, b"new_code")
state = make_state(STORAGE_ACCESS_LIST)
state.get_storage(ADDRESS, 123)
state.set_storage(ADDRESS, 123, 5)
# test without access list
state = make_state([])
with pytest.raises(UnannouncedStateAccess):
state.get_balance(ADDRESS)
with pytest.raises(UnannouncedStateAccess):
state.set_balance(ADDRESS, 3)
with pytest.raises(UnannouncedStateAccess):
state.delta_balance(ADDRESS, 1)
with pytest.raises(UnannouncedStateAccess):
state.get_code(ADDRESS)
with pytest.raises(UnannouncedStateAccess):
state.set_code(ADDRESS, b"new_code")
with pytest.raises(UnannouncedStateAccess):
state.get_storage(ADDRESS, 123)
with pytest.raises(UnannouncedStateAccess):
state.set_storage(ADDRESS, 123, 5)
def test_get_witness_nodes(populated_shard_chaindb_and_root_hash):
chaindb, root_hash = populated_shard_chaindb_and_root_hash
header = CollationHeader(
shard_id=1,
expected_period_number=0,
period_start_prevhash=ZERO_HASH32,
parent_hash=ZERO_HASH32,
number=0,
state_root=root_hash
)
prefixes = [
get_balance_key(A_ADDRESS),
get_balance_key(B_ADDRESS),
get_storage_key(A_ADDRESS, big_endian_to_int(b"key1")),
get_storage_key(B_ADDRESS, big_endian_to_int(b"key1")),
get_storage_key(B_ADDRESS, big_endian_to_int(b"key2")),
get_storage_key(B_ADDRESS, big_endian_to_int(b"key")),
get_storage_key(B_ADDRESS, big_endian_to_int(b"")),
]
witness_nodes = chaindb.get_witness_nodes(header, prefixes)
assert len(witness_nodes) == len(set(witness_nodes)) # no duplicates
assert sorted(witness_nodes) == sorted(witness_nodes) # sorted