Support setting a default handler in the IDT through a macro for automatic generation of wrapper functions

Cargo.toml

@@ -25,22 +25,26 @@ repository = "https://github.com/rust-osdev/x86_64"
 version = "0.12.2"
 edition = "2018"
 
+[workspace]
+
 [dependencies]
 bit_field = "0.9.0"
 bitflags = "1.0.4"
+x86_64-idt-general-handler = { path = "idt-general-handler", optional = true}
 
 [build-dependencies]
 cc = { version = "1.0.37", optional = true }
 
 [features]
-default = [ "nightly", "instructions" ]
+default = [ "nightly", "instructions", "proc_macros" ]
 instructions = []
 deny-warnings = []
 external_asm = [ "cc" ]
 nightly = [ "inline_asm", "const_fn", "abi_x86_interrupt" ]
 inline_asm = []
 abi_x86_interrupt = []
 const_fn = []
+proc_macros = [ "x86_64-idt-general-handler" ]
 
 [package.metadata.release]
 no-dev-version = true

idt-general-handler/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "x86_64-idt-general-handler"
+version = "0.1.0"
+authors = ["Philipp Oppermann <[email protected]>"]
+edition = "2018"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[lib]
+proc-macro = true
+
+[dependencies]
+syn = { version = "1.0.33", features = ["full"] }
+quote = "1.0.7"
+proc-macro2 = "1.0.18"

idt-general-handler/src/lib.rs

@@ -0,0 +1,186 @@
+use proc_macro::TokenStream;
+use quote::{quote, quote_spanned};
+use std::{fmt::Display, ops::RangeInclusive};
+use syn::{
+    parse::Parser, punctuated::Punctuated, spanned::Spanned, token::Comma, Expr, ExprLit, ExprPath,
+    ExprRange, ExprReference, Lit, RangeLimits,
+};
+
+extern crate proc_macro;
+
+const RESERVED_HANDLERS: &[u8] = &[15, 21, 22, 23, 24, 25, 26, 27, 28, 29, 31];
+const HANDLERS_WITH_ERR_CODE: &[u8] = &[8, 10, 11, 12, 13, 14, 17, 30];
+
+#[proc_macro]
+pub fn set_general_handler(item: TokenStream) -> TokenStream {
+    set_general_handler_impl(item).unwrap_or_else(|err| err.to_compile_error().into())
+}
+
+fn set_general_handler_impl(item: TokenStream) -> syn::Result<TokenStream> {
+    let input = Punctuated::<Expr, Comma>::parse_separated_nonempty.parse(item)?;
+    let (idt, function, range) = parse_input(&input)?;
+
+    let mut handlers = Vec::new();
+    for index in range {
+        if RESERVED_HANDLERS.contains(&index) {
+            continue; // skip reserved handlers
+        }
+        let function_call = if HANDLERS_WITH_ERR_CODE.contains(&index) {
+            quote_spanned!(function.span() => {
+                #function(stack_frame, #index.into(), Some(error_code));
+            })
+        } else {
+            quote_spanned!(function.span() => {
+                #function(stack_frame, #index.into(), None);
+            })
+        };
+
+        let stack_frame_arg =
+            quote! { stack_frame: &mut x86_64::structures::idt::InterruptStackFrame };
+
+        let handler_with_err_code = quote! {
+            extern "x86-interrupt" fn handler(#stack_frame_arg, error_code: u64) {
+                #function_call
+            }
+        };
+
+        let set_handler_fn = match index {
+            8 => quote! {
+                extern "x86-interrupt" fn handler(#stack_frame_arg, error_code: u64) -> ! {
+                    #function_call
+                    panic!("General handler returned on double fault");
+                }
+                (#idt).double_fault.set_handler_fn(handler);
+            },
+            10 => quote! {
+                #handler_with_err_code
+                (#idt).invalid_tss.set_handler_fn(handler);
+            },
+            11 => quote! {
+                #handler_with_err_code
+                (#idt).segment_not_present.set_handler_fn(handler);
+            },
+            12 => quote! {
+                #handler_with_err_code
+                (#idt).stack_segment_fault.set_handler_fn(handler);
+            },
+            13 => quote! {
+                #handler_with_err_code
+                (#idt).general_protection_fault.set_handler_fn(handler);
+            },
+            14 => quote! {
+                extern "x86-interrupt" fn handler(#stack_frame_arg, error_code: x86_64::structures::idt::PageFaultErrorCode) {
+                    let error_code = error_code.bits();
+                    #function_call
+                }
+                (#idt).page_fault.set_handler_fn(handler);
+            },
+            17 => quote! {
+                #handler_with_err_code
+                (#idt).alignment_check.set_handler_fn(handler);
+            },
+            18 => quote! {
+                extern "x86-interrupt" fn handler(#stack_frame_arg) -> ! {
+                    #function_call
+                    panic!("General handler returned on machine check exception");
+                }
+                (#idt).machine_check.set_handler_fn(handler);
+            },
+            30 => quote! {
+                #handler_with_err_code
+                (#idt).security_exception.set_handler_fn(handler);
+            },
+            index => quote! {
+                extern "x86-interrupt" fn handler(#stack_frame_arg) {
+                    #function_call
+                }
+                (#idt)[#index.into()].set_handler_fn(handler);
+            },
+        };
+
+        // double `{{` to create new scope
+        handlers.push(quote! {{
+            #set_handler_fn
+        }});
+    }
+
+    let ret = quote! {
+        #(#handlers)*
+    };
+
+    // uncomment to print generated code:
+    // println!("{}", ret);
+
+    Ok(ret.into())
+}
+
+fn parse_input(
+    input: &Punctuated<Expr, Comma>,
+) -> syn::Result<(&ExprReference, &ExprPath, RangeInclusive<u8>)> {
+    if input.len() < 2 || input.len() > 3 {
+        return Err(err(input, "expected 2 or 3 arguments"));
+    }
+
+    let idt = match &input[0] {
+        Expr::Reference(r) if r.mutability.is_some() => r,
+        other => return Err(err(other, "expected a `&mut` reference to an IDT")),
+    };
+
+    let function = match &input[1] {
+        Expr::Path(path) => path,
+        other => return Err(err(other, "expected the name of a handler function")),
+    };
+
+    let range = if input.len() == 3 {
+        match &input[2] {
+            Expr::Lit(lit) => match &lit.lit {
+                Lit::Int(int) => {
+                    let index: u8 = int.base10_parse()?;
+                    index..=index
+                }
+                other => return Err(err(other, "expected index or range")),
+            },
+            Expr::Range(range) => parse_range(&range)?,
+            other => return Err(err(other, "expected index or range")),
+        }
+    } else {
+        0..=255
+    };
+
+    Ok((idt, function, range))
+}
+
+fn parse_range(range: &ExprRange) -> syn::Result<RangeInclusive<u8>> {
+    let range_start: u8 = match range.from.as_deref() {
+        Some(&Expr::Lit(ExprLit {
+            lit: Lit::Int(ref int),
+            ..
+        })) => int.base10_parse()?,
+        Some(other) => return Err(err(other, "Invalid range start")),
+        None => 0,
+    };
+    let mut range_end: u8 = match range.to.as_deref() {
+        Some(&Expr::Lit(ExprLit {
+            lit: Lit::Int(ref int),
+            ..
+        })) => int.base10_parse()?,
+        Some(other) => return Err(err(other, "Invalid range end")),
+        None => 255,
+    };
+    if let ExprRange {
+        limits: RangeLimits::HalfOpen(_),
+        to: Some(_),
+        ..
+    } = range
+    {
+        range_end = range_end
+            .checked_sub(1)
+            .ok_or_else(|| err(&range.to, "Invalid range"))?;
+    };
+
+    Ok(range_start..=range_end)
+}
+
+fn err(tokens: impl quote::ToTokens, message: impl Display) -> syn::Error {
+    syn::Error::new_spanned(tokens, message)
+}

src/structures/idt.rs

@@ -16,6 +16,8 @@ use core::fmt;
 use core::marker::PhantomData;
 use core::ops::Bound::{Excluded, Included, Unbounded};
 use core::ops::{Deref, Index, IndexMut, RangeBounds};
+#[cfg(feature = "proc_macros")]
+pub use x86_64_idt_general_handler::set_general_handler;
 
 /// An Interrupt Descriptor Table with 256 entries.
 ///
@@ -822,11 +824,77 @@ bitflags! {
 #[cfg(test)]
 mod test {
     use super::*;
+    use crate as x86_64;
+
+    fn entry_present(idt: &InterruptDescriptorTable, index: usize) -> bool {
+        let options = match index {
+            8 => &idt.double_fault.options,
+            10 => &idt.invalid_tss.options,
+            11 => &idt.segment_not_present.options,
+            12 => &idt.stack_segment_fault.options,
+            13 => &idt.general_protection_fault.options,
+            14 => &idt.page_fault.options,
+            15 => &idt.reserved_1.options,
+            17 => &idt.alignment_check.options,
+            18 => &idt.machine_check.options,
+            i @ 21..=29 => &idt.reserved_2[i - 21].options,
+            30 => &idt.security_exception.options,
+            31 => &idt.reserved_3.options,
+            other => &idt[other].options,
+        };
+        options.0.get_bit(15)
+    }
 
     #[test]
     fn size_test() {
         use core::mem::size_of;
         assert_eq!(size_of::<Entry<HandlerFunc>>(), 16);
         assert_eq!(size_of::<InterruptDescriptorTable>(), 256 * 16);
     }
+
+    #[cfg(feature = "proc_macros")]
+    #[test]
+    fn general_handlers() {
+        fn general_handler(
+            _stack_frame: &mut InterruptStackFrame,
+            _index: u8,
+            _error_code: Option<u64>,
+        ) {
+        }
+
+        let mut idt = InterruptDescriptorTable::new();
+        set_general_handler!(&mut idt, general_handler, 0);
+        for i in 0..256 {
+            if i == 0 {
+                assert!(entry_present(&idt, i));
+            } else {
+                assert!(!entry_present(&idt, i));
+            }
+        }
+        set_general_handler!(&mut idt, general_handler, 14);
+        for i in 0..256 {
+            if i == 0 || i == 14 {
+                assert!(entry_present(&idt, i));
+            } else {
+                assert!(!entry_present(&idt, i));
+            }
+        }
+        set_general_handler!(&mut idt, general_handler, 32..64);
+        for i in 1..256 {
+            if i == 0 || i == 14 || (i >= 32 && i < 64) {
+                assert!(entry_present(&idt, i));
+            } else {
+                assert!(!entry_present(&idt, i));
+            }
+        }
+        set_general_handler!(&mut idt, general_handler);
+        for i in 0..256 {
+            if i == 15 || i == 31 || (i >= 21 && i <= 29) {
+                // reserved entries should not be set
+                assert!(!entry_present(&idt, i));
+            } else {
+                assert!(entry_present(&idt, i));
+            }
+        }
+    }
 }