From c56cdb373268a06d196f93f0fb31871ae489b986 Mon Sep 17 00:00:00 2001
From: Joel Dice <joel.dice@fermyon.com>
Date: Wed, 1 Nov 2023 16:31:01 -0600
Subject: [PATCH] add `reset_adapter_state` export to adapter (#7444)

This is useful as the last step of component pre-initialization with
e.g. [component-init](https://github.com/dicej/component-init), in which case we
want the adapter to forget about any open handles it has (and force it to
re-request the stdio handles next time they're needed) since it will be talking
to a brand new host at runtime.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>
---
 .../src/lib.rs                                | 68 +++++++++++--------
 1 file changed, 41 insertions(+), 27 deletions(-)

diff --git a/crates/wasi-preview1-component-adapter/src/lib.rs b/crates/wasi-preview1-component-adapter/src/lib.rs
index 66d56da23589..e9f5c6f0f5fe 100644
--- a/crates/wasi-preview1-component-adapter/src/lib.rs
+++ b/crates/wasi-preview1-component-adapter/src/lib.rs
@@ -94,6 +94,14 @@ impl<T, E> TrappingUnwrap<T> for Result<T, E> {
     }
 }
 
+#[no_mangle]
+pub unsafe extern "C" fn reset_adapter_state() {
+    let state = get_state_ptr();
+    if !state.is_null() {
+        State::init(state)
+    }
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn cabi_import_realloc(
     old_ptr: *mut u8,
@@ -2385,8 +2393,8 @@ enum AllocationState {
 
 #[allow(improper_ctypes)]
 extern "C" {
-    fn get_state_ptr() -> *const State;
-    fn set_state_ptr(state: *const State);
+    fn get_state_ptr() -> *mut State;
+    fn set_state_ptr(state: *mut State);
     fn get_allocation_state() -> AllocationState;
     fn set_allocation_state(state: AllocationState);
 }
@@ -2415,7 +2423,7 @@ impl State {
     }
 
     #[cold]
-    fn new() -> &'static State {
+    fn new() -> *mut State {
         #[link(wasm_import_module = "__main_module__")]
         extern "C" {
             fn cabi_realloc(
@@ -2445,31 +2453,37 @@ impl State {
         unsafe { set_allocation_state(AllocationState::StateAllocated) };
 
         unsafe {
-            ret.write(State {
-                magic1: MAGIC,
-                magic2: MAGIC,
-                import_alloc: ImportAlloc::new(),
-                descriptors: RefCell::new(None),
-                path_buf: UnsafeCell::new(MaybeUninit::uninit()),
-                long_lived_arena: BumpArena::new(),
-                args: Cell::new(None),
-                env_vars: Cell::new(None),
-                dirent_cache: DirentCache {
-                    stream: Cell::new(None),
-                    for_fd: Cell::new(0),
-                    cookie: Cell::new(wasi::DIRCOOKIE_START),
-                    cached_dirent: Cell::new(wasi::Dirent {
-                        d_next: 0,
-                        d_ino: 0,
-                        d_type: FILETYPE_UNKNOWN,
-                        d_namlen: 0,
-                    }),
-                    path_data: UnsafeCell::new(MaybeUninit::uninit()),
-                },
-                dotdot: [UnsafeCell::new(b'.'), UnsafeCell::new(b'.')],
-            });
-            &*ret
+            Self::init(ret);
         }
+
+        ret
+    }
+
+    #[cold]
+    unsafe fn init(state: *mut State) {
+        state.write(State {
+            magic1: MAGIC,
+            magic2: MAGIC,
+            import_alloc: ImportAlloc::new(),
+            descriptors: RefCell::new(None),
+            path_buf: UnsafeCell::new(MaybeUninit::uninit()),
+            long_lived_arena: BumpArena::new(),
+            args: Cell::new(None),
+            env_vars: Cell::new(None),
+            dirent_cache: DirentCache {
+                stream: Cell::new(None),
+                for_fd: Cell::new(0),
+                cookie: Cell::new(wasi::DIRCOOKIE_START),
+                cached_dirent: Cell::new(wasi::Dirent {
+                    d_next: 0,
+                    d_ino: 0,
+                    d_type: FILETYPE_UNKNOWN,
+                    d_namlen: 0,
+                }),
+                path_data: UnsafeCell::new(MaybeUninit::uninit()),
+            },
+            dotdot: [UnsafeCell::new(b'.'), UnsafeCell::new(b'.')],
+        });
     }
 
     /// Accessor for the descriptors member that ensures it is properly initialized