1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

#![recursion_limit = "128"]

extern crate proc_macro;

use proc_macro::TokenStream;
use quote::{quote, TokenStreamExt};
use syn::parse::{Parse, ParseStream};
use syn::{parse_macro_input, DeriveInput, Generics, Ident, ItemFn, ItemType, LitStr};

/// Parses a type definition, extracts its identifier and generic parameters
struct TypeDefinition {
    ident: Ident,
    generics: Generics,
}

impl Parse for TypeDefinition {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        if let Ok(d) = DeriveInput::parse(input) {
            Ok(Self {
                ident: d.ident,
                generics: d.generics,
            })
        } else if let Ok(t) = ItemType::parse(input) {
            Ok(Self {
                ident: t.ident,
                generics: t.generics,
            })
        } else {
            Err(input.error("Input is not an alias, enum, struct or union definition"))
        }
    }
}

/// `unsafe_guid` attribute macro, implements the `Identify` trait for any type
/// (mostly works like a custom derive, but also supports type aliases)
#[proc_macro_attribute]
pub fn unsafe_guid(args: TokenStream, input: TokenStream) -> TokenStream {
    // Parse the arguments and input using Syn
    let guid_str = parse_macro_input!(args as LitStr).value();
    let mut result: proc_macro2::TokenStream = input.clone().into();
    let type_definition = parse_macro_input!(input as TypeDefinition);

    // We expect a canonical GUID string, such as "12345678-9abc-def0-fedc-ba9876543210"
    if guid_str.len() != 36 {
        panic!(
            "\"{}\" is not a canonical GUID string (expected 36 bytes, found {})",
            guid_str,
            guid_str.len()
        );
    }
    let mut guid_hex_iter = guid_str.split('-');
    let mut next_guid_int = |expected_num_bits: usize| -> u64 {
        let guid_hex_component = guid_hex_iter.next().unwrap();
        if guid_hex_component.len() != expected_num_bits / 4 {
            panic!(
                "GUID component \"{}\" is not a {}-bit hexadecimal string",
                guid_hex_component, expected_num_bits
            );
        }
        match u64::from_str_radix(guid_hex_component, 16) {
            Ok(number) => number,
            _ => panic!(
                "GUID component \"{}\" is not a hexadecimal number",
                guid_hex_component
            ),
        }
    };

    // The GUID string is composed of a 32-bit integer, three 16-bit ones, and a 48-bit one
    let time_low = next_guid_int(32) as u32;
    let time_mid = next_guid_int(16) as u16;
    let time_high_and_version = next_guid_int(16) as u16;
    let clock_seq_and_variant = next_guid_int(16) as u16;
    let node_64 = next_guid_int(48);

    // Convert the node ID to an array of bytes to comply with Guid::from_values expectations
    let node = [
        (node_64 >> 40) as u8,
        ((node_64 >> 32) % 0x100) as u8,
        ((node_64 >> 24) % 0x100) as u8,
        ((node_64 >> 16) % 0x100) as u8,
        ((node_64 >> 8) % 0x100) as u8,
        (node_64 % 0x100) as u8,
    ];

    // At this point, we know everything we need to implement Identify
    let ident = type_definition.ident.clone();
    let (impl_generics, ty_generics, where_clause) = type_definition.generics.split_for_impl();
    result.append_all(quote! {
        unsafe impl #impl_generics crate::Identify for #ident #ty_generics #where_clause {
            #[doc(hidden)]
            #[allow(clippy::unreadable_literal)]
            const GUID : crate::Guid = crate::Guid::from_values(
                #time_low,
                #time_mid,
                #time_high_and_version,
                #clock_seq_and_variant,
                [#(#node),*],
            );
        }
    });
    result.into()
}

/// Custom derive for the `Protocol` trait
#[proc_macro_derive(Protocol)]
pub fn derive_protocol(item: TokenStream) -> TokenStream {
    // Parse the input using Syn
    let item = parse_macro_input!(item as DeriveInput);

    // Then implement Protocol
    let ident = item.ident.clone();
    let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl();
    let result = quote! {
        // Mark this as a `Protocol` implementation
        impl #impl_generics crate::proto::Protocol for #ident #ty_generics #where_clause {}

        // Most UEFI functions expect to be called on the bootstrap processor.
        impl #impl_generics !Send for #ident #ty_generics #where_clause {}

        // Most UEFI functions do not support multithreaded access.
        impl #impl_generics !Sync for #ident #ty_generics #where_clause {}
    };
    result.into()
}

/// Custom attribute for a UEFI executable entrypoint
#[proc_macro_attribute]
pub fn entry(args: TokenStream, input: TokenStream) -> TokenStream {
    // This code is inspired by the approach in this embedded Rust crate:
    // https://github.com/rust-embedded/cortex-m-rt/blob/965bf1e3291571e7e3b34834864117dc020fb391/macros/src/lib.rs#L85

    if !args.is_empty() {
        panic!("This attribute accepts no arguments");
    }

    let f = parse_macro_input!(input as ItemFn);

    let entry_fn_ident = &f.sig.ident;

    let result = quote!(
        static _UEFI_ENTRY_POINT_TYPE_CHECK: extern "efiapi" fn(uefi::Handle, uefi::table::SystemTable<uefi::table::Boot>) -> uefi::Status = #entry_fn_ident;
        #[no_mangle]
        pub extern "efiapi" #f
    );
    result.into()
}