From 1e59043b01b78960eb7114cdb85597647931ff40 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florian=20N=C3=BCcke?= <florian@nuecke.de>
Date: Sun, 9 Jan 2022 09:13:05 +0100
Subject: [PATCH] Fixed potential crash when computer stops immediately while
 still initializing disk from image.

---
 .../item/AbstractBlockDeviceVMDevice.java     | 24 ++++++++-----
 .../HardDriveVMDeviceWithInitialData.java     | 34 ++++++++++++++-----
 2 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/src/main/java/li/cil/oc2/common/bus/device/item/AbstractBlockDeviceVMDevice.java b/src/main/java/li/cil/oc2/common/bus/device/item/AbstractBlockDeviceVMDevice.java
index 3a6d5f38..77a0268b 100644
--- a/src/main/java/li/cil/oc2/common/bus/device/item/AbstractBlockDeviceVMDevice.java
+++ b/src/main/java/li/cil/oc2/common/bus/device/item/AbstractBlockDeviceVMDevice.java
@@ -89,19 +89,11 @@ public abstract class AbstractBlockDeviceVMDevice<TBlock extends BlockDevice, TI
 
     @Override
     public void suspend() {
-        if (device != null) {
-            try {
-                device.close();
-            } catch (final IOException e) {
-                LOGGER.error(e);
-            }
-        }
+        closeBlockDevice();
 
         if (blobHandle != null) {
             BlobStorage.close(blobHandle);
         }
-
-        device = null;
     }
 
     @Override
@@ -163,6 +155,20 @@ public abstract class AbstractBlockDeviceVMDevice<TBlock extends BlockDevice, TI
 
     protected abstract TBlock createBlockDevice() throws IOException;
 
+    protected void closeBlockDevice() {
+        if (device == null) {
+            return;
+        }
+
+        try {
+            device.close();
+        } catch (final IOException e) {
+            LOGGER.error(e);
+        }
+
+        device = null;
+    }
+
     protected void handleDataAccess() {
     }
 
diff --git a/src/main/java/li/cil/oc2/common/bus/device/item/HardDriveVMDeviceWithInitialData.java b/src/main/java/li/cil/oc2/common/bus/device/item/HardDriveVMDeviceWithInitialData.java
index 829acc83..84cd5c69 100644
--- a/src/main/java/li/cil/oc2/common/bus/device/item/HardDriveVMDeviceWithInitialData.java
+++ b/src/main/java/li/cil/oc2/common/bus/device/item/HardDriveVMDeviceWithInitialData.java
@@ -40,15 +40,7 @@ public final class HardDriveVMDeviceWithInitialData extends HardDriveVMDevice {
 
     @Subscribe
     public void handleResumedRunningEvent(final VMResumedRunningEvent event) {
-        if (copyJob != null) {
-            try {
-                copyJob.get();
-            } catch (final Throwable e) {
-                LOGGER.error(e);
-            } finally {
-                copyJob = null;
-            }
-        }
+        joinCopyJob();
     }
 
     ///////////////////////////////////////////////////////////////////
@@ -71,4 +63,28 @@ public final class HardDriveVMDeviceWithInitialData extends HardDriveVMDevice {
         }
         return device;
     }
+
+    @Override
+    protected void closeBlockDevice() {
+        // Join the copy job before releasing the device to avoid writes from thread to closed device.
+        // Since we use memory mapped memory, closing the device leads to it holding a dead pointer, meaning
+        // further access to it will hard-crash the JVM.
+        joinCopyJob();
+
+        super.closeBlockDevice();
+    }
+
+    ///////////////////////////////////////////////////////////////////
+
+    private void joinCopyJob() {
+        if (copyJob != null) {
+            try {
+                copyJob.get();
+            } catch (final Throwable e) {
+                LOGGER.error(e);
+            } finally {
+                copyJob = null;
+            }
+        }
+    }
 }