Use custom code for WASM custom sections

Use our own code to read/write WASM custom sections. This allows us to remove `wasm-gen` as a dependency, and also to support removing custom sections with `set-metadata`.

Author: Timmmm

Cargo.lock

@@ -38,8 +38,7 @@ serde_json5 = "0.2.1"
 tokio = { version = "1.43.0", features = ["rt", "macros", "fs"] }
 tokio-util = { version = "0.7.13", features = ["io",] }
 walkdir = "2.5.0"
-wasm-gen = "0.1.4"
-wasmtime = { version = "32.0.0", features = ["reexport-wasmparser"] }
+wasmtime = { version = "32.0.0" }
 wasmtime-wasi = "32.0.0"
 
 # OpenSSL is only used on Unix. We don't want to add it as a dependency on Windows.

Cargo.toml

@@ -0,0 +1,66 @@
+/// Convert a u32 to LEB128 encoding. This is used in the WASM binary format.
+/// I got a little bit carried away here making a version without many branches.
+/// I have exhaustively tested that it is correct using the `exhaustive_leb128_test()`
+/// test below.
+pub fn u32_to_leb128(n: u32) -> Vec<u8> {
+    let mut leb = vec![
+        ((n >> 0) as u8 & 0x7F) | (if n >> 7 == 0 { 0 } else { 0x80 }),
+        ((n >> 7) as u8 & 0x7F) | (if n >> 14 == 0 { 0 } else { 0x80 }),
+        ((n >> 14) as u8 & 0x7F) | (if n >> 21 == 0 { 0 } else { 0x80 }),
+        ((n >> 21) as u8 & 0x7F) | (if n >> 28 == 0 { 0 } else { 0x80 }),
+        ((n >> 28) as u8 & 0x7F),
+    ];
+    let num_bytes = n.checked_ilog2().map(|n: u32| n / 7 + 1).unwrap_or(1);
+    leb.resize(num_bytes as usize, 0);
+    leb
+}
+
+/// Convert leb128 to u32, returning the value and number of bytes read.
+/// If there are insufficient bytes we return None.
+pub fn leb128_to_u32(leb: &[u8]) -> Option<(u32, usize)> {
+    let mut n = 0u32;
+    for i in 0..5 {
+        let byte = leb.get(i)?;
+        n |= ((byte & 0x7F) as u32) << (i * 7);
+        if byte & 0x80 == 0 {
+            return Some((n, i + 1));
+        }
+    }
+    Some((n, 5))
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    fn u32_to_leb128_simple(n: u32) -> Vec<u8> {
+        let mut leb = Vec::with_capacity(5);
+
+        for septet in 0..5 {
+            let byte = (n >> (7 * septet)) as u8 & 0x7F;
+            if septet == 4 || (n >> (7 * (septet + 1))) == 0 {
+                // Last byte.
+                leb.push(byte);
+                break;
+            } else {
+                leb.push(byte | 0x80);
+            }
+        }
+        leb
+    }
+
+    /// To run:
+    ///
+    ///   cargo test --release -- --ignored exhaustive_leb128_test
+    ///
+    /// It takes about 5 minutes.
+    #[test]
+    #[ignore]
+    fn exhaustive_leb128_test() {
+        for i in 0..=u32::MAX {
+            let correct_encoding = u32_to_leb128_simple(i);
+            assert_eq!(u32_to_leb128(i), correct_encoding);
+            assert!(matches!(leb128_to_u32(&correct_encoding), Some((v, _)) if v == i));
+        }
+    }
+}

src/leb128.rs

@@ -5,11 +5,13 @@ mod fetch;
 mod file_matching;
 mod git;
 mod hash_adapter;
+mod leb128;
 mod metadata;
 mod serde_glob;
 mod serde_regex;
 mod unique_filename;
 mod wasi_cache;
+mod wasm;
 
 use anyhow::{Result, anyhow, bail};
 use bash_paths::path_to_bash_string;
@@ -21,10 +23,11 @@ use fetch::fetch_linters;
 use file_matching::retain_matching_files;
 use git::git_diff_unstaged;
 use log::info;
-use metadata::{has_metadata, read_metadata};
+use metadata::read_metadata;
 use owo_colors::OwoColorize;
 use std::path::{Path, PathBuf};
 use tokio::fs;
+use wasm::{find_custom_sections, make_custom_section};
 
 #[derive(Parser)]
 #[command(
@@ -358,16 +361,20 @@ async fn subcommand_show_metadata(_cli: &Cli, args: &ShowMetadataArgs) -> Result
 }
 
 async fn subcommand_set_metadata(_cli: &Cli, args: &SetMetadataArgs) -> Result<()> {
-    // TODO (1.0): Remove any existing custom metadata sections.
-
     let mut bytes = fs::read(&args.file).await?;
-    if has_metadata(&bytes)? {
-        bail!("File already has metadata. Removing it is not yet supported.");
-    }
     let metadata_bytes = fs::read(&args.metadata).await?;
 
-    // TODO (1.0): This is simple enough we can do it without an external crate.
-    wasm_gen::write_custom_section(&mut bytes, "nit_metadata", &metadata_bytes);
+    // Find the existing metadata sections.
+    let (section_ranges, _) = find_custom_sections(&bytes, "nit_metadata")?;
+
+    // Remove them all.
+    for range in section_ranges.into_iter().rev() {
+        bytes.drain(range);
+    }
+
+    // Add a new section on the end.
+    let metadata_section = make_custom_section("nit_metadata", &metadata_bytes);
+    bytes.extend_from_slice(&metadata_section);
 
     fs::write(&args.file, bytes).await?;
 

src/main.rs

@@ -1,9 +1,8 @@
 use anyhow::{Context, Result, anyhow, bail};
 use serde::Deserialize;
 use std::path::Path;
-use wasmtime::wasmparser::{Parser, Payload};
 
-use crate::file_matching::MatchExpression;
+use crate::{file_matching::MatchExpression, wasm::find_custom_sections};
 
 #[derive(Debug, Deserialize)]
 pub struct ArgBlock {
@@ -53,35 +52,21 @@ pub struct NitMetadata {
 ///     cargo install wasm-custom-section
 ///
 pub fn read_metadata(wasm_path: &Path) -> Result<NitMetadata> {
-    let module = std::fs::read(wasm_path)?;
+    let wasm_bytes = std::fs::read(wasm_path)?;
 
     // Ideally we wouldn't load the entire file into memory, but
-    // it's probably fine in most cases and wasmparser doesn't provide
-    // a handy parse_reader() method.
+    // it's probably fine in most cases.
 
-    let parser = Parser::new(0);
-    for payload in parser.parse_all(&module) {
-        match payload? {
-            Payload::CustomSection(section) if section.name() == "nit_metadata" => {
-                return Ok(serde_json::from_slice::<NitMetadata>(section.data())
-                    .with_context(|| anyhow!("Reading metadata for {}", wasm_path.display()))?);
-            }
-            _ => {}
-        }
-    }
-    bail!("No nit_metadata section found in the wasm file");
-}
+    let (_, section_contents) = find_custom_sections(&wasm_bytes, "nit_metadata")
+        .context("Finding nit_metadata section")?;
 
-/// Return true if the file has a section called `nit_metadata`.
-pub fn has_metadata(module: &[u8]) -> Result<bool> {
-    let parser = Parser::new(0);
-    for payload in parser.parse_all(&module) {
-        match payload? {
-            Payload::CustomSection(section) if section.name() == "nit_metadata" => {
-                return Ok(true);
-            }
-            _ => {}
-        }
+    if section_contents.is_empty() {
+        bail!("No nit_metadata section found in the wasm file");
     }
-    Ok(false)
+    if section_contents.len() > 1 {
+        bail!("Multiple nit_metadata sections found in the wasm file");
+    }
+
+    Ok(serde_json::from_slice::<NitMetadata>(section_contents[0])
+        .with_context(|| anyhow!("Reading metadata for {}", wasm_path.display()))?)
 }