Skip to content

Commit 0ba22ad

Browse files
canercandancanercandan
authored
Merge pull request #3283 from Sifchain/fix/optimize-invariant-checks
fix: 🐛 optimize invariant checks
2 parents 4e7a4b1 + f9aae9e commit 0ba22ad

File tree

4 files changed

+76
-45
lines changed

4 files changed

+76
-45
lines changed

‎x/clp/keeper/invariants.go

+39
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,45 @@ func (k Keeper) BalanceModuleAccountCheck() sdk.Invariant {
4747
}
4848
}
4949

50+
func (k Keeper) SingleExternalBalanceModuleAccountCheck(externalAsset string) sdk.Invariant {
51+
return func(ctx sdk.Context) (string, bool) {
52+
pool, err := k.GetPool(ctx, externalAsset)
53+
if err == nil {
54+
clpModuleTotalExternalBalance := k.GetBankKeeper().GetBalance(ctx, types.GetCLPModuleAddress(), pool.ExternalAsset.Symbol)
55+
clpModuleTotalExternalBalanceUint := sdk.NewUintFromString(clpModuleTotalExternalBalance.Amount.String())
56+
57+
ok := pool.ExternalAssetBalance.Add(pool.ExternalCustody).Equal(clpModuleTotalExternalBalanceUint)
58+
if !ok {
59+
return fmt.Sprintf("external balance mismatch in pool %s (module: %s != pool: %s)",
60+
pool.ExternalAsset.Symbol,
61+
clpModuleTotalExternalBalanceUint.String(),
62+
pool.ExternalAssetBalance.String()), true
63+
}
64+
}
65+
66+
// Get Rowan Balance from CLP Module
67+
clpModuleTotalNativeBalance := k.GetBankKeeper().GetBalance(ctx, types.GetCLPModuleAddress(), types.GetSettlementAsset().Symbol)
68+
clpModuleTotalNativeBalanceUint := sdk.NewUintFromString(clpModuleTotalNativeBalance.Amount.String())
69+
70+
pools := k.GetPools(ctx)
71+
poolsTotalNativeBalanceUint := sdk.ZeroUint()
72+
poolsTotalNativeCustodyUint := sdk.ZeroUint()
73+
for _, pool := range pools {
74+
poolsTotalNativeBalanceUint = poolsTotalNativeBalanceUint.Add(pool.NativeAssetBalance)
75+
poolsTotalNativeCustodyUint = poolsTotalNativeCustodyUint.Add(pool.NativeCustody)
76+
}
77+
78+
ok := poolsTotalNativeBalanceUint.Add(poolsTotalNativeCustodyUint).Equal(clpModuleTotalNativeBalanceUint)
79+
if !ok {
80+
return fmt.Sprintf("native balance mismatch across all pools (module: %s != pools: %s)",
81+
clpModuleTotalNativeBalanceUint.String(),
82+
poolsTotalNativeBalanceUint.String()), true
83+
}
84+
85+
return "pool and module account balances match", false
86+
}
87+
}
88+
5089
func (k Keeper) UnitsCheck() sdk.Invariant {
5190
return func(ctx sdk.Context) (string, bool) {
5291
pools := k.GetPools(ctx)

‎x/clp/keeper/msg_server.go

+15-37
Original file line numberDiff line numberDiff line change
@@ -379,16 +379,11 @@ func (k msgServer) DecommissionPool(goCtx context.Context, msg *types.MsgDecommi
379379
),
380380
})
381381

382-
res, stop := k.BalanceModuleAccountCheck()(ctx)
382+
res, stop := k.SingleExternalBalanceModuleAccountCheck(msg.Symbol)(ctx)
383383
if stop {
384384
return nil, sdkerrors.Wrap(types.ErrBalanceModuleAccountCheck, res)
385385
}
386386

387-
res, stop = k.UnitsCheck()(ctx)
388-
if stop {
389-
return nil, sdkerrors.Wrap(types.ErrUnitsCheck, res)
390-
}
391-
392387
return &types.MsgDecommissionPoolResponse{}, nil
393388
}
394389

@@ -453,16 +448,11 @@ func (k msgServer) CreatePool(goCtx context.Context, msg *types.MsgCreatePool) (
453448
),
454449
})
455450

456-
res, stop := k.BalanceModuleAccountCheck()(ctx)
451+
res, stop := k.SingleExternalBalanceModuleAccountCheck(msg.ExternalAsset.Symbol)(ctx)
457452
if stop {
458453
return nil, sdkerrors.Wrap(types.ErrBalanceModuleAccountCheck, res)
459454
}
460455

461-
res, stop = k.UnitsCheck()(ctx)
462-
if stop {
463-
return nil, sdkerrors.Wrap(types.ErrUnitsCheck, res)
464-
}
465-
466456
return &types.MsgCreatePoolResponse{}, nil
467457
}
468458

@@ -678,14 +668,17 @@ func (k msgServer) Swap(goCtx context.Context, msg *types.MsgSwap) (*types.MsgSw
678668
}
679669
}
680670

681-
res, stop := k.BalanceModuleAccountCheck()(ctx)
682-
if stop {
683-
return nil, sdkerrors.Wrap(types.ErrBalanceModuleAccountCheck, res)
671+
if !msg.SentAsset.Equals(types.GetSettlementAsset()) {
672+
res, stop := k.SingleExternalBalanceModuleAccountCheck(msg.SentAsset.Symbol)(ctx)
673+
if stop {
674+
return nil, sdkerrors.Wrap(types.ErrBalanceModuleAccountCheck, res)
675+
}
684676
}
685-
686-
res, stop = k.UnitsCheck()(ctx)
687-
if stop {
688-
return nil, sdkerrors.Wrap(types.ErrUnitsCheck, res)
677+
if !msg.ReceivedAsset.Equals(types.GetSettlementAsset()) {
678+
res, stop := k.SingleExternalBalanceModuleAccountCheck(msg.ReceivedAsset.Symbol)(ctx)
679+
if stop {
680+
return nil, sdkerrors.Wrap(types.ErrBalanceModuleAccountCheck, res)
681+
}
689682
}
690683

691684
return &types.MsgSwapResponse{}, nil
@@ -753,16 +746,11 @@ func (k msgServer) AddLiquidity(goCtx context.Context, msg *types.MsgAddLiquidit
753746
k.ProcessRemovalQueue(ctx, msg, newPoolUnits)
754747
}
755748

756-
res, stop := k.BalanceModuleAccountCheck()(ctx)
749+
res, stop := k.SingleExternalBalanceModuleAccountCheck(msg.ExternalAsset.Symbol)(ctx)
757750
if stop {
758751
return nil, sdkerrors.Wrap(types.ErrBalanceModuleAccountCheck, res)
759752
}
760753

761-
res, stop = k.UnitsCheck()(ctx)
762-
if stop {
763-
return nil, sdkerrors.Wrap(types.ErrUnitsCheck, res)
764-
}
765-
766754
return &types.MsgAddLiquidityResponse{}, nil
767755
}
768756

@@ -867,16 +855,11 @@ func (k msgServer) RemoveLiquidityUnits(goCtx context.Context, msg *types.MsgRem
867855
),
868856
})
869857

870-
res, stop := k.BalanceModuleAccountCheck()(ctx)
858+
res, stop := k.SingleExternalBalanceModuleAccountCheck(msg.ExternalAsset.Symbol)(ctx)
871859
if stop {
872860
return nil, sdkerrors.Wrap(types.ErrBalanceModuleAccountCheck, res)
873861
}
874862

875-
res, stop = k.UnitsCheck()(ctx)
876-
if stop {
877-
return nil, sdkerrors.Wrap(types.ErrUnitsCheck, res)
878-
}
879-
880863
return &types.MsgRemoveLiquidityUnitsResponse{}, nil
881864
}
882865

@@ -1031,16 +1014,11 @@ func (k msgServer) RemoveLiquidity(goCtx context.Context, msg *types.MsgRemoveLi
10311014
),
10321015
})
10331016

1034-
res, stop := k.BalanceModuleAccountCheck()(ctx)
1017+
res, stop := k.SingleExternalBalanceModuleAccountCheck(msg.ExternalAsset.Symbol)(ctx)
10351018
if stop {
10361019
return nil, sdkerrors.Wrap(types.ErrBalanceModuleAccountCheck, res)
10371020
}
10381021

1039-
res, stop = k.UnitsCheck()(ctx)
1040-
if stop {
1041-
return nil, sdkerrors.Wrap(types.ErrUnitsCheck, res)
1042-
}
1043-
10441022
return &types.MsgRemoveLiquidityResponse{}, nil
10451023
}
10461024

‎x/margin/keeper/msg_server.go

+21-7
Original file line numberDiff line numberDiff line change
@@ -223,7 +223,7 @@ func (k msgServer) OpenLong(ctx sdk.Context, msg *types.MsgOpen) (*types.MTP, er
223223
return nil, types.ErrMTPUnhealthy
224224
}
225225

226-
res, stop := k.ClpKeeper().BalanceModuleAccountCheck()(ctx)
226+
res, stop := k.ClpKeeper().SingleExternalBalanceModuleAccountCheck(externalAsset)(ctx)
227227
if stop {
228228
return nil, sdkerrors.Wrap(clptypes.ErrBalanceModuleAccountCheck, res)
229229
}
@@ -285,9 +285,16 @@ func (k msgServer) CloseLong(ctx sdk.Context, msg *types.MsgClose) (*types.MTP,
285285
return nil, sdk.ZeroUint(), err
286286
}
287287

288-
res, stop := k.ClpKeeper().BalanceModuleAccountCheck()(ctx)
289-
if stop {
290-
return nil, sdk.ZeroUint(), sdkerrors.Wrap(clptypes.ErrBalanceModuleAccountCheck, res)
288+
if types.StringCompare(mtp.CollateralAsset, nativeAsset) {
289+
res, stop := k.ClpKeeper().SingleExternalBalanceModuleAccountCheck(mtp.CustodyAsset)(ctx)
290+
if stop {
291+
return nil, sdk.ZeroUint(), sdkerrors.Wrap(clptypes.ErrBalanceModuleAccountCheck, res)
292+
}
293+
} else {
294+
res, stop := k.ClpKeeper().SingleExternalBalanceModuleAccountCheck(mtp.CollateralAsset)(ctx)
295+
if stop {
296+
return nil, sdk.ZeroUint(), sdkerrors.Wrap(clptypes.ErrBalanceModuleAccountCheck, res)
297+
}
291298
}
292299

293300
return &mtp, repayAmount, nil
@@ -407,9 +414,16 @@ func (k msgServer) ForceClose(goCtx context.Context, msg *types.MsgForceClose) (
407414

408415
k.EmitAdminClose(ctx, &mtpToClose, repayAmount, msg.Signer)
409416

410-
res, stop := k.ClpKeeper().BalanceModuleAccountCheck()(ctx)
411-
if stop {
412-
return nil, sdkerrors.Wrap(clptypes.ErrBalanceModuleAccountCheck, res)
417+
if types.StringCompare(mtpToClose.CollateralAsset, types.GetSettlementAsset()) {
418+
res, stop := k.ClpKeeper().SingleExternalBalanceModuleAccountCheck(mtpToClose.CustodyAsset)(ctx)
419+
if stop {
420+
return nil, sdkerrors.Wrap(clptypes.ErrBalanceModuleAccountCheck, res)
421+
}
422+
} else {
423+
res, stop := k.ClpKeeper().SingleExternalBalanceModuleAccountCheck(mtpToClose.CollateralAsset)(ctx)
424+
if stop {
425+
return nil, sdkerrors.Wrap(clptypes.ErrBalanceModuleAccountCheck, res)
426+
}
413427
}
414428

415429
return &types.MsgForceCloseResponse{}, nil

‎x/margin/types/expected_keepers.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ type CLPKeeper interface {
2727

2828
GetRemovalQueue(ctx sdk.Context, symbol string) clptypes.RemovalQueue
2929

30-
BalanceModuleAccountCheck() sdk.Invariant
30+
SingleExternalBalanceModuleAccountCheck(externalAsset string) sdk.Invariant
3131
}
3232

3333
type Keeper interface {

0 commit comments

Comments
 (0)