@@ -114,7 +114,7 @@ func (p Programs) ActivateProgram(evm *vm.EVM, address common.Address, runCtx *c
114114 // already activated and up to date
115115 return 0 , codeHash , common.Hash {}, nil , false , ProgramUpToDateError ()
116116 }
117- wasm , err := getWasm (statedb , address , params . MaxWasmSize )
117+ wasm , err := getWasm (statedb , address , params )
118118 if err != nil {
119119 return 0 , codeHash , common.Hash {}, nil , false , err
120120 }
@@ -225,7 +225,7 @@ func (p Programs) CallProgram(
225225 statedb .AddStylusPages (program .footprint )
226226 defer statedb .SetStylusPagesOpen (open )
227227
228- asmMap , err := getCompiledProgram (statedb , moduleHash , contract .Address (), contract .Code , contract .CodeHash , params . MaxWasmSize , params . PageLimit , evm .Context .Time , debugMode , program , runCtx )
228+ asmMap , err := getCompiledProgram (statedb , moduleHash , contract .Address (), contract .Code , contract .CodeHash , params , evm .Context .Time , debugMode , program , runCtx )
229229 var ok bool
230230 var localAsm []byte
231231 if asmMap != nil {
@@ -307,30 +307,105 @@ func evmMemoryCost(size uint64) uint64 {
307307 return linearCost + squareCost
308308}
309309
310- func getWasm (statedb vm.StateDB , program common.Address , maxWasmSize uint32 ) ([]byte , error ) {
310+ func getWasm (statedb vm.StateDB , program common.Address , params * StylusParams ) ([]byte , error ) {
311311 prefixedWasm := statedb .GetCode (program )
312- return getWasmFromContractCode (prefixedWasm , maxWasmSize )
312+ return getWasmFromContractCode (statedb , prefixedWasm , params , true )
313313}
314314
315- func getWasmFromContractCode (prefixedWasm []byte , maxWasmSize uint32 ) ([]byte , error ) {
316- if prefixedWasm == nil {
315+ func getWasmFromContractCode (statedb vm. StateDB , prefixedWasm []byte , params * StylusParams , isActivation bool ) ([]byte , error ) {
316+ if len ( prefixedWasm ) == 0 {
317317 return nil , ProgramNotWasmError ()
318318 }
319- wasm , dictByte , err := state .StripStylusPrefix (prefixedWasm )
319+
320+ if state .IsStylusProgramClassic (prefixedWasm ) {
321+ return handleClassicStylus (prefixedWasm , params .MaxWasmSize )
322+ }
323+
324+ if params .arbosVersion >= gethParams .ArbosVersion_StylusContractLimit {
325+ if state .IsStylusProgramRoot (prefixedWasm ) {
326+ return handleRootStylus (statedb , prefixedWasm , params .MaxWasmSize , params .MaxFragmentCount , isActivation )
327+ }
328+
329+ if state .IsStylusProgramFragment (prefixedWasm ) {
330+ return nil , errors .New ("fragmented stylus programs cannot be activated directly; activate the root program instead" )
331+ }
332+ }
333+
334+ return nil , ProgramNotWasmError ()
335+ }
336+
337+ func handleClassicStylus (data []byte , maxSize uint32 ) ([]byte , error ) {
338+ wasm , dictByte , err := state .StripStylusPrefix (data )
339+ if err != nil {
340+ return nil , err
341+ }
342+
343+ dict , err := getStylusCompressionDict (dictByte )
344+ if err != nil {
345+ return nil , err
346+ }
347+
348+ return arbcompress .DecompressWithDictionary (wasm , int (maxSize ), dict )
349+ }
350+
351+ func handleRootStylus (statedb vm.StateDB , data []byte , maxSize uint32 , maxFragments uint16 , isActivation bool ) ([]byte , error ) {
352+ root , err := state .NewStylusRoot (data )
353+ if err != nil {
354+ return nil , err
355+ }
356+
357+ if isActivation {
358+ if root .DecompressedLength > maxSize {
359+ return nil , fmt .Errorf ("invalid wasm: decompressedLength %d is greater then MaxWasmSize %d" , root .DecompressedLength , maxSize )
360+ }
361+ if len (root .Addresses ) > int (maxFragments ) {
362+ return nil , fmt .Errorf ("invalid wasm: fragment count exceeds limit of %d" , maxFragments )
363+ }
364+ }
365+
366+ if len (root .Addresses ) == 0 {
367+ return nil , fmt .Errorf ("invalid wasm: fragment count cannot be zero" )
368+ }
369+
370+ var compressedWasm []byte
371+ for _ , addr := range root .Addresses {
372+ fragCode := statedb .GetCode (addr )
373+
374+ payload , err := state .StripStylusFragmentPrefix (fragCode )
375+ if err != nil {
376+ return nil , err
377+ }
378+
379+ compressedWasm = append (compressedWasm , payload ... )
380+ }
381+
382+ dict , err := getStylusCompressionDict (root .DictionaryType )
320383 if err != nil {
321384 return nil , err
322385 }
323386
324- var dict arbcompress.Dictionary
325- switch dictByte {
387+ wasm , err := arbcompress .DecompressWithDictionary (compressedWasm , int (root .DecompressedLength ), dict )
388+ if err != nil {
389+ return nil , err
390+ }
391+
392+ if len (wasm ) != int (root .DecompressedLength ) {
393+ return nil , fmt .Errorf ("invalid wasm: decompressed length %d does not match expected length %d" , len (wasm ), root .DecompressedLength )
394+ }
395+
396+ return wasm , nil
397+ }
398+
399+ // Named return parameters allow us to return the zero-value for 'dict' implicitly on error
400+ func getStylusCompressionDict (id byte ) (dict arbcompress.Dictionary , err error ) {
401+ switch id {
326402 case 0 :
327- dict = arbcompress .EmptyDictionary
403+ return arbcompress .EmptyDictionary , nil
328404 case 1 :
329- dict = arbcompress .StylusProgramDictionary
405+ return arbcompress .StylusProgramDictionary , nil
330406 default :
331- return nil , fmt .Errorf ("unsupported dictionary %v " , dictByte )
407+ return dict , fmt .Errorf ("unsupported dictionary type: %d " , id )
332408 }
333- return arbcompress .DecompressWithDictionary (wasm , int (maxWasmSize ), dict )
334409}
335410
336411// Gets a program entry, which may be expired or not yet activated.
0 commit comments