arb_stylus/
lib.rs

1//! Stylus WASM smart contract runtime.
2//!
3//! Provides the execution pipeline for Stylus programs: WASM compilation
4//! and caching, ink metering, host I/O functions, and EVM interop.
5
6pub mod cache;
7pub mod config;
8pub mod env;
9pub mod error;
10pub mod evm_api;
11pub mod evm_api_impl;
12#[allow(unused_mut)]
13pub mod host;
14pub mod ink;
15pub mod meter;
16pub mod middleware;
17pub mod multi_gas;
18pub mod native;
19pub mod pricing;
20pub mod run;
21pub mod trace;
22
23pub use cache::InitCache;
24pub use config::{CompileConfig, StylusConfig};
25pub use error::{MaybeEscape, StylusError};
26pub use evm_api::EvmApi;
27pub use evm_api_impl::StylusEvmApi;
28pub use ink::{Gas, Ink};
29pub use meter::{MachineMeter, MeteredMachine, STYLUS_ENTRY_POINT};
30pub use native::{NativeInstance, compile_module};
31pub use run::RunProgram;
32
33/// Prefix bytes that identify a Stylus WASM program in contract bytecode.
34///
35/// The discriminant is `[0xEF, 0xF0, 0x00]`. The `0xEF` byte conflicts with
36/// EIP-3541, so EIP-3541 must be disabled for Stylus-era blocks to allow
37/// deployment. The third byte `0x00` is the EOF version marker.
38pub const STYLUS_DISCRIMINANT: [u8; 3] = [0xEF, 0xF0, 0x00];
39
40/// Returns `true` if the bytecode is a Stylus WASM program.
41///
42/// Checks for the 3-byte discriminant prefix `[0xEF, 0xF0, 0x00]`.
43pub fn is_stylus_program(bytecode: &[u8]) -> bool {
44    bytecode.len() >= 4 && bytecode[..3] == STYLUS_DISCRIMINANT
45}
46
47/// Strips the 4-byte Stylus prefix from contract bytecode.
48///
49/// Returns `(stripped_bytecode, version_byte)` or an error if the bytecode
50/// is too short or doesn't have the Stylus discriminant.
51pub fn strip_stylus_prefix(bytecode: &[u8]) -> Result<(&[u8], u8), StylusError> {
52    if bytecode.len() < 4 {
53        return Err(StylusError::InvalidProgram(
54            "bytecode too short for Stylus prefix",
55        ));
56    }
57    if bytecode[..3] != STYLUS_DISCRIMINANT {
58        return Err(StylusError::InvalidProgram(
59            "bytecode does not have Stylus discriminant",
60        ));
61    }
62    let version = bytecode[3];
63    Ok((&bytecode[4..], version))
64}
65
66/// Root Stylus program prefix: `[0xEF, 0xF0, 0x02]`.
67pub const STYLUS_ROOT_DISCRIMINANT: [u8; 3] = [0xEF, 0xF0, 0x02];
68
69/// Fragment prefix: `[0xEF, 0xF0, 0x01]`.
70pub const STYLUS_FRAGMENT_DISCRIMINANT: [u8; 3] = [0xEF, 0xF0, 0x01];
71
72/// Returns `true` if the bytecode is a classic Stylus program (`[0xEF, 0xF0, 0x00, ...]`).
73pub fn is_stylus_classic(bytecode: &[u8]) -> bool {
74    bytecode.len() > 3 && bytecode[..3] == STYLUS_DISCRIMINANT
75}
76
77/// Returns `true` if the bytecode is a Stylus root program (`[0xEF, 0xF0, 0x02, ...]`).
78pub fn is_stylus_root(bytecode: &[u8]) -> bool {
79    bytecode.len() > 3 && bytecode[..3] == STYLUS_ROOT_DISCRIMINANT
80}
81
82/// Returns `true` if the bytecode is a Stylus fragment (`[0xEF, 0xF0, 0x01, ...]`).
83pub fn is_stylus_fragment(bytecode: &[u8]) -> bool {
84    bytecode.len() > 3 && bytecode[..3] == STYLUS_FRAGMENT_DISCRIMINANT
85}
86
87/// Returns `true` if the bytecode is a runnable Stylus program: a classic or a
88/// root program (a fragment is not runnable on its own). Root code can only
89/// have been deployed at the contract-limit version, so no version gate is
90/// needed here.
91pub fn is_stylus_runnable(bytecode: &[u8]) -> bool {
92    is_stylus_classic(bytecode) || is_stylus_root(bytecode)
93}
94
95/// Returns `true` if the bytecode is a deployable Stylus component: a classic
96/// or root program, or (at the contract-limit version) a fragment. Used to
97/// permit storing such code despite its `0xEF` prefix, mirroring
98/// `IsStylusComponentPrefix`.
99pub fn is_stylus_component(bytecode: &[u8], arbos_version: u64) -> bool {
100    use arb_chainspec::arbos_version as av;
101    if arbos_version < av::ARBOS_VERSION_STYLUS_CONTRACT_LIMIT {
102        return is_stylus_deployable(bytecode, arbos_version);
103    }
104    is_stylus_deployable(bytecode, arbos_version) || is_stylus_fragment(bytecode)
105}
106
107/// Returns `true` if the bytecode is a deployable Stylus program.
108pub fn is_stylus_deployable(bytecode: &[u8], arbos_version: u64) -> bool {
109    use arb_chainspec::arbos_version as av;
110    if arbos_version < av::ARBOS_VERSION_STYLUS {
111        return false;
112    }
113    if arbos_version < av::ARBOS_VERSION_STYLUS_CONTRACT_LIMIT {
114        return is_stylus_classic(bytecode);
115    }
116    is_stylus_classic(bytecode) || is_stylus_root(bytecode)
117}
118
119/// Decompress a Stylus WASM program from its contract bytecode.
120///
121/// The bytecode format is `[0xEF, 0xF0, 0x00, dict_byte, ...compressed_wasm]`.
122pub fn decompress_wasm(bytecode: &[u8]) -> Result<Vec<u8>, StylusError> {
123    if bytecode.len() < 4 || bytecode[..3] != STYLUS_DISCRIMINANT {
124        return Err(StylusError::InvalidProgram("not a Stylus program"));
125    }
126    let dict_byte = bytecode[3];
127    let compressed = &bytecode[4..];
128
129    let dict = match dict_byte {
130        0 => nitro_brotli::Dictionary::Empty,
131        1 => nitro_brotli::Dictionary::StylusProgram,
132        _ => return Err(StylusError::InvalidProgram("unsupported dictionary type")),
133    };
134
135    nitro_brotli::decompress(compressed, dict)
136        .map_err(|e| StylusError::Decompression(format!("{e:?}")))
137}
138
139/// Compress raw WASM into classic Stylus contract bytecode.
140///
141/// Produces `[0xEF, 0xF0, 0x00, 0x00, ...brotli(wasm)]` (empty dictionary),
142/// the inverse of [`decompress_wasm`]. This is the on-chain form a Stylus
143/// deploy must store for activation to reconstruct the program; raw WASM after
144/// the discriminant is rejected because byte 3 is read as the dictionary type
145/// and the remainder is brotli-decompressed.
146pub fn compress_classic_program_code(wasm: &[u8]) -> Result<Vec<u8>, StylusError> {
147    let compressed = nitro_brotli::compress(wasm, 11, 22, nitro_brotli::Dictionary::Empty)
148        .map_err(|e| StylusError::Decompression(format!("{e:?}")))?;
149    let mut out = Vec::with_capacity(4 + compressed.len());
150    out.extend_from_slice(&STYLUS_DISCRIMINANT);
151    out.push(0);
152    out.extend_from_slice(&compressed);
153    Ok(out)
154}
155
156/// A parsed Stylus root program. The on-chain layout is
157/// `[0xEF, 0xF0, 0x02, dict, decompressed_len(4, big-endian), addr×20...]`,
158/// where each 20-byte address points to a fragment holding part of the
159/// compressed WASM.
160#[derive(Debug, Clone)]
161pub struct StylusRoot {
162    pub dictionary: u8,
163    pub decompressed_length: u32,
164    pub addresses: Vec<alloy_primitives::Address>,
165}
166
167impl StylusRoot {
168    /// Parse a root program's contract bytecode.
169    pub fn parse(bytecode: &[u8]) -> Result<Self, StylusError> {
170        if !is_stylus_root(bytecode) {
171            return Err(StylusError::InvalidProgram("not a Stylus program root"));
172        }
173        if bytecode.len() < 8 {
174            return Err(StylusError::InvalidProgram("Stylus root too short"));
175        }
176        let address_data = &bytecode[8..];
177        if !address_data.len().is_multiple_of(20) {
178            return Err(StylusError::InvalidProgram(
179                "Stylus root address data misaligned",
180            ));
181        }
182        let addresses = address_data
183            .chunks_exact(20)
184            .map(alloy_primitives::Address::from_slice)
185            .collect();
186        Ok(Self {
187            dictionary: bytecode[3],
188            decompressed_length: u32::from_be_bytes([
189                bytecode[4],
190                bytecode[5],
191                bytecode[6],
192                bytecode[7],
193            ]),
194            addresses,
195        })
196    }
197}
198
199/// Reconstruct the WASM of a root Stylus program: read each fragment's deployed
200/// bytecode via `read_code`, strip its prefix, concatenate the compressed
201/// payloads, and decompress with the root's dictionary. When `enforce` is set
202/// (i.e. activation), the decompressed-length and fragment-count limits are
203/// applied; reads otherwise only reconstruct the program.
204pub fn get_wasm_from_root(
205    root: &[u8],
206    max_wasm_size: u32,
207    max_fragments: u8,
208    enforce: bool,
209    mut read_code: impl FnMut(alloy_primitives::Address) -> Result<Vec<u8>, StylusError>,
210) -> Result<Vec<u8>, StylusError> {
211    let parsed = StylusRoot::parse(root)?;
212    if enforce {
213        if parsed.decompressed_length > max_wasm_size {
214            return Err(StylusError::InvalidProgram(
215                "decompressed length exceeds max wasm size",
216            ));
217        }
218        if parsed.addresses.len() > max_fragments as usize {
219            return Err(StylusError::InvalidProgram("fragment count exceeds limit"));
220        }
221    }
222    if parsed.addresses.is_empty() {
223        return Err(StylusError::InvalidProgram("fragment count cannot be zero"));
224    }
225    let dict = match parsed.dictionary {
226        0 => nitro_brotli::Dictionary::Empty,
227        1 => nitro_brotli::Dictionary::StylusProgram,
228        _ => return Err(StylusError::InvalidProgram("unsupported dictionary type")),
229    };
230    let mut compressed = Vec::new();
231    for addr in &parsed.addresses {
232        let fragment = read_code(*addr)?;
233        if fragment.len() <= 3 || fragment[..3] != STYLUS_FRAGMENT_DISCRIMINANT {
234            return Err(StylusError::InvalidProgram(
235                "referenced code is not a Stylus fragment",
236            ));
237        }
238        compressed.extend_from_slice(&fragment[3..]);
239    }
240    let wasm = nitro_brotli::decompress(&compressed, dict)
241        .map_err(|e| StylusError::Decompression(format!("{e:?}")))?;
242    if wasm.len() != parsed.decompressed_length as usize {
243        return Err(StylusError::InvalidProgram(
244            "decompressed length does not match the declared length",
245        ));
246    }
247    Ok(wasm)
248}
249
250/// Gas charged for reading one fragment of `code_size` bytes during activation,
251/// matching `fragmentReadGasCost`: a cold (or warm) account access plus the
252/// per-word copy cost. The fragment-read charger uses this for both the
253/// preflight affordability check (against the max code size) and the actual
254/// per-fragment charge.
255pub fn fragment_read_gas(warm: bool, code_size: u64) -> u64 {
256    const WARM_ACCESS: u64 = 100; // WarmStorageReadCostEIP2929
257    const COLD_ACCESS: u64 = 2_600; // ColdAccountAccessCostEIP2929
258    const COPY: u64 = 3; // CopyGas
259    let base = if warm { WARM_ACCESS } else { COLD_ACCESS };
260    let words = code_size.div_ceil(32);
261    base.saturating_add(words.saturating_mul(COPY))
262}
263
264/// Activate a Stylus program.
265///
266/// `wasm` must be the decompressed WASM bytes (call `decompress_wasm` first).
267/// `gas` is decremented by the activation cost.
268pub fn activate_program(
269    wasm: &[u8],
270    codehash: &[u8; 32],
271    stylus_version: u16,
272    arbos_version: u64,
273    page_limit: u16,
274    debug: bool,
275    gas: &mut u64,
276) -> Result<arbos::programs::types::ActivationResult, StylusError> {
277    let codehash_bytes32 = nitro_arbutil::Bytes32(*codehash);
278    let (module, stylus_data) = nitro_prover::machine::Module::activate(
279        wasm,
280        &codehash_bytes32,
281        stylus_version,
282        arbos_version,
283        page_limit,
284        debug,
285        gas,
286    )
287    .map_err(|e| StylusError::Activation(format!("{e}")))?;
288
289    Ok(arbos::programs::types::ActivationResult {
290        module_hash: alloy_primitives::B256::from(module.hash().0),
291        init_gas: stylus_data.init_cost,
292        cached_init_gas: stylus_data.cached_init_cost,
293        asm_estimate: stylus_data.asm_estimate,
294        footprint: stylus_data.footprint,
295    })
296}
297
298#[cfg(test)]
299mod stylus_root_tests {
300    use alloy_primitives::Address;
301
302    use super::*;
303
304    fn make_fragment(chunk: &[u8]) -> Vec<u8> {
305        let mut f = STYLUS_FRAGMENT_DISCRIMINANT.to_vec();
306        f.extend_from_slice(chunk);
307        f
308    }
309
310    fn make_root(dict: u8, decompressed_len: u32, addrs: &[Address]) -> Vec<u8> {
311        let mut r = STYLUS_ROOT_DISCRIMINANT.to_vec();
312        r.push(dict);
313        r.extend_from_slice(&decompressed_len.to_be_bytes());
314        for a in addrs {
315            r.extend_from_slice(a.as_slice());
316        }
317        r
318    }
319
320    #[test]
321    fn parse_extracts_fields() {
322        let a = Address::repeat_byte(0xab);
323        let root = make_root(1, 0x1234_5678, &[a]);
324        let p = StylusRoot::parse(&root).unwrap();
325        assert_eq!(p.dictionary, 1);
326        assert_eq!(p.decompressed_length, 0x1234_5678);
327        assert_eq!(p.addresses, vec![a]);
328    }
329
330    #[test]
331    fn root_reconstructs_wasm_across_fragments() {
332        let payload =
333            b"\x00asm\x01\x00\x00\x00 stylus wasm body bytes for the fragment round trip".to_vec();
334        let compressed =
335            nitro_brotli::compress(&payload, 0, 22, nitro_brotli::Dictionary::Empty).unwrap();
336        let mid = compressed.len() / 2;
337        let frag0 = make_fragment(&compressed[..mid]);
338        let frag1 = make_fragment(&compressed[mid..]);
339        let a0 = Address::repeat_byte(0x11);
340        let a1 = Address::repeat_byte(0x22);
341        let root = make_root(0, payload.len() as u32, &[a0, a1]);
342        let out = get_wasm_from_root(&root, 100_000, 4, true, |a| {
343            Ok(if a == a0 {
344                frag0.clone()
345            } else {
346                frag1.clone()
347            })
348        })
349        .unwrap();
350        assert_eq!(out, payload);
351    }
352
353    #[test]
354    fn decompressed_length_mismatch_rejected_even_unenforced() {
355        let payload =
356            b"\x00asm\x01\x00\x00\x00 stylus wasm body bytes for the fragment round trip".to_vec();
357        let compressed =
358            nitro_brotli::compress(&payload, 0, 22, nitro_brotli::Dictionary::Empty).unwrap();
359        let frag = make_fragment(&compressed);
360        let a = Address::repeat_byte(0x11);
361        // Declare a length one byte longer than the fragments actually decompress to.
362        let root = make_root(0, payload.len() as u32 + 1, &[a]);
363        let err = get_wasm_from_root(&root, 100_000, 4, false, |_| Ok(frag.clone())).unwrap_err();
364        assert!(matches!(
365            err,
366            StylusError::InvalidProgram("decompressed length does not match the declared length")
367        ));
368    }
369
370    #[test]
371    fn fragment_count_zero_rejected() {
372        let root = make_root(0, 10, &[]);
373        let err = get_wasm_from_root(&root, 100_000, 4, true, |_| Ok(Vec::new())).unwrap_err();
374        assert!(matches!(err, StylusError::InvalidProgram(_)));
375    }
376
377    #[test]
378    fn fragment_count_over_limit_rejected_when_enforced() {
379        let addrs: Vec<Address> = (0..5).map(Address::repeat_byte).collect();
380        let root = make_root(0, 10, &addrs);
381        let err =
382            get_wasm_from_root(&root, 100_000, 4, true, |_| Ok(make_fragment(&[]))).unwrap_err();
383        assert!(matches!(
384            err,
385            StylusError::InvalidProgram("fragment count exceeds limit")
386        ));
387    }
388
389    #[test]
390    fn decompressed_length_over_max_rejected_when_enforced() {
391        let root = make_root(0, 1000, &[Address::repeat_byte(1)]);
392        let err = get_wasm_from_root(&root, 500, 4, true, |_| Ok(make_fragment(&[]))).unwrap_err();
393        assert!(matches!(
394            err,
395            StylusError::InvalidProgram("decompressed length exceeds max wasm size")
396        ));
397    }
398
399    #[test]
400    fn fragment_read_gas_matches_reference_constants() {
401        // cold: 2600 + ceil(64/32)*3 = 2606; warm: 100 + 2*3 = 106.
402        assert_eq!(fragment_read_gas(false, 64), 2_600 + 2 * 3);
403        assert_eq!(fragment_read_gas(true, 64), 100 + 2 * 3);
404        // zero-length code: just the account access.
405        assert_eq!(fragment_read_gas(false, 0), 2_600);
406        // partial word rounds up.
407        assert_eq!(fragment_read_gas(true, 33), 100 + 2 * 3);
408    }
409
410    #[test]
411    fn non_fragment_code_rejected() {
412        let root = make_root(0, 10, &[Address::repeat_byte(1)]);
413        let err = get_wasm_from_root(&root, 100_000, 4, true, |_| {
414            Ok(vec![0xEF, 0xF0, 0x00, 0x00, 0x99])
415        })
416        .unwrap_err();
417        assert!(matches!(
418            err,
419            StylusError::InvalidProgram("referenced code is not a Stylus fragment")
420        ));
421    }
422}