3 files changed, 40 insertions, 0 deletions
diff --git a/src/arch/x86_64/interrupt/mod.rs b/src/arch/x86_64/interrupt/mod.rs
index ed08223..a3d9f89 100644
--- a/src/arch/x86_64/interrupt/mod.rs
+++ b/src/arch/x86_64/interrupt/mod.rs
@@ -1,6 +1,7 @@
 pub mod pic_8259;
 pub mod pit;
 use crate::arch::x86_64::arch_regs::TrapFrame;
+use crate::arch::x86_64::is_int_enabled;
 use crate::arch::x86_64::paging::fault;
 use crate::io::*;
 use core::arch::asm;
@@ -103,6 +104,27 @@ pub fn interrupt_disable() {
 	}
 }
 
+#[inline]
+/// irq_save() disables all interrupts and returns the previous state
+pub fn irq_save() -> bool {
+	if is_int_enabled() {
+		interrupt_disable();
+		return true;
+	} else {
+		return false;
+	}
+}
+
+#[inline]
+/// irq_restore only re-enable irq if was_enabled==true.
+/// it will not disable irq regardless the was_enabled value. This function
+/// should only be called to restore irq based on previous irq_save();
+pub fn irq_restore(was_enabled: bool) {
+	if was_enabled {
+		interrupt_enable();
+	}
+}
+
 #[inline(always)]
 fn _idt_init() {
 	println!("[init] idt: vectors_start: 0x{:x}", vectors_start as usize);
diff --git a/src/arch/x86_64/mod.rs b/src/arch/x86_64/mod.rs
index ee51338..38b08b1 100644
--- a/src/arch/x86_64/mod.rs
+++ b/src/arch/x86_64/mod.rs
@@ -4,3 +4,19 @@ pub mod io_port;
 pub mod mem;
 pub mod misc;
 pub mod paging;
+use core::arch::asm;
+
+pub const RFLAGS_IF_MASK: u64 = 1 << 9;
+#[inline]
+pub fn read_rflags() -> u64 {
+	let rflags;
+	unsafe {
+		asm!("pushfq; popq {}", out(reg) rflags);
+	}
+	rflags
+}
+
+pub fn is_int_enabled() -> bool {
+	let rf = read_rflags();
+	return (rf & RFLAGS_IF_MASK) != 0;
+}
diff --git a/src/proc/sched.rs b/src/proc/sched.rs
index 613336e..701f165 100644
--- a/src/proc/sched.rs
+++ b/src/proc/sched.rs
@@ -39,6 +39,8 @@ impl Scheduler {
 
 	// pop front, push back
 	pub fn do_schedule() {
+		// TODO: remove this spinlock, because we should protect the scheduler
+		// with irq_save/restore
 		if SCHEDULER.is_locked() {
 			panic!("scheduler lock has been taken, something wrong");
 		}