2 files changed, 40 insertions, 20 deletions

diff --git a/src/shared/control_regs.rs b/src/shared/control_regs.rs
index 306e52f..1a1f6a3 100644
--- a/src/shared/control_regs.rs
+++ b/src/shared/control_regs.rs
@@ -1,5 +1,5 @@
 //! Functions to read and write control registers.
-//! See Intel Vol. 3a Section 2.5, especially Figure 2.6.
+//! See Intel Vol. 3a Section 2.5, especially Figure 2-7.
 
 bitflags! {
     pub flags Cr0: usize {
@@ -19,12 +19,15 @@ bitflags! {
 
 bitflags! {
     pub flags Cr4: usize {
+        const CR4_ENABLE_PROTECTION_KEY = 1 << 22,
         const CR4_ENABLE_SMAP = 1 << 21,
         const CR4_ENABLE_SMEP = 1 << 20,
         const CR4_ENABLE_OS_XSAVE = 1 << 18,
         const CR4_ENABLE_PCID = 1 << 17,
+        const CR4_ENABLE_FSGSBASE = 1 << 16,
         const CR4_ENABLE_SMX = 1 << 14,
         const CR4_ENABLE_VMX = 1 << 13,
+        const CR4_ENABLE_UMIP = 1 << 11,
         const CR4_UNMASKED_SSE = 1 << 10,
         const CR4_ENABLE_SSE = 1 << 9,
         const CR4_ENABLE_PPMC = 1 << 8,
@@ -39,6 +42,20 @@ bitflags! {
     }
 }
 
+bitflags! {
+    pub flags Xcr0: u64 {
+        const XCR0_PKRU_STATE = 1 << 9,
+        const XCR0_HI16_ZMM_STATE = 1 << 7,
+        const XCR0_ZMM_HI256_STATE = 1 << 6,
+        const XCR0_OPMASK_STATE = 1 << 5,
+        const XCR0_BNDCSR_STATE = 1 << 4,
+        const XCR0_BNDREG_STATE = 1 << 3,
+        const XCR0_AVX_STATE = 1 << 2,
+        const XCR0_SSE_STATE = 1 << 1,
+        const XCR0_FPU_MMX_STATE = 1 << 0,
+    }
+}
+
 
 /// Read cr0
 pub unsafe fn cr0() -> Cr0 {
@@ -82,3 +99,20 @@ pub unsafe fn cr4() -> Cr4 {
 pub unsafe fn cr4_write(val: Cr4) {
     asm!("mov $0, %cr4" :: "r" (val.bits) : "memory");
 }
+
+/// Read Extended Control Register XCR0.
+/// Only supported if CR4_ENABLE_OS_XSAVE is set.
+pub unsafe fn xcr0() -> Xcr0 {
+    let high: u32;
+    let low: u32;
+    asm!("xgetbv" : "={eax}"(low), "={edx}"(high) : "{ecx}"(0));
+    Xcr0::from_bits_truncate((high as u64) << 32 | low as u64)
+}
+
+/// Write to Extended Control Register XCR0.
+/// Only supported if CR4_ENABLE_OS_XSAVE is set.
+pub unsafe fn xcr0_write(val: Xcr0) {
+    let high: u32 = (val.bits >> 32) as u32;
+    let low: u32 = val.bits as u32;
+    asm!("xsetbv" :: "{eax}"(low), "{ecx}"(0), "{edx}"(high));
+}
diff --git a/src/shared/dtables.rs b/src/shared/dtables.rs
index c5337fb..5ec2117 100644
--- a/src/shared/dtables.rs
+++ b/src/shared/dtables.rs
@@ -17,8 +17,11 @@ pub struct DescriptorTablePointer<Entry> {
 }
 
 impl<T> DescriptorTablePointer<T> {
-    fn new(slice: &[T]) -> Self {
-        let len = slice.len() * size_of::<T>();
+    pub fn new(slice: &[T]) -> Self {
+        // GDT, LDT, and IDT all expect the limit to be set to "one less".
+        // See Intel 3a, Section 3.5.1 "Segment Descriptor Tables" and
+        // Section 6.10 "Interrupt Descriptor Table (IDT)".
+        let len = slice.len() * size_of::<T>() - 1;
         assert!(len < 0x10000);
         DescriptorTablePointer {
             base: slice.as_ptr(),
@@ -27,23 +30,6 @@ impl<T> DescriptorTablePointer<T> {
     }
 }
 
-impl DescriptorTablePointer<SegmentDescriptor> {
-    pub fn new_gdtp(gdt: &[SegmentDescriptor]) -> Self {
-        let mut p = Self::new(gdt);
-        p.limit -= 1;
-        p
-    }
-    pub fn new_ldtp(ldt: &[SegmentDescriptor]) -> Self {
-        Self::new(ldt)
-    }
-}
-
-impl DescriptorTablePointer<IdtEntry> {
-    pub fn new_idtp(idt: &[IdtEntry]) -> Self {
-        Self::new(idt)
-    }
-}
-
 
 /// Load GDT table.
 pub unsafe fn lgdt(gdt: &DescriptorTablePointer<SegmentDescriptor>) {