@@ -29,7 +29,7 @@ use aurora_engine_sdk::{
29
29
io:: { StorageIntermediate , IO } ,
30
30
promise:: PromiseHandler ,
31
31
} ;
32
- use aurora_engine_types:: parameters:: engine:: FullAccessKeyArgs ;
32
+ use aurora_engine_types:: parameters:: engine:: { FullAccessKeyArgs , UpgradeParams } ;
33
33
use aurora_engine_types:: types:: { NearGas , ZERO_YOCTO } ;
34
34
use aurora_engine_types:: {
35
35
borsh:: BorshDeserialize ,
@@ -38,7 +38,7 @@ use aurora_engine_types::{
38
38
NewCallArgs , PausePrecompilesCallArgs , RelayerKeyArgs , RelayerKeyManagerArgs ,
39
39
SetOwnerArgs , SetUpgradeDelayBlocksArgs , StartHashchainArgs ,
40
40
} ,
41
- promise:: { PromiseAction , PromiseBatchAction , PromiseCreateArgs } ,
41
+ promise:: { PromiseAction , PromiseBatchAction } ,
42
42
} ,
43
43
storage:: { self , KeyPrefix } ,
44
44
types:: { Address , Yocto } ,
@@ -48,7 +48,7 @@ use function_name::named;
48
48
49
49
const CODE_KEY : & [ u8 ; 4 ] = b"CODE" ;
50
50
const CODE_STAGE_KEY : & [ u8 ; 10 ] = b"CODE_STAGE" ;
51
- const GAS_FOR_STATE_MIGRATION : NearGas = NearGas :: new ( 100_000_000_000_000 ) ;
51
+ const GAS_FOR_STATE_MIGRATION : NearGas = NearGas :: new ( 50_000_000_000_000 ) ;
52
52
53
53
#[ named]
54
54
pub fn new < I : IO + Copy , E : Env > ( mut io : I , env : & E ) -> Result < ( ) , ContractError > {
@@ -189,23 +189,30 @@ pub fn upgrade<I: IO + Copy, E: Env, H: PromiseHandler>(
189
189
require_running ( & state) ?;
190
190
require_owner_only ( & state, & env. predecessor_account_id ( ) ) ?;
191
191
192
- let code = io. read_input ( ) . to_vec ( ) ;
193
- let current_account_id = env. current_account_id ( ) ;
194
- let batch = PromiseBatchAction {
195
- target_account_id : current_account_id. clone ( ) ,
196
- actions : vec ! [ PromiseAction :: DeployContract { code } ] ,
197
- } ;
198
- let state_migration_callback = PromiseCreateArgs {
199
- target_account_id : current_account_id,
200
- method : "state_migration" . to_string ( ) ,
201
- args : vec ! [ ] ,
202
- attached_balance : ZERO_YOCTO ,
203
- attached_gas : GAS_FOR_STATE_MIGRATION ,
192
+ let input = io. read_input ( ) . to_vec ( ) ;
193
+ let ( code, state_migration_gas) = match UpgradeParams :: try_from_slice ( & input) {
194
+ Ok ( args) => (
195
+ args. code ,
196
+ args. state_migration_gas
197
+ . map_or ( GAS_FOR_STATE_MIGRATION , NearGas :: new) ,
198
+ ) ,
199
+ Err ( _) => ( input, GAS_FOR_STATE_MIGRATION ) , // Backward compatibility
204
200
} ;
205
- let promise_id = unsafe {
206
- let base_id = handler. promise_create_batch ( & batch) ;
207
- handler. promise_attach_callback ( base_id, & state_migration_callback)
201
+
202
+ let target_account_id = env. current_account_id ( ) ;
203
+ let batch = PromiseBatchAction {
204
+ target_account_id,
205
+ actions : vec ! [
206
+ PromiseAction :: DeployContract { code } ,
207
+ PromiseAction :: FunctionCall {
208
+ name: "state_migration" . to_string( ) ,
209
+ args: vec![ ] ,
210
+ attached_yocto: ZERO_YOCTO ,
211
+ gas: state_migration_gas,
212
+ } ,
213
+ ] ,
208
214
} ;
215
+ let promise_id = unsafe { handler. promise_create_batch ( & batch) } ;
209
216
210
217
handler. promise_return ( promise_id) ;
211
218
0 commit comments