Fix two semi-related savestate bugs in waterbox.

We had two instances where the act of sealing could lose snapshots that were never recovered.  It's oof.

waterbox/waterboxhost/src/memory_block/mod.rs

@@ -739,12 +739,16 @@ impl MemoryBlock {
 
 	/// release pages, assuming the range has been fully validated already
 	fn free_pages_impl(range: &mut PageRange, advise_only: bool) {
-		// we do not save the current state of unmapped pages, and if they are later remapped,
-		// the expectation is that they will start out as zero filled.  accordingly, the most
-		// sensible way to do this is to zero them now
 		unsafe {
-			let addr = range.mirror_addr();
-			addr.zero();
+			// We do not save the current state of unmapped pages, and if they are later remapped,
+			// the expectation is that they will start out as zero filled.  Accordingly, the most
+			// sensible way to do this is to zero them now.
+			// Since this will mutate the current memory, if they have no snapshot stored we must store one now.
+			for (maddr, p) in range.iter_mut_with_mirror_addr() {
+				p.maybe_snapshot(maddr.start);
+			}
+			range.mirror_addr().zero();
+
 			// simple state size optimization: we can undirty pages in this case depending on the initial state
 			#[cfg(not(feature = "no-dirty-detection"))]
 			for p in range.iter_mut() {
@@ -800,10 +804,18 @@ impl MemoryBlock {
 		}
 		self.get_stack_dirty();
 
-		for p in self.pages.iter_mut() {
+		for (maddr, p) in self.page_range().iter_mut_with_mirror_addr() {
 			if p.dirty && !p.invisible {
 				p.dirty = false;
 				p.snapshot = Snapshot::None;
+
+				// Just as we needed to precapture RWStack snapshots when allocating on Windows, we also do when sealing
+				#[cfg(windows)]
+				if p.status == PageAllocation::Allocated(Protection::RWStack) {
+					unsafe {
+						p.maybe_snapshot(maddr.start)
+					}
+				}
 			}
 		}
 

waterbox/waterboxhost/src/memory_block/tests.rs

@@ -423,6 +423,64 @@ fn test_state_invisible() -> TestResult {
 	}
 }
 
+// Snapshot loss in normal memory:
+// If memory was allocated and had non-zero values in it when originally sealed, but then was deallocated
+// without ever being written to, that deallocation has to store the old state in the snapshot.  This used
+// to work but was broken when mirror addresses were added.
+#[test]
+fn test_state_snapshot_loss_norm() -> TestResult {
+	unsafe {
+		let addr = AddressRange { start: 0x36a00000000, size: 0x10000 };
+		let mut b = MemoryBlock::new(addr);
+		b.activate();
+		let ptr = b.addr.slice_mut();
+		b.mmap_fixed(addr, Protection::RW, true)?;
+		ptr[0xfff2] = 2;
+
+		b.seal()?;
+
+		let mut state0 = Vec::new();
+		b.save_state(&mut state0)?;
+
+		b.munmap(addr)?;
+
+		b.load_state(&mut state0.as_slice())?;
+		assert_eq!(ptr[0xfff2], 2);
+
+		Ok(())
+	}
+}
+
+// Snapshot loss in stack memory (windows only):
+// On Windows, we must pregrab snapshots for all stack area when it is allocated, because we can't reliably
+// handler it.  But this also has to rehappen when we seal the memory and potentially wipe the old snapshot.
+#[test]
+fn test_state_snapshot_loss_stack() -> TestResult {
+	unsafe {
+		let addr = AddressRange { start: 0x36d00000000, size: 0x10000 };
+		let mut b = MemoryBlock::new(addr);
+		b.activate();
+		let ptr = b.addr.slice_mut();
+		b.mmap_fixed(addr, Protection::RWStack, true)?;
+		ptr[0xfff2] = 2;
+
+		b.seal()?;
+
+		let mut state0 = Vec::new();
+		b.save_state(&mut state0)?;
+		// no pages should be in the state
+		assert!(state0.len() < 0x1000);
+
+		ptr[0xfff3] = 3;
+
+		b.load_state(&mut state0.as_slice())?;
+		assert_eq!(ptr[0xfff2], 2);
+		assert_eq!(ptr[0xfff3], 0);
+
+		Ok(())
+	}
+}
+
 #[test]
 fn test_dontneed() -> TestResult {
 	unsafe {