@@ -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,108 @@ 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 )
320339 if err != nil {
321340 return nil , err
322341 }
323342
324- var dict arbcompress.Dictionary
325- switch dictByte {
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+ rawAddresses , dictByte , decompressedLength , err := state .StripStylusRootPrefix (data )
353+ if err != nil {
354+ return nil , err
355+ }
356+ if isActivation && decompressedLength > maxSize {
357+ return nil , fmt .Errorf ("invalid wasm: decompressedLength %d is greater then MaxWasmSize %d" , decompressedLength , maxSize )
358+ }
359+ if isActivation && len (rawAddresses )/ common .AddressLength > int (maxFragments ) {
360+ return nil , fmt .Errorf ("invalid wasm: fragment count exceeds limit of %d" , maxFragments )
361+ }
362+
363+ if len (rawAddresses )/ common .AddressLength == 0 {
364+ return nil , fmt .Errorf ("invalid wasm: fragment count cannot be zero" )
365+ }
366+
367+ var compressedWasm []byte
368+ for i := 0 ; i < len (rawAddresses ); i += common .AddressLength {
369+ addr := common .BytesToAddress (rawAddresses [i : i + common .AddressLength ])
370+
371+ fragCode := statedb .GetCode (addr )
372+ payload , err := state .StripStylusFragmentPrefix (fragCode )
373+ if err != nil {
374+ return nil , err
375+ }
376+
377+ if len (payload ) == 0 {
378+ continue
379+ }
380+
381+ compressedWasm = append (compressedWasm , payload ... )
382+ }
383+
384+ dict , err := getStylusCompressionDict (dictByte )
385+ if err != nil {
386+ return nil , err
387+ }
388+
389+ wasm , err := arbcompress .DecompressWithDictionary (compressedWasm , int (decompressedLength ), dict )
390+
391+ if err != nil {
392+ return nil , err
393+ }
394+
395+ if len (wasm ) != int (decompressedLength ) {
396+ return nil , fmt .Errorf ("invalid wasm: decompressed length %d does not match expected length %d" , len (wasm ), decompressedLength )
397+ }
398+
399+ return wasm , nil
400+ }
401+
402+ // Named return parameters allow us to return the zero-value for 'dict' implicitly on error
403+ func getStylusCompressionDict (id byte ) (dict arbcompress.Dictionary , err error ) {
404+ switch id {
326405 case 0 :
327- dict = arbcompress .EmptyDictionary
406+ return arbcompress .EmptyDictionary , nil
328407 case 1 :
329- dict = arbcompress .StylusProgramDictionary
408+ return arbcompress .StylusProgramDictionary , nil
330409 default :
331- return nil , fmt .Errorf ("unsupported dictionary %v " , dictByte )
410+ return dict , fmt .Errorf ("unsupported dictionary type: %d " , id )
332411 }
333- return arbcompress .DecompressWithDictionary (wasm , int (maxWasmSize ), dict )
334412}
335413
336414// Gets a program entry, which may be expired or not yet activated.
0 commit comments