Prevent drop on AddrSpace

Cargo.toml

@@ -94,6 +94,7 @@ multi_core = ["acpi"]
 profiling = []
 #TODO: remove when threading issues are fixed
 pti = []
+drop_panic = []
 qemu_debug = []
 serial_debug = []
 system76_ec_debug = []

src/context/memory.rs

@@ -3,6 +3,7 @@ use arrayvec::ArrayVec;
 use core::{
     cmp,
     fmt::Debug,
+    mem::ManuallyDrop,
     num::NonZeroUsize,
     ops::Bound,
     sync::atomic::{AtomicU32, Ordering},
@@ -133,6 +134,9 @@ impl AddrSpaceWrapper {
             }
         }
     }
+    pub fn into_drop(self, token: &mut CleanLockToken) {
+        self.inner.into_inner().into_drop(token);
+    }
 }
 
 #[derive(Debug)]
@@ -721,6 +725,28 @@ impl AddrSpace {
 
         Ok(selected_span.base)
     }
+
+    pub fn into_drop(self, token: &mut CleanLockToken) {
+        ManuallyDrop::new(self).inner_drop(token);
+    }
+
+    fn inner_drop(&mut self, token: &mut CleanLockToken) {
+        for mut grant in core::mem::take(&mut self.grants).into_iter() {
+            // Unpinning the grant is allowed, because pinning only occurs in UserScheme calls to
+            // prevent unmapping the mapped range twice (which would corrupt only the scheme
+            // provider), but it won't be able to double free any range after this address space
+            // has been dropped!
+            grant.info.unpin();
+
+            // TODO: Optimize away clearing the actual page tables? Since this address space is no
+            // longer arc-rwlock wrapped, it cannot be referenced `External`ly by borrowing grants,
+            // so it should suffice to iterate over PageInfos and decrement and maybe deallocate
+            // the underlying pages (and send some funmaps).
+            let res = grant.unmap(&mut self.table.utable, &mut NopFlusher);
+
+            let _ = res.unmap(token);
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -2257,23 +2283,11 @@ pub struct Table {
 
 impl Drop for AddrSpace {
     fn drop(&mut self) {
-        //TODO: DANGER: unmap requires a CleanLockToken, this cheats!
         let mut token = unsafe { CleanLockToken::new() };
-
-        for mut grant in core::mem::take(&mut self.grants).into_iter() {
-            // Unpinning the grant is allowed, because pinning only occurs in UserScheme calls to
-            // prevent unmapping the mapped range twice (which would corrupt only the scheme
-            // provider), but it won't be able to double free any range after this address space
-            // has been dropped!
-            grant.info.unpin();
-
-            // TODO: Optimize away clearing the actual page tables? Since this address space is no
-            // longer arc-rwlock wrapped, it cannot be referenced `External`ly by borrowing grants,
-            // so it should suffice to iterate over PageInfos and decrement and maybe deallocate
-            // the underlying pages (and send some funmaps).
-            let res = grant.unmap(&mut self.table.utable, &mut NopFlusher);
-
-            let _ = res.unmap(&mut token);
+        self.inner_drop(&mut token);
+        #[cfg(feature = "drop_panic")]
+        {
+            panic!("AddrSpace dropped");
         }
     }
 }

src/scheme/proc.rs

@@ -436,7 +436,7 @@ impl KernelScheme for ProcScheme {
                         arg1,
                     },
             } => {
-                let _ = try_stop_context(context, token, |context: &mut Context| {
+                let old_ctx = try_stop_context(context, token, |context: &mut Context| {
                     let regs = context.regs_mut().ok_or(Error::new(EBADFD))?;
                     regs.set_instr_pointer(new_ip);
                     regs.set_stack_pointer(new_sp);
@@ -449,6 +449,11 @@ impl KernelScheme for ProcScheme {
 
                     Ok(context.set_addr_space(Some(new)))
                 })?;
+                if let Some(old_ctx) = old_ctx {
+                    if let Some(addrspace) = Arc::into_inner(old_ctx) {
+                        addrspace.into_drop(token);
+                    }
+                }
                 let _ = ptrace::send_event(
                     crate::syscall::ptrace_event!(PTRACE_EVENT_ADDRSPACE_SWITCH, 0),
                     token,
@@ -457,7 +462,11 @@ impl KernelScheme for ProcScheme {
             Handle {
                 kind: ContextHandle::AddrSpace { addrspace } | ContextHandle::MmapMinAddr(addrspace),
                 ..
-            } => drop(addrspace),
+            } => {
+                if let Some(addrspace) = Arc::into_inner(addrspace) {
+                    addrspace.into_drop(token);
+                }
+            }
 
             Handle {
                 kind: ContextHandle::AwaitingFiletableChange { new_ft },

src/scheme/user.rs

@@ -1323,8 +1323,14 @@ impl<const READ: bool, const WRITE: bool> CaptureGuard<READ, WRITE> {
 
         Ok(())
     }
-    fn release(mut self) -> Result<()> {
-        self.release_inner()
+    pub fn release(mut self, token: &mut CleanLockToken) -> Result<()> {
+        self.release_inner()?;
+        if let Some(addrsp) = self.addrsp.take() {
+            if let Some(addrsp) = Arc::into_inner(addrsp) {
+                addrsp.into_drop(token);
+            }
+        }
+        Ok(())
     }
 }
 impl<const READ: bool, const WRITE: bool> Drop for CaptureGuard<READ, WRITE> {
@@ -1372,7 +1378,7 @@ impl KernelScheme for UserScheme {
             token,
         );
 
-        address.release()?;
+        address.release(token)?;
 
         match result? {
             Response::Regular(res, fl, _) => Ok({
@@ -1637,7 +1643,7 @@ impl KernelScheme for UserScheme {
             token,
         );
 
-        address.release()?;
+        address.release(token)?;
 
         match result? {
             Response::Regular(res, fl, _) => Ok({
@@ -1665,7 +1671,7 @@ impl KernelScheme for UserScheme {
                 token,
             )?
             .as_regular();
-        address.release()?;
+        address.release(token)?;
         result
     }
 
@@ -1698,7 +1704,7 @@ impl KernelScheme for UserScheme {
                 token,
             )?
             .as_regular();
-        address.release()?;
+        address.release(token)?;
 
         result
     }
@@ -1732,7 +1738,7 @@ impl KernelScheme for UserScheme {
                 token,
             )?
             .as_regular();
-        address.release()?;
+        address.release(token)?;
 
         result
     }
@@ -1755,7 +1761,7 @@ impl KernelScheme for UserScheme {
                 token,
             )?
             .as_regular();
-        address.release()?;
+        address.release(token)?;
         result
     }
     fn getdents(
@@ -1789,7 +1795,7 @@ impl KernelScheme for UserScheme {
                 token,
             )?
             .as_regular();
-        address.release()?;
+        address.release(token)?;
         result
     }
     fn kfstat(&self, file: usize, stat: UserSliceWo, token: &mut CleanLockToken) -> Result<()> {
@@ -1806,7 +1812,7 @@ impl KernelScheme for UserScheme {
                 token,
             )?
             .as_regular();
-        address.release()?;
+        address.release(token)?;
         result.map(|_| ())
     }
     fn kfstatvfs(&self, file: usize, stat: UserSliceWo, token: &mut CleanLockToken) -> Result<()> {
@@ -1823,7 +1829,7 @@ impl KernelScheme for UserScheme {
                 token,
             )?
             .as_regular();
-        address.release()?;
+        address.release(token)?;
         result.map(|_| ())
     }
     fn kfmap(

src/syscall/process.rs

@@ -48,7 +48,9 @@ pub fn exit_this_context(excp: Option<syscall::Exception>, token: &mut CleanLock
 
     // Files must be closed while context is valid so that messages can be passed
     close_files.force_close_all(token);
-    drop(addrspace_opt);
+    if let Some(addrspace) = addrspace_opt {
+        addrspace.into_drop(token);
+    }
     // TODO: Should status == Status::HardBlocked be handled differently?
     let owner = {
         let mut guard = context_lock.write(token.token());