[PATCH 2/3] rust: Add support for deriving `AsBytes` and `FromBytes`

[Matthew Maurer <mmaurer@google.com>]

This provides a derive macro for `AsBytes` and `FromBytes` for structs
only. For both, it checks the respective trait on every underlying
field. For `AsBytes`, it emits a const-time padding check that will fail
the compilation if derived on a type with padding.

Signed-off-by: Matthew Maurer <mmaurer@google.com>
---
 rust/macros/lib.rs       | 60 ++++++++++++++++++++++++++++++++++++++++++++++++
 rust/macros/transmute.rs | 58 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 118 insertions(+)

diff --git a/rust/macros/lib.rs b/rust/macros/lib.rs
index b38002151871a33f6b4efea70be2deb6ddad38e2..cec8bd3a742b3f6b3680c4286f73c8fa550a58a0 100644
--- a/rust/macros/lib.rs
+++ b/rust/macros/lib.rs
@@ -20,9 +20,11 @@
 mod kunit;
 mod module;
 mod paste;
+mod transmute;
 mod vtable;
 
 use proc_macro::TokenStream;
+use syn::{parse_macro_input, DeriveInput};
 
 /// Declares a kernel module.
 ///
@@ -475,3 +477,61 @@ pub fn paste(input: TokenStream) -> TokenStream {
 pub fn kunit_tests(attr: TokenStream, ts: TokenStream) -> TokenStream {
     kunit::kunit_tests(attr, ts)
 }
+
+/// Implements `FromBytes` for a struct.
+///
+/// It will fail compilation if the struct you are deriving on cannot be determined to implement
+/// `FromBytes` safely. It may still fail for some types which would be safe to implement
+/// `FromBytes` for, in which case you will need to write the implementation and justification
+/// yourself.
+///
+/// Main reasons your type may be rejected:
+/// * Not a `struct`
+/// * One of the fields is not `FromBytes`
+///
+/// # Examples
+///
+/// ```
+/// #[derive(FromBytes)]
+/// #[repr(C)]
+/// struct Foo {
+///   x: u32,
+///   y: u16,
+///   z: u16,
+/// }
+/// ```
+#[proc_macro_derive(FromBytes)]
+pub fn derive_from_bytes(tokens: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(tokens as DeriveInput);
+    transmute::from_bytes(input).into()
+}
+
+/// Implements `AsBytes` for a struct.
+///
+/// It will fail compilation if the struct you are deriving on cannot be determined to implement
+/// `AsBytes` safely. It may still fail for some structures which would be safe to implement
+/// `AsBytes`, in which case you will need to write the implementation and justification
+/// yourself.
+///
+/// Main reasons your type may be rejected:
+/// * Not a `struct`
+/// * One of the fields is not `AsBytes`
+/// * Your struct has generic parameters
+/// * There is padding somewhere in your struct
+///
+/// # Examples
+///
+/// ```
+/// #[derive(AsBytes)]
+/// #[repr(C)]
+/// struct Foo {
+///   x: u32,
+///   y: u16,
+///   z: u16,
+/// }
+/// ```
+#[proc_macro_derive(AsBytes)]
+pub fn derive_as_bytes(tokens: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(tokens as DeriveInput);
+    transmute::as_bytes(input).into()
+}
diff --git a/rust/macros/transmute.rs b/rust/macros/transmute.rs
new file mode 100644
index 0000000000000000000000000000000000000000..43cf36a1334f1fed23c0e777026392f987f78d8d
--- /dev/null
+++ b/rust/macros/transmute.rs
@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: GPL-2.0
+
+use proc_macro2::TokenStream;
+use syn::{parse_quote, DeriveInput, Fields, Ident, ItemConst, Path, WhereClause};
+
+fn all_fields_impl(fields: &Fields, trait_: &Path) -> WhereClause {
+    let tys = fields.iter().map(|field| &field.ty);
+    parse_quote! {
+        where #(for<'a> #tys: #trait_),*
+    }
+}
+
+fn struct_padding_check(fields: &Fields, name: &Ident) -> ItemConst {
+    let tys = fields.iter().map(|field| &field.ty);
+    parse_quote! {
+        const _: () = {
+            assert!(#(core::mem::size_of::<#tys>())+* == core::mem::size_of::<#name>());
+        };
+    }
+}
+
+pub(crate) fn as_bytes(input: DeriveInput) -> TokenStream {
+    if !input.generics.params.is_empty() {
+        return quote::quote! { compile_error!("#[derive(AsBytes)] does not support generics") };
+    }
+    let syn::Data::Struct(ref ds) = &input.data else {
+        return quote::quote! { compile_error!("#[derive(AsBytes)] only supports structs") };
+    };
+    let name = input.ident;
+    let trait_ = parse_quote! { ::kernel::transmute::AsBytes };
+    let where_clause = all_fields_impl(&ds.fields, &trait_);
+    let padding_check = struct_padding_check(&ds.fields, &name);
+    quote::quote! {
+        #padding_check
+        // SAFETY: #name has no padding and all of its fields implement `AsBytes`
+        unsafe impl #trait_ for #name #where_clause {}
+    }
+}
+
+pub(crate) fn from_bytes(input: DeriveInput) -> TokenStream {
+    let syn::Data::Struct(ref ds) = &input.data else {
+        return quote::quote! { compile_error!("#[derive(FromBytes)] only supports structs") };
+    };
+    let (impl_generics, ty_generics, base_where_clause) = input.generics.split_for_impl();
+    let name = input.ident;
+    let trait_ = parse_quote! { ::kernel::transmute::FromBytes };
+    let mut where_clause = all_fields_impl(&ds.fields, &trait_);
+    if let Some(base_clause) = base_where_clause {
+        where_clause
+            .predicates
+            .extend(base_clause.predicates.clone())
+    };
+    quote::quote! {
+        // SAFETY: All fields of #name implement `FromBytes` and it is a struct, so there is no
+        // implicit discriminator.
+        unsafe impl #impl_generics #trait_ for #name #ty_generics #where_clause {}
+    }
+}

-- 
2.52.0.305.g3fc767764a-goog