Skip to content

Commit

Permalink
Support gzipped WASM inputs in utils::parse_wasm and `utils::parse_…
Browse files Browse the repository at this point in the history
…wasm_file` (#55)
  • Loading branch information
aterga committed Mar 20, 2024
1 parent 61692f4 commit b1c63d5
Show file tree
Hide file tree
Showing 5 changed files with 92 additions and 6 deletions.
33 changes: 33 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ walrus = "0.20.1"
candid = "0.10"
rustc-demangle = "0.1"
thiserror = "1.0.35"
libflate = "1.3.0"

wasm-opt = { version = "0.113.0", optional = true }
tempfile = { version = "3.5.0", optional = true }
Expand Down
35 changes: 29 additions & 6 deletions src/utils.rs
Original file line number Diff line number Diff line change
@@ -1,26 +1,49 @@
use crate::Error;
use libflate::gzip;
use std::borrow::Cow;
use std::collections::HashMap;
use std::io::{self, Read};
use walrus::*;

pub const WASM_MAGIC_BYTES: &[u8] = &[0, 97, 115, 109];

pub const GZIPPED_WASM_MAGIC_BYTES: &[u8] = &[31, 139, 8, 0];

fn wasm_parser_config(keep_name_section: bool) -> ModuleConfig {
let mut config = walrus::ModuleConfig::new();
config.generate_name_section(keep_name_section);
config.generate_producers_section(false);
config
}

pub fn parse_wasm(wasm: &[u8], keep_name_section: bool) -> Result<Module, Error> {
fn decompress(bytes: &[u8]) -> Result<Vec<u8>, std::io::Error> {
let mut decoder = gzip::Decoder::new(bytes)?;
let mut decoded_data = Vec::new();
decoder.read_to_end(&mut decoded_data)?;
Ok(decoded_data)
}

pub fn parse_wasm(bytes: &[u8], keep_name_section: bool) -> Result<Module, Error> {
let wasm = if bytes.starts_with(WASM_MAGIC_BYTES) {
Ok(Cow::Borrowed(bytes))
} else if bytes.starts_with(GZIPPED_WASM_MAGIC_BYTES) {
decompress(bytes).map(Cow::Owned)
} else {
Err(io::Error::new(
io::ErrorKind::InvalidInput,
"Input must be either gzipped or uncompressed WASM.",
))
}
.map_err(Error::IO)?;
let config = wasm_parser_config(keep_name_section);
config
.parse(wasm)
.parse(&wasm)
.map_err(|e| Error::WasmParse(e.to_string()))
}

pub fn parse_wasm_file(file: std::path::PathBuf, keep_name_section: bool) -> Result<Module, Error> {
let config = wasm_parser_config(keep_name_section);
config
.parse_file(file)
.map_err(|e| Error::WasmParse(e.to_string()))
let bytes = std::fs::read(file).map_err(Error::IO)?;
parse_wasm(&bytes[..], keep_name_section)
}

#[derive(Clone, Copy, PartialEq, Eq)]
Expand Down
29 changes: 29 additions & 0 deletions tests/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,11 @@ fn instrumentation() {
.assert()
.success();
assert_wasm("wat-instrument.wasm");
wasm_input("wat.wasm.gz", true)
.arg("instrument")
.assert()
.success();
assert_wasm("wat-instrument.wasm");
wasm_input("rust.wasm", true)
.arg("instrument")
.assert()
Expand All @@ -86,6 +91,11 @@ fn shrink() {
.assert()
.success();
assert_wasm("wat-shrink.wasm");
wasm_input("wat.wasm.gz", true)
.arg("shrink")
.assert()
.success();
assert_wasm("wat-shrink.wasm");
wasm_input("rust.wasm", true)
.arg("shrink")
.assert()
Expand Down Expand Up @@ -147,6 +157,14 @@ fn resource() {
.assert()
.success();
assert_wasm("wat-limit.wasm");
wasm_input("wat.wasm.gz", true)
.arg("resource")
.arg("--remove-cycles-transfer")
.arg("--limit-stable-memory-page")
.arg("32")
.assert()
.success();
assert_wasm("wat-limit.wasm");
wasm_input("rust.wasm", true)
.arg("resource")
.arg("--remove-cycles-transfer")
Expand Down Expand Up @@ -213,6 +231,11 @@ Custom sections with size: []
.assert()
.stdout(expected)
.success();
wasm_input("wat.wasm.gz", false)
.arg("info")
.assert()
.stdout(expected)
.success();
}

#[test]
Expand Down Expand Up @@ -256,6 +279,12 @@ fn json_info() {
.assert()
.stdout(expected)
.success();
wasm_input("wat.wasm.gz", false)
.arg("info")
.arg("--json")
.assert()
.stdout(expected)
.success();
}

#[test]
Expand Down
Binary file added tests/wat.wasm.gz
Binary file not shown.

0 comments on commit b1c63d5

Please sign in to comment.