Skip to content

Commit ff7ef20

Browse files
OBorceOBorce
committed
Switch to BigUint for address balances to avoid overflow
1 parent 7fd42d4 commit ff7ef20

21 files changed

Lines changed: 252 additions & 186 deletions

File tree

‎Cargo.lock‎

Lines changed: 6 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,7 @@ dyn-clone = "1.0"
168168
# Note: enum iteration is also provided by strum. However, with strum it's a bit more limited,
169169
# e.g. there is no way to to go from one enum value to the next one, like enum_iterator::Sequence::next()
170170
# does.
171+
num-bigint = "0.4.6"
171172
enum-iterator = "2.0"
172173
enumflags2 = "0.7"
173174
expect-test = "1.3"

‎api-server/api-server-common/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ bb8-postgres = "0.8"
2222
clap = { workspace = true, features = ["derive"] }
2323
futures = { workspace = true, default-features = false }
2424
itertools.workspace = true
25+
num-bigint.workspace = true
2526
parity-scale-codec.workspace = true
2627
thiserror.workspace = true
2728
tokio = { workspace = true, features = ["full"] }

‎api-server/api-server-common/src/storage/impls/in_memory/mod.rs‎

Lines changed: 19 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,6 @@ use std::{
2222
sync::Arc,
2323
};
2424

25-
use itertools::Itertools as _;
26-
2725
use common::{
2826
address::Address,
2927
chain::{
@@ -44,6 +42,9 @@ use crate::storage::storage_api::{
4442

4543
use super::CURRENT_STORAGE_VERSION;
4644

45+
use itertools::Itertools as _;
46+
use num_bigint::BigUint;
47+
4748
#[derive(Debug, Clone, PartialEq, Eq)]
4849
struct TokenTransactionOrderedByTxId(TokenTransaction);
4950

@@ -63,8 +64,9 @@ impl Ord for TokenTransactionOrderedByTxId {
6364
struct ApiServerInMemoryStorage {
6465
block_table: BTreeMap<Id<Block>, BlockWithExtraData>,
6566
block_aux_data_table: BTreeMap<Id<Block>, BlockAuxData>,
66-
address_balance_table: BTreeMap<String, BTreeMap<CoinOrTokenId, BTreeMap<BlockHeight, Amount>>>,
67-
address_locked_balance_table: BTreeMap<String, BTreeMap<(CoinOrTokenId, BlockHeight), Amount>>,
67+
address_balance_table:
68+
BTreeMap<String, BTreeMap<CoinOrTokenId, BTreeMap<BlockHeight, BigUint>>>,
69+
address_locked_balance_table: BTreeMap<String, BTreeMap<(CoinOrTokenId, BlockHeight), BigUint>>,
6870
address_transactions_table: BTreeMap<String, BTreeMap<BlockHeight, Vec<Id<Transaction>>>>,
6971
token_transactions_table:
7072
BTreeMap<TokenId, BTreeMap<BlockHeight, BTreeSet<TokenTransactionOrderedByTxId>>>,
@@ -123,13 +125,13 @@ impl ApiServerInMemoryStorage {
123125
&self,
124126
address: &str,
125127
coin_or_token_id: CoinOrTokenId,
126-
) -> Result<Option<Amount>, ApiServerStorageError> {
128+
) -> Result<Option<BigUint>, ApiServerStorageError> {
127129
self.address_balance_table
128130
.get(address)
129131
.and_then(|by_coin_or_token| by_coin_or_token.get(&coin_or_token_id))
130132
.map_or_else(
131133
|| Ok(None),
132-
|by_height| Ok(by_height.values().last().copied()),
134+
|by_height| Ok(by_height.values().last().cloned()),
133135
)
134136
}
135137

@@ -155,7 +157,7 @@ impl ApiServerInMemoryStorage {
155157
(
156158
*coin_or_token_id,
157159
AmountWithDecimals {
158-
amount: *by_height.values().last().expect("not empty"),
160+
amount: by_height.values().last().expect("not empty").clone(),
159161
decimals: number_of_decimals,
160162
},
161163
)
@@ -170,14 +172,14 @@ impl ApiServerInMemoryStorage {
170172
&self,
171173
address: &str,
172174
coin_or_token_id: CoinOrTokenId,
173-
) -> Result<Option<Amount>, ApiServerStorageError> {
175+
) -> Result<Option<BigUint>, ApiServerStorageError> {
174176
self.address_locked_balance_table.get(address).map_or_else(
175177
|| Ok(None),
176178
|balance| {
177179
let range_begin = (coin_or_token_id, BlockHeight::zero());
178180
let range_end = (coin_or_token_id, BlockHeight::max());
179181
let range = balance.range(range_begin..=range_end);
180-
Ok(range.last().map(|(_, v)| *v))
182+
Ok(range.last().map(|(_, v)| v.clone()))
181183
},
182184
)
183185
}
@@ -927,7 +929,7 @@ impl ApiServerInMemoryStorage {
927929
fn set_address_balance_at_height(
928930
&mut self,
929931
address: &Address<Destination>,
930-
amount: Amount,
932+
amount: BigUint,
931933
coin_or_token_id: CoinOrTokenId,
932934
block_height: BlockHeight,
933935
) -> Result<(), ApiServerStorageError> {
@@ -937,8 +939,8 @@ impl ApiServerInMemoryStorage {
937939
.entry(coin_or_token_id)
938940
.or_default()
939941
.entry(block_height)
940-
.and_modify(|e| *e = amount)
941-
.or_insert(amount);
942+
.and_modify(|e| *e = amount.clone())
943+
.or_insert(amount.clone());
942944

943945
self.update_nft_owner(coin_or_token_id, amount, address, block_height);
944946

@@ -950,7 +952,7 @@ impl ApiServerInMemoryStorage {
950952
fn update_nft_owner(
951953
&mut self,
952954
coin_or_token_id: CoinOrTokenId,
953-
amount: Amount,
955+
amount: BigUint,
954956
address: &Address<Destination>,
955957
block_height: BlockHeight,
956958
) {
@@ -960,7 +962,7 @@ impl ApiServerInMemoryStorage {
960962

961963
if let Some(by_height) = self.nft_token_issuances.get_mut(&token_id) {
962964
let last = by_height.values().last().expect("not empty");
963-
let owner = (amount > Amount::ZERO).then_some(address.as_object().clone());
965+
let owner = (amount > BigUint::ZERO).then_some(address.as_object().clone());
964966
let new = NftWithOwner {
965967
nft: last.nft.clone(),
966968
owner,
@@ -972,16 +974,16 @@ impl ApiServerInMemoryStorage {
972974
fn set_address_locked_balance_at_height(
973975
&mut self,
974976
address: &Address<Destination>,
975-
amount: Amount,
977+
amount: BigUint,
976978
coin_or_token_id: CoinOrTokenId,
977979
block_height: BlockHeight,
978980
) -> Result<(), ApiServerStorageError> {
979981
self.address_locked_balance_table
980982
.entry(address.to_string())
981983
.or_default()
982984
.entry((coin_or_token_id, block_height))
983-
.and_modify(|e| *e = amount)
984-
.or_insert(amount);
985+
.and_modify(|e| *e = amount.clone())
986+
.or_insert(amount.clone());
985987

986988
self.update_nft_owner(coin_or_token_id, amount, address, block_height);
987989

‎api-server/api-server-common/src/storage/impls/in_memory/transactional/read.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@ use crate::storage::storage_api::{
3232

3333
use super::ApiServerInMemoryStorageTransactionalRo;
3434

35+
use num_bigint::BigUint;
36+
3537
#[async_trait::async_trait]
3638
impl ApiServerStorageRead for ApiServerInMemoryStorageTransactionalRo<'_> {
3739
async fn is_initialized(&self) -> Result<bool, ApiServerStorageError> {
@@ -42,7 +44,7 @@ impl ApiServerStorageRead for ApiServerInMemoryStorageTransactionalRo<'_> {
4244
&self,
4345
address: &str,
4446
coin_or_token_id: CoinOrTokenId,
45-
) -> Result<Option<Amount>, ApiServerStorageError> {
47+
) -> Result<Option<BigUint>, ApiServerStorageError> {
4648
self.transaction.get_address_balance(address, coin_or_token_id)
4749
}
4850

@@ -57,7 +59,7 @@ impl ApiServerStorageRead for ApiServerInMemoryStorageTransactionalRo<'_> {
5759
&self,
5860
address: &str,
5961
coin_or_token_id: CoinOrTokenId,
60-
) -> Result<Option<Amount>, ApiServerStorageError> {
62+
) -> Result<Option<BigUint>, ApiServerStorageError> {
6163
self.transaction.get_address_locked_balance(address, coin_or_token_id)
6264
}
6365

‎api-server/api-server-common/src/storage/impls/in_memory/transactional/write.rs‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ use common::{
3434

3535
use super::ApiServerInMemoryStorageTransactionalRw;
3636

37+
use num_bigint::BigUint;
38+
3739
#[async_trait::async_trait]
3840
impl ApiServerStorageWrite for ApiServerInMemoryStorageTransactionalRw<'_> {
3941
async fn reinitialize_storage(
@@ -74,7 +76,7 @@ impl ApiServerStorageWrite for ApiServerInMemoryStorageTransactionalRw<'_> {
7476
async fn set_address_balance_at_height(
7577
&mut self,
7678
address: &Address<Destination>,
77-
amount: Amount,
79+
amount: BigUint,
7880
coin_or_token_id: CoinOrTokenId,
7981
block_height: BlockHeight,
8082
) -> Result<(), ApiServerStorageError> {
@@ -89,7 +91,7 @@ impl ApiServerStorageWrite for ApiServerInMemoryStorageTransactionalRw<'_> {
8991
async fn set_address_locked_balance_at_height(
9092
&mut self,
9193
address: &Address<Destination>,
92-
amount: Amount,
94+
amount: BigUint,
9395
coin_or_token_id: CoinOrTokenId,
9496
block_height: BlockHeight,
9597
) -> Result<(), ApiServerStorageError> {
@@ -327,7 +329,7 @@ impl ApiServerStorageRead for ApiServerInMemoryStorageTransactionalRw<'_> {
327329
&self,
328330
address: &str,
329331
coin_or_token_id: CoinOrTokenId,
330-
) -> Result<Option<Amount>, ApiServerStorageError> {
332+
) -> Result<Option<BigUint>, ApiServerStorageError> {
331333
self.transaction.get_address_balance(address, coin_or_token_id)
332334
}
333335

@@ -342,7 +344,7 @@ impl ApiServerStorageRead for ApiServerInMemoryStorageTransactionalRw<'_> {
342344
&self,
343345
address: &str,
344346
coin_or_token_id: CoinOrTokenId,
345-
) -> Result<Option<Amount>, ApiServerStorageError> {
347+
) -> Result<Option<BigUint>, ApiServerStorageError> {
346348
self.transaction.get_address_locked_balance(address, coin_or_token_id)
347349
}
348350

‎api-server/api-server-common/src/storage/impls/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
// See the License for the specific language governing permissions and
1414
// limitations under the License.
1515

16-
pub const CURRENT_STORAGE_VERSION: u32 = 24;
16+
pub const CURRENT_STORAGE_VERSION: u32 = 25;
1717

1818
pub mod in_memory;
1919
pub mod postgres;

0 commit comments

Comments
 (0)