Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit bf067e3

Browse files
minaminaominaminao
committed
remove redundancy
1 parent 896528a commit bf067e3

File tree

9 files changed

+86
-460
lines changed

9 files changed

+86
-460
lines changed

‎src/OakSecurityCosmWasmCTF/01-Mjolnir/src/exploit.rs‎

Lines changed: 4 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -4,78 +4,11 @@
44
pub mod exploit {
55
use crate::{
66
contract::{DENOM, LOCK_PERIOD, MINIMUM_DEPOSIT_AMOUNT},
7-
msg::{ExecuteMsg, InstantiateMsg},
7+
integration_tests::tests::{proper_instantiate, USER},
8+
msg::ExecuteMsg,
89
};
9-
use cosmwasm_std::{coin, Addr, Empty, Uint128};
10-
use cw_multi_test::{App, Contract, ContractWrapper, Executor};
11-
12-
pub fn challenge_contract() -> Box<dyn Contract<Empty>> {
13-
let contract = ContractWrapper::new(
14-
crate::contract::execute,
15-
crate::contract::instantiate,
16-
crate::contract::query,
17-
);
18-
Box::new(contract)
19-
}
20-
21-
pub const USER: &str = "user";
22-
pub const ADMIN: &str = "admin";
23-
24-
pub fn proper_instantiate() -> (App, Addr) {
25-
let mut app = App::default();
26-
let cw_template_id = app.store_code(challenge_contract());
27-
28-
// init contract
29-
let msg = InstantiateMsg {};
30-
let contract_addr = app
31-
.instantiate_contract(
32-
cw_template_id,
33-
Addr::unchecked(ADMIN),
34-
&msg,
35-
&[],
36-
"test",
37-
None,
38-
)
39-
.unwrap();
40-
41-
// mint funds to contract
42-
app = mint_tokens(
43-
app,
44-
contract_addr.to_string(),
45-
MINIMUM_DEPOSIT_AMOUNT * Uint128::new(10),
46-
);
47-
48-
// mint funds to user
49-
app = mint_tokens(app, USER.to_string(), MINIMUM_DEPOSIT_AMOUNT);
50-
51-
// deposit
52-
let msg = ExecuteMsg::Deposit {};
53-
let sender = Addr::unchecked(USER);
54-
app.execute_contract(
55-
sender.clone(),
56-
contract_addr.clone(),
57-
&msg,
58-
&[coin(MINIMUM_DEPOSIT_AMOUNT.u128(), DENOM)],
59-
)
60-
.unwrap();
61-
62-
// verify no funds
63-
let balance = app.wrap().query_balance(USER, DENOM).unwrap().amount;
64-
assert_eq!(balance, Uint128::zero());
65-
66-
(app, contract_addr)
67-
}
68-
69-
pub fn mint_tokens(mut app: App, recipient: String, amount: Uint128) -> App {
70-
app.sudo(cw_multi_test::SudoMsg::Bank(
71-
cw_multi_test::BankSudo::Mint {
72-
to_address: recipient.to_owned(),
73-
amount: vec![coin(amount.u128(), DENOM)],
74-
},
75-
))
76-
.unwrap();
77-
app
78-
}
10+
use cosmwasm_std::{Addr, Uint128};
11+
use cw_multi_test::Executor;
7912

8013
#[test]
8114
fn exploit() {

‎src/OakSecurityCosmWasmCTF/02-Gungnir/src/exploit.rs‎

Lines changed: 8 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -4,60 +4,15 @@
44
pub mod exploit {
55
use crate::{
66
contract::{DENOM, LOCK_PERIOD},
7-
msg::{ExecuteMsg, InstantiateMsg, QueryMsg},
7+
integration_tests::tests::{base_scenario, INITIAL_USER_AMOUNT, USER},
8+
msg::{ExecuteMsg, QueryMsg},
89
};
9-
use cosmwasm_std::{coin, Addr, Empty, Uint128};
10-
use cw_multi_test::{App, Contract, ContractWrapper, Executor};
11-
12-
pub fn challenge_contract() -> Box<dyn Contract<Empty>> {
13-
let contract = ContractWrapper::new(
14-
crate::contract::execute,
15-
crate::contract::instantiate,
16-
crate::contract::query,
17-
);
18-
Box::new(contract)
19-
}
20-
21-
pub const USER: &str = "user";
22-
pub const ADMIN: &str = "admin";
23-
const INITIAL_AMOUNT: Uint128 = Uint128::new(1_000);
24-
25-
pub fn proper_instantiate() -> (App, Addr) {
26-
let mut app = App::default();
27-
let cw_template_id = app.store_code(challenge_contract());
28-
29-
// init contract
30-
let msg = InstantiateMsg {};
31-
let contract_addr = app
32-
.instantiate_contract(
33-
cw_template_id,
34-
Addr::unchecked(ADMIN),
35-
&msg,
36-
&[],
37-
"test",
38-
None,
39-
)
40-
.unwrap();
41-
42-
app = mint_tokens(app, USER.to_string(), INITIAL_AMOUNT);
43-
44-
(app, contract_addr)
45-
}
46-
47-
pub fn mint_tokens(mut app: App, recipient: String, amount: Uint128) -> App {
48-
app.sudo(cw_multi_test::SudoMsg::Bank(
49-
cw_multi_test::BankSudo::Mint {
50-
to_address: recipient,
51-
amount: vec![coin(amount.u128(), DENOM)],
52-
},
53-
))
54-
.unwrap();
55-
app
56-
}
10+
use cosmwasm_std::{coin, Addr};
11+
use cw_multi_test::Executor;
5712

5813
#[test]
5914
fn exploit() {
60-
let (mut app, contract_addr) = proper_instantiate();
15+
let (mut app, contract_addr) = base_scenario();
6116

6217
let player = Addr::unchecked(USER);
6318

@@ -67,12 +22,12 @@ pub mod exploit {
6722
player.clone(),
6823
contract_addr.clone(),
6924
&msg,
70-
&[coin(INITIAL_AMOUNT.u128(), DENOM)],
25+
&[coin(INITIAL_USER_AMOUNT.u128(), DENOM)],
7126
)
7227
.unwrap();
7328

7429
let msg = ExecuteMsg::Stake {
75-
lock_amount: INITIAL_AMOUNT.u128(),
30+
lock_amount: INITIAL_USER_AMOUNT.u128(),
7631
};
7732
app.execute_contract(player.clone(), contract_addr.clone(), &msg, &[])
7833
.unwrap();
@@ -84,7 +39,7 @@ pub mod exploit {
8439

8540
// normal stake
8641
let msg = ExecuteMsg::Unstake {
87-
unlock_amount: INITIAL_AMOUNT.u128() + 1,
42+
unlock_amount: INITIAL_USER_AMOUNT.u128() + 1,
8843
};
8944
app.execute_contract(player.clone(), contract_addr.clone(), &msg, &[])
9045
.unwrap();

‎src/OakSecurityCosmWasmCTF/02-Gungnir/src/integration_tests.rs‎

Lines changed: 25 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -51,13 +51,20 @@ pub mod tests {
5151
app
5252
}
5353

54-
#[test]
55-
fn basic_flow() {
54+
pub const INITIAL_USER_AMOUNT: Uint128 = Uint128::new(1_000);
55+
56+
pub fn base_scenario() -> (App, Addr) {
5657
let (mut app, contract_addr) = proper_instantiate();
5758

58-
let amount = Uint128::new(1_000);
59+
app = mint_tokens(app, USER.to_string(), INITIAL_USER_AMOUNT);
60+
61+
(app, contract_addr)
62+
}
63+
64+
#[test]
65+
fn basic_flow() {
66+
let (mut app, contract_addr) = base_scenario();
5967

60-
app = mint_tokens(app, USER.to_string(), amount);
6168
let sender = Addr::unchecked(USER);
6269

6370
// deposit funds
@@ -66,7 +73,7 @@ pub mod tests {
6673
sender.clone(),
6774
contract_addr.clone(),
6875
&msg,
69-
&[coin(amount.u128(), DENOM)],
76+
&[coin(INITIAL_USER_AMOUNT.u128(), DENOM)],
7077
)
7178
.unwrap();
7279

@@ -82,18 +89,18 @@ pub mod tests {
8289
.wrap()
8390
.query_wasm_smart(contract_addr.clone(), &msg)
8491
.unwrap();
85-
assert_eq!(user.total_tokens, amount);
92+
assert_eq!(user.total_tokens, INITIAL_USER_AMOUNT);
8693

8794
// cannot stake more than deposited
8895
let msg = ExecuteMsg::Stake {
89-
lock_amount: amount.u128() + 1,
96+
lock_amount: INITIAL_USER_AMOUNT.u128() + 1,
9097
};
9198
app.execute_contract(sender.clone(), contract_addr.clone(), &msg, &[])
9299
.unwrap_err();
93100

94101
// normal stake
95102
let msg = ExecuteMsg::Stake {
96-
lock_amount: amount.u128(),
103+
lock_amount: INITIAL_USER_AMOUNT.u128(),
97104
};
98105
app.execute_contract(sender.clone(), contract_addr.clone(), &msg, &[])
99106
.unwrap();
@@ -106,17 +113,19 @@ pub mod tests {
106113
.wrap()
107114
.query_wasm_smart(contract_addr.clone(), &msg)
108115
.unwrap();
109-
assert_eq!(voting_power, amount.u128());
116+
assert_eq!(voting_power, INITIAL_USER_AMOUNT.u128());
110117

111118
// cannot unstake before maturity
112119
let msg = ExecuteMsg::Unstake {
113-
unlock_amount: amount.u128(),
120+
unlock_amount: INITIAL_USER_AMOUNT.u128(),
114121
};
115122
app.execute_contract(sender.clone(), contract_addr.clone(), &msg, &[])
116123
.unwrap_err();
117124

118125
// cannot withdraw while staked
119-
let msg = ExecuteMsg::Withdraw { amount };
126+
let msg = ExecuteMsg::Withdraw {
127+
amount: INITIAL_USER_AMOUNT,
128+
};
120129
app.execute_contract(sender.clone(), contract_addr.clone(), &msg, &[])
121130
.unwrap_err();
122131

@@ -127,7 +136,7 @@ pub mod tests {
127136

128137
// normal unstake
129138
let msg = ExecuteMsg::Unstake {
130-
unlock_amount: amount.u128(),
139+
unlock_amount: INITIAL_USER_AMOUNT.u128(),
131140
};
132141
app.execute_contract(sender.clone(), contract_addr.clone(), &msg, &[])
133142
.unwrap();
@@ -143,12 +152,14 @@ pub mod tests {
143152
assert_eq!(voting_power, 0_u128);
144153

145154
// normal withdraw
146-
let msg = ExecuteMsg::Withdraw { amount };
155+
let msg = ExecuteMsg::Withdraw {
156+
amount: INITIAL_USER_AMOUNT,
157+
};
147158
app.execute_contract(sender, contract_addr, &msg, &[])
148159
.unwrap();
149160

150161
// funds are received
151162
let balance = app.wrap().query_balance(USER, DENOM).unwrap().amount;
152-
assert_eq!(balance, amount);
163+
assert_eq!(balance, INITIAL_USER_AMOUNT);
153164
}
154165
}

0 commit comments

Comments
 (0)