From a31cc7e9f47ec29663f307bc85bd2f914522918a Mon Sep 17 00:00:00 2001
From: Pawel Wieczorkiewicz <wipawel@grsecurity.net>
Date: Mon, 13 Nov 2023 11:13:32 +0100
Subject: [PATCH] lib,irq: implement interrupts enable/disable with restore

Signed-off-by: Pawel Wieczorkiewicz <wipawel@grsecurity.net>
---
 arch/x86/apic.c                           |  4 ++--
 include/lib.h                             | 23 +++++++++++++++++++++++
 tests/test_cond_branch_mispredictions.c   |  4 ++--
 tests/test_uncond_branch_mispredictions.c |  4 ++--
 4 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/arch/x86/apic.c b/arch/x86/apic.c
index 7688eb1f..8bd9e8f9 100644
--- a/arch/x86/apic.c
+++ b/arch/x86/apic.c
@@ -179,7 +179,7 @@ void init_apic_timer(void) {
     printk("Initializing local APIC timer\n");
 
     /* Enable interrupts for calibration */
-    sti();
+    unsigned long flags = interrupts_enable_save();
 
     /* Spend 200ms calibrating the timer, 10 iterations of 20ms each */
     for (i = 0; i < CAL_ITERATIONS; ++i) {
@@ -199,7 +199,7 @@ void init_apic_timer(void) {
         min_ticks = min(min_ticks, elapsed_ticks);
     }
 
-    cli();
+    interrupts_restore(flags);
 
     /* Interrupt every min_ticks ticks */
     apic_write(APIC_TMR_DCR, APIC_TIMER_DIVIDE_BY_16);
diff --git a/include/lib.h b/include/lib.h
index fe0c5be6..0f5c47aa 100644
--- a/include/lib.h
+++ b/include/lib.h
@@ -187,10 +187,33 @@ static inline void interrupts_enable(void) {
     sti();
 }
 
+static inline unsigned long interrupts_enable_save(void) {
+    unsigned long flags = read_eflags();
+    interrupts_enable();
+    return flags;
+}
+
 static inline void interrupts_disable(void) {
     cli();
 }
 
+static inline unsigned long interrupts_disable_save(void) {
+    unsigned long flags = read_eflags();
+    interrupts_disable();
+    return flags;
+}
+
+static inline void interrupts_restore(unsigned long flags) {
+    if (flags & X86_EFLAGS_IF)
+        sti();
+    else
+        cli();
+}
+
+static inline bool interrupts_enabled(void) {
+    return read_eflags() & X86_EFLAGS_IF;
+}
+
 static inline unsigned long read_cs(void) {
     unsigned long cs;
 
diff --git a/tests/test_cond_branch_mispredictions.c b/tests/test_cond_branch_mispredictions.c
index 9f5f4d30..c9b9c15b 100644
--- a/tests/test_cond_branch_mispredictions.c
+++ b/tests/test_cond_branch_mispredictions.c
@@ -215,13 +215,13 @@ int __aligned(PAGE_SIZE) test_cond_branch_mispredictions(void *unused) {
     printk("Testing conditional branch %s BTB flushing\n",
            WITH_BTB_FLUSH ? "with" : "without");
 
-    cli();
+    unsigned long flags = interrupts_disable_save();
     test_cond_forward_branch_cl0(LOOP_ITERATIONS);
     test_cond_forward_branch_cl1(LOOP_ITERATIONS);
 
     test_cond_backward_branch_cl0(LOOP_ITERATIONS);
     test_cond_backward_branch_cl1(LOOP_ITERATIONS);
-    sti();
+    interrupts_restore(flags);
 
     return 0;
 }
diff --git a/tests/test_uncond_branch_mispredictions.c b/tests/test_uncond_branch_mispredictions.c
index d1720f85..6be92edf 100644
--- a/tests/test_uncond_branch_mispredictions.c
+++ b/tests/test_uncond_branch_mispredictions.c
@@ -231,13 +231,13 @@ int __aligned(PAGE_SIZE) test_uncond_branch_mispredictions(void *unused) {
     printk("Testing direct unconditional " STR(BRANCH) " %s BTB flushing\n",
            WITH_BTB_FLUSH ? "with" : "without");
 
-    cli();
+    unsigned long flags = interrupts_disable_save();
     test_uncond_forward_branch_cl0(LOOP_ITERATIONS);
     test_uncond_forward_branch_cl1(LOOP_ITERATIONS);
 
     test_uncond_backward_branch_cl0(LOOP_ITERATIONS);
     test_uncond_backward_branch_cl1(LOOP_ITERATIONS);
-    sti();
+    interrupts_restore(flags);
 
     return 0;
 }