Initial kernel support for security contexts (#167)

* Adding initial security context APIs for kernel.

* Add security context mgmt for threads.

* Support multi-security context in kernel.

* Hook up attach and upcall for sctx.

* Fix target caching.

* Cleanup aarch64 switch.

src/kernel/src/arch/aarch64/context.rs

@@ -3,28 +3,33 @@ use arm64::registers::{TTBR0_EL1, TTBR1_EL1};
 use crate::{
     arch::memory::pagetables::{Entry, EntryFlags, Table},
     memory::{
-        frame::{alloc_frame, PhysicalFrameFlags},
+        frame::{alloc_frame, free_frame, get_frame, PhysicalFrameFlags},
         pagetables::{
-            DeferredUnmappingOps, Mapper, MapReader, MappingCursor, MappingSettings,
+            DeferredUnmappingOps, MapReader, Mapper, MappingCursor, MappingSettings,
             PhysAddrProvider,
         },
         PhysAddr,
     },
     mutex::Mutex,
     spinlock::Spinlock,
+    VirtAddr,
 };
 
 // this does not need to be pub
 pub struct ArchContextInner {
     // we have a single mapper that covers one part of the address space
     mapper: Mapper,
 }
+
 pub struct ArchContext {
-    kernel: u64, // TODO: do we always need a copy?
-    user: PhysAddr,
+    pub target: ArchContextTarget,
     inner: Mutex<ArchContextInner>,
 }
 
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
+// TODO: can we get the kernel tables elsewhere?
+pub struct ArchContextTarget(PhysAddr);
+
 // default kernel mapper that is shared among all kernel instances of ArchContext
 lazy_static::lazy_static! {
     static ref KERNEL_MAPPER: Spinlock<Mapper> = {
@@ -37,17 +42,19 @@ lazy_static::lazy_static! {
         for idx in (Table::PAGE_TABLE_ENTRIES/2)..Table::PAGE_TABLE_ENTRIES {
             // write out PT entries for a top level table
             // whose entries point to another zeroed page
-            m.set_top_level_table(idx, 
+            m.set_top_level_table(idx,
                 Entry::new(
                     alloc_frame(PhysicalFrameFlags::ZEROED)
-                        .start_address(), 
+                        .start_address(),
                     // intermediate here means another page table
                     EntryFlags::intermediate()
                 )
             );
         }
         Spinlock::new(m)
     };
+
+    static ref KERNEL_TABLE_ADDR: PhysAddr = KERNEL_MAPPER.lock().root_address();
 }
 
 impl Default for ArchContext {
@@ -60,9 +67,9 @@ impl ArchContext {
     /// Construct a new context for the kernel.
     pub fn new_kernel() -> Self {
         let inner = ArchContextInner::new();
+        let target = ArchContextTarget(inner.mapper.root_address());
         Self {
-            kernel: KERNEL_MAPPER.lock().root_address().raw(),
-            user: inner.mapper.root_address(),
+            target,
             inner: Mutex::new(inner),
         }
     }
@@ -71,25 +78,34 @@ impl ArchContext {
         Self::new_kernel()
     }
 
-    #[allow(named_asm_labels)]
     pub fn switch_to(&self) {
-        // TODO: make sure the TTBR1_EL1 switch only happens once
+        unsafe {
+            Self::switch_to_target(&self.target);
+        }
+    }
+
+    #[allow(named_asm_labels)]
+    /// Switch to a target context.
+    ///
+    /// # Safety
+    /// This function must be called with a target that comes from an ArchContext that lives long enough.
+    pub unsafe fn switch_to_target(tgt: &ArchContextTarget) {
+        // TODO: If the incoming target is already the current user table, this should be a no-op. Also, we don't
+        // need to set the kernel tables each time.
         // write TTBR1
-        TTBR1_EL1.set_baddr(self.kernel);
+        TTBR1_EL1.set_baddr(KERNEL_TABLE_ADDR.raw());
         // write TTBR0
-        TTBR0_EL1.set_baddr(self.user.raw());
-        unsafe { 
-            core::arch::asm!(
-                // ensure that all previous instructions have completed
-                "isb",
-                // invalidate all tlb entries (locally)
-                "tlbi vmalle1",
-                // ensure tlb invalidation completes
-                "dsb nsh",
-                // ensure dsb instruction completes
-                "isb",
-            );
-        }
+        TTBR0_EL1.set_baddr(tgt.0.raw());
+        core::arch::asm!(
+            // ensure that all previous instructions have completed
+            "isb",
+            // invalidate all tlb entries (locally)
+            "tlbi vmalle1",
+            // ensure tlb invalidation completes
+            "dsb nsh",
+            // ensure dsb instruction completes
+            "isb",
+        );
     }
 
     pub fn map(
@@ -133,11 +149,9 @@ impl ArchContext {
 
 impl ArchContextInner {
     fn new() -> Self {
-        // we need to create a new mapper object by allocating 
+        // we need to create a new mapper object by allocating
         // some memory for the page table.
-        let mapper = Mapper::new(
-            alloc_frame(PhysicalFrameFlags::ZEROED).start_address()
-        );
+        let mapper = Mapper::new(alloc_frame(PhysicalFrameFlags::ZEROED).start_address());
         Self { mapper }
     }
 
@@ -158,3 +172,19 @@ impl ArchContextInner {
         self.mapper.unmap(cursor)
     }
 }
+
+impl Drop for ArchContextInner {
+    fn drop(&mut self) {
+        // Unmap all user memory to clear any allocated page tables.
+        self.mapper
+            .unmap(MappingCursor::new(
+                VirtAddr::start_user_memory(),
+                VirtAddr::end_user_memory() - VirtAddr::start_user_memory(),
+            ))
+            .run_all();
+        // Manually free the root.
+        if let Some(frame) = get_frame(self.mapper.root_address()) {
+            free_frame(frame);
+        }
+    }
+}

src/kernel/src/arch/amd64/context.rs

@@ -1,11 +1,12 @@
 use crate::{
     arch::memory::pagetables::{Entry, EntryFlags},
     memory::{
-        frame::{alloc_frame, PhysicalFrameFlags},
+        frame::{alloc_frame, free_frame, get_frame, PhysicalFrameFlags},
         pagetables::{
             DeferredUnmappingOps, MapReader, Mapper, MappingCursor, MappingSettings,
             PhysAddrProvider,
         },
+        VirtAddr,
     },
     mutex::Mutex,
     spinlock::Spinlock,
@@ -16,10 +17,14 @@ pub struct ArchContextInner {
 }
 
 pub struct ArchContext {
-    target: u64,
+    pub target: ArchContextTarget,
     inner: Mutex<ArchContextInner>,
 }
 
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
+#[repr(transparent)]
+pub struct ArchContextTarget(u64);
+
 lazy_static::lazy_static! {
     static ref KERNEL_MAPPER: Spinlock<Mapper> = {
         let mut m = Mapper::new(alloc_frame(PhysicalFrameFlags::ZEROED).start_address());
@@ -43,17 +48,26 @@ impl ArchContext {
 
     pub fn new() -> Self {
         let inner = ArchContextInner::new();
-        let target = inner.mapper.root_address().into();
+        let target = ArchContextTarget(inner.mapper.root_address().into());
         Self {
             target,
             inner: Mutex::new(inner),
         }
     }
 
-    #[allow(named_asm_labels)]
     pub fn switch_to(&self) {
+        unsafe { Self::switch_to_target(&self.target) }
+    }
+
+    /// Switch to a given set of page tables.
+    ///
+    /// # Safety
+    /// The specified target must be a root page table that will live as long as we are switched to it.
+    pub unsafe fn switch_to_target(tgt: &ArchContextTarget) {
         unsafe {
-            x86::controlregs::cr3_write(self.target);
+            if tgt.0 != x86::controlregs::cr3() {
+                x86::controlregs::cr3_write(tgt.0);
+            }
         }
     }
 
@@ -124,3 +138,19 @@ impl ArchContextInner {
         self.mapper.unmap(cursor)
     }
 }
+
+impl Drop for ArchContextInner {
+    fn drop(&mut self) {
+        // Unmap all user memory to clear any allocated page tables.
+        self.mapper
+            .unmap(MappingCursor::new(
+                VirtAddr::start_user_memory(),
+                VirtAddr::end_user_memory() - VirtAddr::start_user_memory(),
+            ))
+            .run_all();
+        // Manually free the root.
+        if let Some(frame) = get_frame(self.mapper.root_address()) {
+            free_frame(frame);
+        }
+    }
+}

src/kernel/src/arch/amd64/thread.rs

@@ -5,7 +5,7 @@ use core::{
 
 use twizzler_abi::{
     arch::XSAVE_LEN,
-    object::{MAX_SIZE, NULLPAGE_SIZE},
+    object::{ObjID, MAX_SIZE, NULLPAGE_SIZE},
     upcall::{
         UpcallData, UpcallFrame, UpcallHandlerFlags, UpcallInfo, UpcallTarget, UPCALL_EXIT_CODE,
     },
@@ -147,6 +147,7 @@ fn set_upcall<T: UpcallAble + Copy>(
     regs: &mut T,
     target: UpcallTarget,
     info: UpcallInfo,
+    source_ctx: ObjID,
     sup: bool,
 ) -> bool
 where
@@ -182,7 +183,7 @@ where
         } else {
             UpcallHandlerFlags::empty()
         },
-        source_ctx: 0.into(),
+        source_ctx,
     };
 
     // Step 1: determine where we are going to put the frame. If we have
@@ -298,6 +299,11 @@ impl Thread {
     /// generated to this thread after calling this function but before returning
     /// to userspace will cause the thread to immediately abort.
     pub fn restore_upcall_frame(&self, frame: &UpcallFrame) {
+        let res = self.secctx.switch_context(frame.prior_ctx);
+        if matches!(res, crate::security::SwitchResult::NotAttached) {
+            logln!("warning -- tried to restore thread to non-attached security context");
+            crate::thread::exit(UPCALL_EXIT_CODE);
+        }
         // We restore this in the syscall return code path, since
         // we know that's where we are coming from, and we actually need
         // to use the ISR return mechanism (see the syscall code).
@@ -312,17 +318,18 @@ impl Thread {
             logln!("warning -- thread aborted due to upcall generation during frame restoration");
             crate::thread::exit(UPCALL_EXIT_CODE);
         }
+        let source_ctx = self.secctx.active_id();
         match *self.arch.entry_registers.borrow() {
             Registers::None => {
                 panic!("tried to upcall to a thread that hasn't started yet");
             }
             Registers::Interrupt(int, _) => {
                 let int = unsafe { &mut *int };
-                set_upcall(int, target, info, sup);
+                set_upcall(int, target, info, source_ctx, sup);
             }
             Registers::Syscall(sys, _) => {
                 let sys = unsafe { &mut *sys };
-                set_upcall(sys, target, info, sup);
+                set_upcall(sys, target, info, source_ctx, sup);
             }
         }
     }

src/kernel/src/main.rs

@@ -40,6 +40,7 @@ mod panic;
 mod processor;
 mod queue;
 mod sched;
+pub mod security;
 mod spinlock;
 mod syscall;
 mod thread;

src/kernel/src/memory/context.rs

@@ -20,9 +20,7 @@ use crate::syscall::object::ObjectHandle;
 use self::virtmem::KernelObjectVirtHandle;
 
 impl ObjectHandle for ContextRef {
-    fn create_with_handle(_obj: ObjectRef) -> Self {
-        Arc::new(Context::new())
-    }
+    type HandleType = Context;
 }
 
 pub mod virtmem;
@@ -40,7 +38,7 @@ pub trait UserContext {
     type MappingInfo;
 
     /// Switch to this context.
-    fn switch_to(&self);
+    fn switch_to(&self, sctx: ObjID);
     /// Insert a range of an object into the context. The implementation may choose to use start and len as hints, but
     /// should keep in mind that calls to `insert_object` may be generated by faults, and so should strive to resolve
     /// the fault by correctly mapping the object as requested.
@@ -59,6 +57,7 @@ pub trait UserContext {
 }
 
 /// A struct containing information about how an object is inserted within a context.
+#[derive(Clone)]
 pub struct ObjectContextInfo {
     object: ObjectRef,
     perms: Protections,