Fix hotpatched assets (#4326)

* wip: fixes to get docsite working again

* Fix asset location with PIC in hot patched builds

* Add an hidden cli flag to exit early in hot patch tests on failure

* add a test for hot patching with assets

* fix clippy

* fix looks_like_manganis_symbol

* longer timeout for hot patching in CI

* log page content

* log assets

* Log asset after write

* record a playwright trace

* Explicitly enable cargo incremental for hot patching

* fix retry for patch test

* reset asset

* disable hot patching test for new exports

* add nightly playwright

* use different toolchain names

* use beta instead of nightly for less thrashing

* add failing test

* workaround hot reloading bug

* restore swapping asset tests

* add toolchain

---------

Co-authored-by: Jonathan Kelley <jkelleyrtp@gmail.com>

packages/cli/src/build/assets.rs

@@ -41,7 +41,6 @@ use manganis::BundledAsset;
 use object::{File, Object, ObjectSection, ObjectSymbol, ReadCache, ReadRef, Section, Symbol};
 use pdb::FallibleIterator;
 use rayon::iter::{IntoParallelRefMutIterator, ParallelIterator};
-use wasmparser::BinaryReader;
 
 /// Extract all manganis symbols and their sections from the given object file.
 fn manganis_symbols<'a, 'b, R: ReadRef<'a>>(
@@ -50,7 +49,7 @@ fn manganis_symbols<'a, 'b, R: ReadRef<'a>>(
     file.symbols()
         .filter(|symbol| {
             if let Ok(name) = symbol.name() {
-                name.contains("__MANGANIS__")
+                looks_like_manganis_symbol(name)
             } else {
                 false
             }
@@ -62,6 +61,10 @@ fn manganis_symbols<'a, 'b, R: ReadRef<'a>>(
         })
 }
 
+fn looks_like_manganis_symbol(name: &str) -> bool {
+    name.contains("__MANGANIS__")
+}
+
 /// Find the offsets of any manganis symbols in the given file.
 fn find_symbol_offsets<'a, R: ReadRef<'a>>(
     path: &Path,
@@ -175,105 +178,113 @@ fn find_native_symbol_offsets<'a, R: ReadRef<'a>>(file: &File<'a, R>) -> Result<
     Ok(offsets)
 }
 
+fn eval_walrus_global_expr(module: &walrus::Module, expr: &walrus::ConstExpr) -> Option<u64> {
+    match expr {
+        walrus::ConstExpr::Value(walrus::ir::Value::I32(value)) => Some(*value as u64),
+        walrus::ConstExpr::Value(walrus::ir::Value::I64(value)) => Some(*value as u64),
+        walrus::ConstExpr::Global(id) => {
+            let global = module.globals.get(*id);
+            if let walrus::GlobalKind::Local(pointer) = &global.kind {
+                eval_walrus_global_expr(module, pointer)
+            } else {
+                None
+            }
+        }
+        _ => None,
+    }
+}
+
 /// Find the offsets of any manganis symbols in the wasm file.
 fn find_wasm_symbol_offsets<'a, R: ReadRef<'a>>(
     file_contents: &[u8],
     file: &File<'a, R>,
 ) -> Result<Vec<u64>> {
-    // Parse the wasm file to find the globals
-    let parser = wasmparser::Parser::new(0);
+    let Some(section) = file
+        .sections()
+        .find(|section| section.name() == Ok("<data>"))
+    else {
+        tracing::error!("Failed to find <data> section in WASM file");
+        return Ok(Vec::new());
+    };
+    let Some((_, section_range_end)) = section.file_range() else {
+        tracing::error!("Failed to find file range for <data> section in WASM file");
+        return Ok(Vec::new());
+    };
+    let section_size = section.data()?.len() as u64;
+    let section_start = section_range_end - section_size;
 
-    // All integer literal global values in the wasm file
-    let mut global_values = Vec::new();
-    for section in parser.parse_all(file_contents) {
-        let Ok(wasmparser::Payload::GlobalSection(global_section)) = section else {
+    // Translate the section_relative_address to the file offset
+    // WASM files have a section address of 0 in object, reparse the data section with wasmparser
+    // to get the correct address and section start
+    // Note: We need to reparse just the data section with wasmparser to get the file offset because walrus does
+    // not expose the file offset information
+    let reader = wasmparser::DataSectionReader::new(wasmparser::BinaryReader::new(
+        &file_contents[section_start as usize..section_range_end as usize],
+        0,
+    ))
+    .context("Failed to create WASM data section reader")?;
+    let main_memory = reader
+        .into_iter()
+        .next()
+        .context("Failed find main memory from WASM data section")?
+        .context("Failed to read main memory from WASM data section")?;
+    // main_memory.data is a slice somewhere in file_contents. Find out the offset in the file
+    let data_start_offset = (main_memory.data.as_ptr() as u64)
+        .checked_sub(file_contents.as_ptr() as u64)
+        .expect("Data section start offset should be within the file contents");
+
+    // Parse the wasm file to find the globals
+    let module = walrus::Module::from_buffer(file_contents).unwrap();
+    let mut offsets = Vec::new();
+
+    // Find the main memory offset
+    let main_memory = module
+        .data
+        .iter()
+        .next()
+        .context("Failed to find main memory in WASM module")?;
+
+    let walrus::DataKind::Active {
+        offset: main_memory_offset,
+        ..
+    } = main_memory.kind
+    else {
+        tracing::error!("Failed to find main memory offset in WASM module");
+        return Ok(Vec::new());
+    };
+
+    // In the hot patch build, the main memory offset is a global from the main module and each global
+    // is it's own global. Use an offset of 0 instead if we can't evaluate the global
+    let main_memory_offset =
+        eval_walrus_global_expr(&module, &main_memory_offset).unwrap_or_default();
+
+    for export in module.exports.iter() {
+        if !looks_like_manganis_symbol(&export.name) {
+            continue;
+        }
+
+        let walrus::ExportItem::Global(global) = export.item else {
             continue;
         };
 
-        global_values = global_section
-            .into_iter()
-            .map(|global| {
-                let global = global.ok()?;
-                match global.init_expr.get_operators_reader().into_iter().next() {
-                    Some(Ok(wasmparser::Operator::I32Const { value })) => Some(value as u64),
-                    Some(Ok(wasmparser::Operator::I64Const { value })) => Some(value as u64),
-                    _ => None,
-                }
-            })
-            .collect::<Vec<_>>();
-    }
-    let mut offsets = Vec::new();
+        let walrus::GlobalKind::Local(pointer) = module.globals.get(global).kind else {
+            continue;
+        };
 
-    for (symbol, section) in manganis_symbols(file) {
-        let virtual_address = symbol.address();
-
-        let Some((_, section_range_end)) = section.file_range() else {
+        let Some(virtual_address) = eval_walrus_global_expr(&module, &pointer) else {
             tracing::error!(
-                "Found __MANGANIS__ symbol {:?} in section {}, but the section has no file range",
-                symbol.name(),
-                section.index()
+                "Found __MANGANIS__ symbol {:?} in WASM file, but the global expression could not be evaluated",
+                export.name
             );
             continue;
         };
-        let section_size = section.data()?.len() as u64;
-        let section_start = section_range_end - section_size;
-        // Translate the section_relative_address to the file offset
-        // WASM files have a section address of 0 in object, reparse the data section with wasmparser
-        // to get the correct address and section start
-        let reader = wasmparser::DataSectionReader::new(BinaryReader::new(
-            &file_contents[section_start as usize..section_range_end as usize],
-            0,
-        ))
-        .context("Failed to create WASM data section reader")?;
-        let main_memory = reader
-            .into_iter()
-            .next()
-            .context("Failed find main memory from WASM data section")?
-            .context("Failed to read main memory from WASM data section")?;
-        let main_memory_offset = match main_memory.kind {
-            wasmparser::DataKind::Active { offset_expr, .. } => {
-                match offset_expr.get_operators_reader().into_iter().next() {
-                    Some(Ok(wasmparser::Operator::I32Const { value })) => -value as i128,
-                    Some(Ok(wasmparser::Operator::I64Const { value })) => -value as i128,
-                    Some(Ok(wasmparser::Operator::GlobalGet { global_index })) => {
-                        let Some(value) =
-                            global_values.get(global_index as usize).copied().flatten()
-                        else {
-                            tracing::error!(
-                                "Found __MANGANIS__ symbol {:?} in WASM file, but the global index {} is not found",
-                                symbol.name(),
-                                global_index
-                            );
-                            continue;
-                        };
-                        value as i128
-                    }
-                    offset_expr => {
-                        tracing::error!(
-                            "Found __MANGANIS__ symbol {:?} in WASM file, but the offset expression is not a constant is is {:?}",
-                            symbol.name(),
-                            offset_expr
-                        );
-                        continue;
-                    }
-                }
-            }
-            _ => {
-                tracing::error!(
-                    "Found __MANGANIS__ symbol {:?} in WASM file, but the data section is not active",
-                    symbol.name()
-                );
-                continue;
-            }
-        };
-        // main_memory.data is a slice somewhere in file_contents. Find out the offset in the file
-        let data_start_offset = (main_memory.data.as_ptr() as u64)
-            .checked_sub(file_contents.as_ptr() as u64)
-            .expect("Data section start offset should be within the file contents");
-        let section_relative_address: u64 = ((virtual_address as i128) + main_memory_offset)
+
+        let section_relative_address: u64 = ((virtual_address as i128)
+            - main_memory_offset as i128)
             .try_into()
             .expect("Virtual address should be greater than or equal to section address");
         let file_offset = data_start_offset + section_relative_address;
+
         offsets.push(file_offset);
     }
 
@@ -304,6 +315,7 @@ pub(crate) fn extract_assets_from_file(path: impl AsRef<Path>) -> Result<AssetMa
 
         if let Some((_, bundled_asset)) = const_serialize::deserialize_const!(BundledAsset, buffer)
         {
+            tracing::debug!("Found asset at offset {offset}: {:?}", bundled_asset);
             assets.push(bundled_asset);
         } else {
             tracing::warn!("Found an asset at offset {offset} that could not be deserialized. This may be caused by a mismatch between your dioxus and dioxus-cli versions.");
@@ -317,6 +329,7 @@ pub(crate) fn extract_assets_from_file(path: impl AsRef<Path>) -> Result<AssetMa
 
     // Write back the assets to the binary file
     for (offset, asset) in offsets.into_iter().zip(&assets) {
+        tracing::debug!("Writing asset to offset {offset}: {:?}", asset);
         let new_data = ConstVec::new();
         let new_data = const_serialize::serialize_const(asset, new_data);
 
@@ -324,6 +337,9 @@ pub(crate) fn extract_assets_from_file(path: impl AsRef<Path>) -> Result<AssetMa
         // Write the modified binary data back to the file
         file.write_all(new_data.as_ref())?;
     }
+    // Ensure the file is flushed to disk
+    file.sync_all()
+        .context("Failed to sync file after writing assets")?;
 
     // If the file is a macos binary, we need to re-sign the modified binary
     if object_file.format() == object::BinaryFormat::MachO {

packages/cli/src/build/builder.rs

@@ -587,13 +587,18 @@ impl AppBuilder {
         let original = self.build.main_exe();
         let new = self.build.patch_exe(res.time_start);
         let triple = self.build.triple.clone();
-        let original_artifacts = self.artifacts.as_ref().unwrap();
         let asset_dir = self.build.asset_dir();
 
         for bundled in res.assets.assets() {
+            let original_artifacts = self.artifacts.as_mut().unwrap();
+
             if original_artifacts.assets.contains(bundled) {
                 continue;
             }
+
+            // If this is a new asset, insert it into the artifacts so we can track it when hot reloading
+            original_artifacts.assets.insert_asset(*bundled);
+
             let from = dunce::canonicalize(PathBuf::from(bundled.absolute_source_path()))?;
 
             let to = asset_dir.join(bundled.bundled_path());

packages/cli/src/build/patch.rs

@@ -597,26 +597,41 @@ fn create_wasm_jump_table(patch: &Path, cache: &HotpatchModuleCache) -> Result<J
         {
             continue;
         }
+        let name = import.name.as_str().to_string();
 
         if let Some(table_idx) = name_to_ifunc_old.get(import.name.as_str()) {
-            let name = import.name.as_str().to_string();
             new.imports.delete(env_func_import);
             convert_import_to_ifunc_call(
                 &mut new,
                 ifunc_table_initializer,
                 func_id,
                 *table_idx,
-                name,
+                name.clone(),
             );
         }
+
+        if name_is_bindgen_symbol(&name) {
+            new.imports.delete(env_func_import);
+            convert_import_to_ifunc_call(&mut new, ifunc_table_initializer, func_id, 0, name);
+        }
     }
 
     // Wire up the preserved intrinsic functions that we saved before running wasm-bindgen to the expected
     // imports from the patch.
     for import_id in wbg_funcs {
         let import = new.imports.get_mut(import_id);
+        let ImportKind::Function(func_id) = import.kind else {
+            continue;
+        };
+
         import.module = "env".into();
         import.name = format!("__saved_wbg_{}", import.name);
+
+        if name_is_bindgen_symbol(&import.name) {
+            let name = import.name.as_str().to_string();
+            new.imports.delete(import_id);
+            convert_import_to_ifunc_call(&mut new, ifunc_table_initializer, func_id, 0, name);
+        }
     }
 
     // Wipe away the unnecessary sections

packages/cli/src/build/request.rs

@@ -791,19 +791,12 @@ impl BuildRequest {
                 ctx.status_start_bundle();
 
                 self.write_executable(ctx, &artifacts.exe, &mut artifacts.assets)
-                    .await
-                    .context("Failed to write main executable")?;
-                self.write_frameworks(ctx, &artifacts.direct_rustc)
-                    .await
-                    .context("Failed to write frameworks")?;
-                self.write_assets(ctx, &artifacts.assets)
-                    .await
-                    .context("Failed to write assets")?;
+                    .await?;
+                self.write_frameworks(ctx, &artifacts.direct_rustc).await?;
+                self.write_assets(ctx, &artifacts.assets).await?;
                 self.write_metadata().await?;
                 self.optimize(ctx).await?;
-                self.assemble(ctx)
-                    .await
-                    .context("Failed to assemble app bundle")?;
+                self.assemble(ctx).await?;
 
                 tracing::debug!("Bundle created at {}", self.root_dir().display());
             }
@@ -1420,7 +1413,7 @@ impl BuildRequest {
         }
 
         // Now extract the assets from the fat binary
-        self.collect_assets(&self.patch_exe(artifacts.time_start), ctx)?;
+        artifacts.assets = self.collect_assets(&self.patch_exe(artifacts.time_start), ctx)?;
 
         // If this is a web build, reset the index.html file in case it was modified by SSG
         self.write_index_html(&artifacts.assets)

packages/cli/src/cli/serve.rs

@@ -72,6 +72,11 @@ pub(crate) struct ServeArgs {
     #[clap(long)]
     pub(crate) force_sequential: bool,
 
+    /// Exit the CLI after running into an error. This is mainly used to test hot patching internally
+    #[clap(long)]
+    #[clap(hide = true)]
+    pub(crate) exit_on_error: bool,
+
     /// Platform-specific arguments for the build
     #[clap(flatten)]
     pub(crate) platform_args: CommandWithPlatformOverrides<PlatformServeArgs>,

packages/cli/src/serve/mod.rs

@@ -1,6 +1,7 @@
 use crate::{
     styles::{GLOW_STYLE, LINK_STYLE},
-    AppBuilder, BuildId, BuildMode, BuilderUpdate, Platform, Result, ServeArgs, TraceController,
+    AppBuilder, BuildId, BuildMode, BuilderUpdate, Error, Platform, Result, ServeArgs,
+    TraceController,
 };
 
 mod ansi_buffer;
@@ -11,6 +12,7 @@ mod runner;
 mod server;
 mod update;
 
+use dioxus_dx_wire_format::BuildStage;
 pub(crate) use output::*;
 pub(crate) use runner::*;
 pub(crate) use server::*;
@@ -38,6 +40,7 @@ pub(crate) use update::*;
 ///   to a dynamic one on the fly.
 pub(crate) async fn serve_all(args: ServeArgs, tracer: &mut TraceController) -> Result<()> {
     // Load the args into a plan, resolving all tooling, build dirs, arguments, decoding the multi-target, etc
+    let exit_on_error = args.exit_on_error;
     let mut builder = AppServer::start(args).await?;
     let mut devserver = WebServer::start(&builder)?;
     let mut screen = Output::start(builder.interactive).await?;
@@ -138,12 +141,33 @@ pub(crate) async fn serve_all(args: ServeArgs, tracer: &mut TraceController) ->
 
                 // And then open the app if it's ready
                 match update {
+                    BuilderUpdate::Progress {
+                        stage: BuildStage::Failed,
+                    } => {
+                        if exit_on_error {
+                            return Err(Error::Cargo(format!(
+                                "Build failed for platform: {platform}"
+                            )));
+                        }
+                    }
+                    BuilderUpdate::Progress {
+                        stage: BuildStage::Aborted,
+                    } => {
+                        if exit_on_error {
+                            return Err(Error::Cargo(format!(
+                                "Build aborted for platform: {platform}"
+                            )));
+                        }
+                    }
                     BuilderUpdate::Progress { .. } => {}
                     BuilderUpdate::CompilerMessage { message } => {
                         screen.push_cargo_log(message);
                     }
                     BuilderUpdate::BuildFailed { err } => {
                         tracing::error!("Build failed: {}", err);
+                        if exit_on_error {
+                            return Err(err);
+                        }
                     }
                     BuilderUpdate::BuildReady { bundle } => match bundle.mode {
                         BuildMode::Thin { ref cache, .. } => {
@@ -191,12 +215,20 @@ pub(crate) async fn serve_all(args: ServeArgs, tracer: &mut TraceController) ->
                             );
                         } else {
                             tracing::error!("Application [{platform}] exited with error: {status}");
+                            if exit_on_error {
+                                return Err(Error::Runtime(format!(
+                                    "Application [{platform}] exited with error: {status}"
+                                )));
+                            }
                         }
                     }
                     BuilderUpdate::ProcessWaitFailed { err } => {
                         tracing::warn!(
                             "Failed to wait for process - maybe it's hung or being debugged?: {err}"
                         );
+                        if exit_on_error {
+                            return Err(err.into());
+                        }
                     }
                 }
             }