WebAssembly integration

Understanding Scramjet’s Rust-powered WASM rewriter for high-performance JavaScript transformation

Scramjet’s JavaScript rewriter is implemented in Rust and compiled to WebAssembly for maximum performance. This architecture enables AST-level transformations with near-native speed.

Architecture overview

The WASM rewriter is a Cargo workspace with multiple components:

rewriter/
├── Cargo.toml          # Workspace configuration
├── js/                 # JavaScript rewriter logic
├── wasm/               # WASM bindings
├── transform/          # AST transformations
└── native/             # Native binary for testing

Key dependencies

[workspace.dependencies]
oxc = { version = "0.77.2", features = ["ast_visit"] }

[profile.release]
opt-level = 3
lto = true              # Link-time optimization
codegen-units = 1       # Single codegen unit for better optimization
panic = "abort"         # Smaller binary size

Scramjet uses OXC, a blazingly fast JavaScript parser written in Rust, achieving performance comparable to native compilers.

WASM module structure

Initialization

The WASM module is loaded and initialized in src/shared/rewriters/wasm.ts:

import { initSync, Rewriter } from "../../../rewriter/wasm/out/wasm.js";

let wasm_u8: Uint8Array;

// Load WASM binary (embedded or fetched)
if (REWRITERWASM) {
  // Embedded at build time
  wasm_u8 = Uint8Array.from(atob(REWRITERWASM), c => c.charCodeAt(0));
} else if (self.WASM) {
  // Injected by service worker
  wasm_u8 = Uint8Array.from(atob(self.WASM), c => c.charCodeAt(0));
}

function initWasm() {
  // Validate WASM magic bytes
  const MAGIC = "\0asm".split("").map(x => x.charCodeAt(0));
  if (![...wasm_u8.slice(0, 4)].every((x, i) => x === MAGIC[i])) {
    throw new Error("Invalid WASM binary");
  }
  
  initSync({
    module: new WebAssembly.Module(wasm_u8),
  });
}

The WASM module must be loaded before any JavaScript rewriting occurs. Scramjet automatically handles this in service workers via asyncSetWasm().

Service worker WASM loading

In service workers, WASM is fetched dynamically:

export async function asyncSetWasm() {
  const buf = await fetch(config.files.wasm).then(r => r.arrayBuffer());
  wasm_u8 = new Uint8Array(buf);
}

// Called during service worker initialization
await scramjet.loadConfig();
// -> setConfig(config);
// -> await asyncSetWasm();

Rewriter class

The WASM Rewriter class exposes a JavaScript-friendly API:

// Simplified from rewriter/wasm/src/lib.rs
#[wasm_bindgen]
pub struct Rewriter {
    alloc: Allocator,
    scramjet: Object,
    js: JsRewriter,
}

#[wasm_bindgen]
impl Rewriter {
    #[wasm_bindgen(constructor)]
    pub fn new(scramjet: Object) -> Result<Self> {
        Ok(Self {
            alloc: Allocator::default(),
            js: create_js(&scramjet)?,
            scramjet,
        })
    }
    
    #[wasm_bindgen]
    pub fn rewrite_js(
        &mut self,
        js: String,
        base: String,
        url: String,
        module: bool,
    ) -> Result<JsRewriterOutput> {
        let flags = get_js_flags(&self.scramjet, base, module)?;
        let out = self.js.rewrite(&self.alloc, &js, flags)?;
        
        let ret = create_js_output(out, url, js);
        self.alloc.reset(); // Reuse allocator
        ret
    }
    
    #[wasm_bindgen]
    pub fn rewrite_js_bytes(
        &mut self,
        js: Vec<u8>,
        base: String,
        url: String,
        module: bool,
    ) -> Result<JsRewriterOutput> {
        // SAFETY: Assumes valid UTF-8
        let js = unsafe { String::from_utf8_unchecked(js) };
        self.rewrite_js(js, base, url, module)
    }
}

TypeScript bindings

The Rust code generates TypeScript definitions:

export class Rewriter {
  constructor(scramjet: object);
  
  rewrite_js(
    js: string,
    base: string,
    url: string,
    module: boolean
  ): JsRewriterOutput;
  
  rewrite_js_bytes(
    js: Uint8Array,
    base: string,
    url: string,
    module: boolean
  ): JsRewriterOutput;
}

export interface JsRewriterOutput {
  js: Uint8Array;
  map: Uint8Array;
  scramtag: string;
  errors: string[];
}

Rewriter configuration

The Rewriter constructor accepts a configuration object:

const rewriter = new Rewriter({
  config: scramjetConfig,
  shared: {
    rewrite: {
      htmlRules,
      rewriteUrl,
      rewriteCss,
      rewriteJs,
      getHtmlInjectCode(cookieStore, foundHead) {
        let inject = getInjectScripts(
          cookieStore,
          src => `<script src="${src}"></script>`
        ).join("");
        return foundHead ? `<head>${inject}</head>` : inject;
      },
    },
  },
  flagEnabled,
  codec: {
    encode: codecEncode,
    decode: codecDecode,
  },
});

Shared functions

The rewriter calls back into JavaScript for certain operations:

• URL rewriting: rewriteUrl(url, meta)
• CSS rewriting: rewriteCss(css, meta) (for <style> tags)
• JS rewriting: rewriteJs(js, url, meta) (for nested scripts)
• HTML injection: getHtmlInjectCode(cookieStore, foundHead)

Why callbacks instead of pure Rust?

While implementing everything in Rust would be faster, JavaScript callbacks provide:

1. Flexibility: Configuration (codec, flags) can be changed without recompiling WASM
2. Code sharing: URL/CSS rewriting logic is shared between client and worker contexts
3. Smaller binary: Avoiding duplicate implementations reduces WASM size
4. Maintainability: Complex logic (HTML parsing) stays in TypeScript where it’s easier to debug

AST transformation

The Rust rewriter performs AST-level transformations using OXC:

use oxc::{
    allocator::Allocator,
    ast_visit::Visit,
    parser::{Parser, ParseOptions},
    span::SourceType,
};

pub fn rewrite<'alloc>(
    &self,
    alloc: &'alloc Allocator,
    js: &str,
    flags: Flags,
) -> Result<RewriteResult<'alloc>> {
    let source_type = SourceType::unambiguous()
        .with_javascript(true)
        .with_module(flags.is_module)
        .with_standard(true);
    
    let parsed = Parser::new(alloc, js, source_type)
        .with_options(ParseOptions {
            allow_v8_intrinsics: true,
            allow_return_outside_function: true,
            ..Default::default()
        })
        .parse();
    
    if parsed.panicked {
        return Err(RewriterError::OxcPanicked(format_errors(&parsed.errors)));
    }
    
    let mut visitor = Visitor {
        alloc,
        config: &self.cfg,
        rewriter: &self.url,
        flags,
        jschanges,
        error: None,
    };
    
    visitor.visit_program(&parsed.program);
    
    let changed = visitor.jschanges.perform(js, &self.cfg, &visitor.flags)?;
    
    Ok(RewriteResult {
        js: changed.source,
        sourcemap: changed.map,
        errors: parsed.errors,
        flags: visitor.flags,
    })
}

Visitor pattern

The visitor walks the AST and collects transformations:

impl<'a> Visit<'a> for Visitor<'a> {
    fn visit_call_expression(&mut self, expr: &CallExpression<'a>) {
        // Intercept: fetch(url)
        if is_fetch_call(expr) {
            let url_arg = &expr.arguments[0];
            self.rewrite_url_argument(url_arg);
        }
    }
    
    fn visit_new_expression(&mut self, expr: &NewExpression<'a>) {
        // Intercept: new WebSocket(url)
        if is_websocket_constructor(expr) {
            let url_arg = &expr.arguments[0];
            self.rewrite_url_argument(url_arg);
        }
    }
    
    fn visit_member_expression(&mut self, expr: &MemberExpression<'a>) {
        // Intercept: location.href
        if is_location_href(expr) {
            self.wrap_location_access(expr);
        }
    }
}

Building the WASM module

Scramjet includes a build script for the WASM module:

# rewriter/wasm/build.sh
cd rewriter/wasm/
wasm-pack build --target web --out-dir out
cd ../..

Build configuration

For production builds, Scramjet uses aggressive optimizations:

[profile.release]
opt-level = 3        # Maximum optimization
lto = true           # Link-time optimization
codegen-units = 1    # Single codegen unit
panic = "abort"      # Smaller binary (no unwinding)

The production WASM binary is approximately 500KB gzipped, which is loaded once and cached by the browser.

Sourcemap generation

The rewriter generates sourcemaps for debugging:

if (flagEnabled("sourcemaps", meta.base)) {
  const pushmap = globalThis[config.globals.pushsourcemapfn];
  if (pushmap) {
    // Direct function call (faster)
    pushmap(Array.from(res.map), res.tag);
  } else {
    // Inline sourcemap as code
    const sourcemapfn = `${config.globals.pushsourcemapfn}([${res.map.join(",")}], "${res.tag}");`;
    
    // Preserve "use strict" directive
    const strictMode = /^\s*(['"])use strict\1;?/;
    if (strictMode.test(newjs)) {
      newjs = newjs.replace(strictMode, `$&\n${sourcemapfn}`);
    } else {
      newjs = `${sourcemapfn}\n${newjs}`;
    }
  }
}

Sourcemaps enable:

• Accurate stack traces in DevTools
• Breakpoint debugging in original code
• Proper error line numbers

Error handling

Parse errors

The rewriter collects parse errors without failing:

const result = rewriteJs(code, url, meta, isModule);

if (flagEnabled("rewriterLogs", meta.base)) {
  for (const error of result.errors) {
    console.error("oxc parse error", error);
  }
}

Panic recovery

If OXC panics, the rewriter catches it:

if parsed.panicked {
    use std::fmt::Write;
    
    let mut errors = String::new();
    for error in parsed.errors {
        writeln!(errors, "{error}")?;
    }
    return Err(RewriterError::OxcPanicked(errors));
}

In TypeScript:

try {
  return rewriteJs(code, url, meta, isModule);
} catch (err) {
  console.warn("Failed rewriting JS:", err.message);
  
  if (flagEnabled("allowInvalidJs", meta.base)) {
    return code; // Return original
  }
  throw err;
}

Performance optimizations

Rewriter pooling

Scramjet maintains a pool of rewriter instances:

let rewriters: Array<{ rewriter: Rewriter; inUse: boolean }> = [];

export function getRewriter(meta: URLMeta): [Rewriter, () => void] {
  initWasm();
  
  let obj = rewriters.find(x => !x.inUse);
  
  if (!obj) {
    if (flagEnabled("rewriterLogs", meta.base)) {
      console.log(`Creating new rewriter, ${rewriters.length} already exist`);
    }
    
    const rewriter = new Rewriter({ config, shared, flagEnabled, codec });
    obj = { rewriter, inUse: false };
    rewriters.push(obj);
  }
  
  obj.inUse = true;
  return [obj.rewriter, () => (obj.inUse = false)];
}

Usage

function rewriteJsWasm(
  input: string | Uint8Array,
  source: string | null,
  meta: URLMeta,
  module: boolean
): RewriterResult {
  let [rewriter, release] = getRewriter(meta);
  
  try {
    const out = typeof input === "string"
      ? rewriter.rewrite_js(input, meta.base.href, source, module)
      : rewriter.rewrite_js_bytes(input, meta.base.href, source, module);
    
    return {
      js: typeof input === "string" ? textDecoder.decode(out.js) : out.js,
      map: out.map,
      tag: out.scramtag,
      errors: out.errors,
    };
  } finally {
    release(); // Return to pool
  }
}

Always release the rewriter back to the pool using the release() function to avoid memory leaks and rewriter exhaustion.

Testing the rewriter

Scramjet includes a native test runner:

# Build and run native binary
cd rewriter/native
cargo run --release

This allows testing the rewriter without compiling to WASM, which is useful for:

• Debugging Rust code
• Running benchmarks
• Integration tests

Debugging tips

Enable rewriter logs

import { flagEnabled } from "@/shared";

if (flagEnabled("rewriterLogs", meta.base)) {
  console.log("Rewriter created");
  console.log("Parse errors:", result.errors);
}

Inspect WASM binary

function initWasm() {
  const MAGIC = [0x00, 0x61, 0x73, 0x6D]; // "\0asm"
  
  console.log("WASM size:", wasm_u8.length, "bytes");
  console.log("Magic bytes:", [...wasm_u8.slice(0, 4)]);
  
  if (![...wasm_u8.slice(0, 4)].every((x, i) => x === MAGIC[i])) {
    console.error("Invalid WASM:", textDecoder.decode(wasm_u8.slice(0, 100)));
    throw new Error("Invalid WASM binary");
  }
}

Compare input/output

const original = code;
const rewritten = rewriteJs(code, url, meta, false);

console.log("Original:", original);
console.log("Rewritten:", rewritten);
console.log("Size change:", rewritten.length - original.length, "bytes");