[PATCH 06/20] rust: pin-init: internal: pin_data: self-referential drop order checks

[Gary Guo]

Check drop order to ensure that usage of lifetime inside self-referential
struct is consistent with the order that the fields will dropped in drop
glue.

First, fields are checked according to their index to ensure that if `a`
borrows from `b`, `b` must outlive `a`. This is simple and produces a very
good diagnostic when misused.

Lifetime bounds can also be indirectly crafted with implied bounds that
make fields well-formed. For example, in this struct

    struct Foo {
        x: &'b &'a (),
        a: String,
        y: PrintOnDrop<&'b str>,
        b: String,
    }

`&'b &'a ()` will imply that `a` outlive `b`, which is inconsistent with
the actual drop order. For this case, create a `__drop_order_check`
function with field lifetimes and outlive relationship of them as generic
parameter, and ask Rust to prove that the types are well-formed inside the
generated function, to ensure that the bad implied bounds cannot happen.

The `__drop_order_check` also need to correlate lifetimes or types captured
by generics and the field lifetimes. Do this by inserting outlive bounds
when a field mentions a specific type or lifetime parameter.

Signed-off-by: Gary Guo <gary@xxxxxxxxxxx>
---
 rust/pin-init/internal/src/pin_data.rs | 212 ++++++++++++++++++++++++++++++++-
 rust/pin-init/internal/src/util.rs     |  64 +++++++++-
 2 files changed, 267 insertions(+), 9 deletions(-)

diff --git a/rust/pin-init/internal/src/pin_data.rs b/rust/pin-init/internal/src/pin_data.rs
index dc8f530e00fb..8a7a4f6d230a 100644
--- a/rust/pin-init/internal/src/pin_data.rs
+++ b/rust/pin-init/internal/src/pin_data.rs
@@ -2,8 +2,8 @@
 
 use std::collections::{BTreeMap, BTreeSet};
 
-use proc_macro2::TokenStream;
-use quote::{format_ident, quote, ToTokens};
+use proc_macro2::{Span, TokenStream};
+use quote::{format_ident, quote, quote_spanned, ToTokens};
 use syn::{
     parse::{End, Nothing, Parse},
     parse_quote, parse_quote_spanned,
@@ -11,8 +11,8 @@
     spanned::Spanned,
     visit::Visit,
     visit_mut::VisitMut,
-    Field, Fields, Generics, Ident, Index, Item, ItemStruct, Lifetime, Member, PathSegment, Type,
-    TypePath,
+    Field, Fields, GenericParam, Generics, Ident, Index, Item, ItemStruct, Lifetime, LifetimeParam,
+    Member, PathSegment, Type, TypePath,
 };
 
 use crate::{
@@ -115,14 +115,19 @@ struct FieldInfo {
     pinned: bool,
     borrowed: Option<BorrowedInfo>,
     captures: BTreeSet<Capture>,
+    generic_lt_captures: BTreeSet<Lifetime>,
+    generic_ty_captures: BTreeSet<Ident>,
 }
 
 struct StructInfo {
     args: Args,
     struct_: ItemStruct,
     fields: Vec<FieldInfo>,
+    field_idx_map: BTreeMap<Ident, usize>,
     is_tuple_struct: bool,
     self_referential: bool,
+    /// Field lifetime generics.
+    field_lts: Generics,
 }
 
 pub(crate) fn expand_with_cfg(
@@ -215,6 +220,8 @@ fn expand(
     // Collect all bound lifetimes from generics.
     let bound_lifetimes: BTreeSet<&Lifetime> =
         struct_.generics.lifetimes().map(|x| &x.lifetime).collect();
+    // Collect all type parameters from generics.
+    let type_params: BTreeSet<&Ident> = struct_.generics.type_params().map(|x| &x.ident).collect();
     // Collect all fields.
     let field_idx_map: BTreeMap<Ident, usize> = struct_
         .fields
@@ -248,6 +255,9 @@ fn expand(
             let mut captures = BTreeSet::new();
             let wildcard_variance = Variance::default();
 
+            let mut generic_lt_captures = BTreeSet::new();
+            let mut generic_ty_captures = BTreeSet::new();
+
             // Infer lifetime based on the field referenced.
             // Bound lifetimes from struct generics take priority.
             //
@@ -261,7 +271,12 @@ fn expand(
             // would not be inferred as self-referential because `'a` is already bound by the
             // struct generics.
             Lifetime::visitor(|lt| {
-                if bound_lifetimes.contains(lt) || captures.contains(lt) {
+                if bound_lifetimes.contains(lt) {
+                    generic_lt_captures.insert(lt.clone());
+                    return;
+                }
+
+                if captures.contains(lt) {
                     return;
                 }
 
@@ -284,12 +299,21 @@ fn expand(
                 implicitly_borrowed.insert(capture.lifetime.ident.clone());
             }
 
+            GenericParam::maybe_type_params_visitor(|ident| {
+                if type_params.contains(ident) {
+                    generic_ty_captures.insert(ident.clone());
+                }
+            })
+            .visit_type(&field.ty);
+
             FieldInfo {
                 field,
                 member,
                 pinned,
                 borrowed: None,
                 captures,
+                generic_lt_captures,
+                generic_ty_captures,
             }
         })
         .collect();
@@ -322,6 +346,58 @@ fn expand(
     })
     .visit_generics(&struct_.generics);
 
+    // Create a lifetime parameter for each field.
+    let borrowed_fields: Vec<_> = fields.iter().filter_map(|f| f.borrowed.as_ref()).collect();
+    let mut field_lts = Generics {
+        lt_token: None,
+        params: borrowed_fields
+            .iter()
+            .map(|borrowed| {
+                GenericParam::Lifetime(LifetimeParam {
+                    attrs: Vec::new(),
+                    lifetime: borrowed.lifetime.clone(),
+                    colon_token: None,
+                    bounds: Default::default(),
+                })
+            })
+            .collect(),
+        gt_token: None,
+        where_clause: None,
+    };
+
+    // Insert necessary bounds to make types well-formed.
+    for field in fields.iter() {
+        let Some(borrowed) = &field.borrowed else {
+            continue;
+        };
+        let field_lt = &borrowed.lifetime;
+
+        // For each borrowed field that borrows from other fields, we need to insert outlive bounds.
+        for capture in &field.captures {
+            let lt = &capture.lifetime;
+            field_lts
+                .make_where_clause()
+                .predicates
+                .push(parse_quote!(#lt: #field_lt));
+        }
+
+        // For each borrowed field that references a generic, we also need to insert their outlive
+        // bounds so they can refer to generics.
+        for lt in field.generic_lt_captures.iter() {
+            field_lts
+                .make_where_clause()
+                .predicates
+                .push(parse_quote!(#lt: #field_lt));
+        }
+
+        for ty in field.generic_ty_captures.iter() {
+            field_lts
+                .make_where_clause()
+                .predicates
+                .push(parse_quote!(#ty: #field_lt));
+        }
+    }
+
     struct_.fields = Fields::Unit;
     let info = StructInfo {
         self_referential: fields
@@ -330,7 +406,9 @@ fn expand(
         args,
         struct_,
         fields,
+        field_idx_map,
         is_tuple_struct,
+        field_lts,
     };
 
     for field in &info.fields {
@@ -356,6 +434,7 @@ fn expand(
     let struct_def = generate_struct_def(&info);
     let unpin_impl = generate_unpin_impl(&info);
     let drop_impl = generate_drop_impl(&info);
+    let drop_order_check = generate_drop_order_check(dcx, &info);
     let projections = generate_projections(&info);
     let the_pin_data = generate_the_pin_data(&info);
 
@@ -364,6 +443,7 @@ fn expand(
         // We put the rest into this const item, because it then will not be accessible to anything
         // outside.
         const _: () = {
+            #drop_order_check
             #projections
             #the_pin_data
             #unpin_impl
@@ -568,6 +648,128 @@ impl #impl_generics
     }
 }
 
+fn generate_drop_order_check(dcx: &mut DiagCtxt, info: &StructInfo) -> TokenStream {
+    let ItemStruct {
+        ident: struct_name,
+        generics,
+        ..
+    } = &info.struct_;
+
+    // If the struct is not self-referential then we can just skip.
+    if !info.self_referential {
+        return quote!();
+    }
+
+    // Make sure fields are dropped earlier than the fields that they borrow.
+    for (i, field) in info.fields.iter().enumerate() {
+        let ident = field.member.as_ident();
+        for capture in &field.captures {
+            let borrowed_field = &capture.lifetime.ident;
+
+            if let Some(&borrowed_idx) = info.field_idx_map.get(borrowed_field) {
+                if i == borrowed_idx {
+                    // We need a strict outlive relationship, in case the lifetime is needed by the
+                    // field's drop glue.
+                    dcx.error(
+                        borrowed_field,
+                        format!("field `{ident}` cannot borrow from itself"),
+                    );
+                } else if i > borrowed_idx {
+                    dcx.error(
+                        borrowed_field,
+                        format!("field `{ident}` borrows `{borrowed_field}`, but drops later"),
+                    );
+                }
+            }
+        }
+    }
+
+    // The check above is necessary, but not sufficient.
+    //
+    // Consider this case:
+    // ```
+    // struct Foo {
+    //     x: &'b &'a (),
+    //     a: String,
+    //     y: PrintOnDrop<&'b str>,
+    //     b: String,
+    // }
+    // ```
+    // we need to ensure that `b` will strictly outlive `a`.
+    //
+    // Rust needs to ensure that types are well-formed; in the above example, `&'b &'a ()` is
+    // well-formed only if `a` outlive `b`. To avoid requiring everyone from having to express this
+    // bound explicitly when declaring a struct, the `'b: 'a` bound is inferred by the Rust
+    // compiler. However this causes an issue, where now `&'a str` can be coerced to `&'b str`
+    // because compiler thinks that it shorten the lifetime. We'll be able to put a reference to `a`
+    // into `y`; but `a` drops first, so when `y` drops, it accesses `a` and causes a
+    // use-after-free!
+    //
+    // Therefore, we must ensure the types contained within the struct has their implied bound being
+    // consistent with the actual lifetime relationship. We create a `__drop_order_check` function,
+    // with known lifetime bounds as bounds on the function, and asks Rust to *prove* that the types
+    // are wellformed, given the bounds that we understand.
+
+    let generics_with_field_lt = CombinedGenerics(vec![&info.field_lts, generics]);
+
+    let (_, ty_generics, _) = generics.split_for_impl();
+    let (impl_generics_with_field_lt, _, whr_with_field_lt) =
+        generics_with_field_lt.split_for_impl();
+
+    // Prove the wellformedness of struct fields with regarding to the bounds of
+    // `__drop_order_check`.
+    //
+    // Consider this case:
+    // ```
+    // struct Foo {
+    //     x: &'b &'a (),
+    //     a: String,
+    //     y: PrintOnDrop<&'b str>,
+    //     b: String,
+    // }
+    // ```
+    // we need to ensure that `b` will strictly outlive `a`.
+    //
+    // Rust needs to ensure that types are well-formed; in the above example, `&'b &'a ()` is
+    // well-formed only if `a` outlive `b`. To avoid requiring everyone from having to express this
+    // bound explicitly when declaring a struct, the `'b: 'a` bound is inferred by the Rust
+    // compiler. However this causes an issue, where now `&'a str` can be coerced to `&'b str`
+    // because compiler thinks that it shorten the lifetime. We'll be able to put a reference to `a`
+    // into `y`; but `a` drops first, so when `y` drops, it accesses `a` and causes a
+    // use-after-free!
+    //
+    // Rust needs to *prove* the wellformedness of the type below, taking into account only the
+    // explicitly defined bounds plus the bounds implied by the lifetime-erased struct (but not
+    // the full implied bound between the field lifetimes).
+    let wf_proofs = info.fields.iter().rev().map(|f| {
+        let ty = &f.field.ty;
+        let ident = f.member.as_ident();
+        if let Some(borrowed) = &f.borrowed {
+            let lt = &borrowed.lifetime;
+            quote!(
+                let #ident: &#lt mut #ty = loop {};
+            )
+        } else {
+            quote!(
+                let #ident: #ty = loop {};
+            )
+        }
+    });
+
+    let struct_span = struct_name.span().resolved_at(Span::mixed_site());
+    quote_spanned! {struct_span =>
+        #[allow(non_snake_case, unused)]
+        fn __drop_order_check #impl_generics_with_field_lt (
+            // This must be present so the function can *assume* the implied bounds on the erased
+            // struct. For example, if the struct has `&'a T`, Rust will infer `T: 'a`; we still
+            // want to assume these bounds as they are not relevant to the field lifetimes.
+            _: &#struct_name #ty_generics,
+        ) #whr_with_field_lt {
+            #(#wf_proofs)*
+        }
+    }
+}
+
 fn generate_projections(info: &StructInfo) -> TokenStream {
     let ItemStruct {
         vis,
diff --git a/rust/pin-init/internal/src/util.rs b/rust/pin-init/internal/src/util.rs
index 67ebb333710f..3dc72e162e1e 100644
--- a/rust/pin-init/internal/src/util.rs
+++ b/rust/pin-init/internal/src/util.rs
@@ -5,7 +5,8 @@
 use proc_macro2::{Ident, TokenStream};
 use quote::{format_ident, ToTokens};
 use syn::{
-    visit::Visit, Attribute, BoundLifetimes, GenericParam, Generics, Index, Lifetime, Member, Token,
+    visit::Visit, Attribute, BoundLifetimes, GenericParam, Generics, Index, Lifetime, Member,
+    Token, TypePath,
 };
 
 use crate::DiagCtxt;
@@ -85,6 +86,7 @@ fn display_name(&self) -> String {
 pub(crate) struct CombinedGenerics<'a>(pub(crate) Vec<&'a Generics>);
 pub(crate) struct CombinedImplGenerics<'a>(&'a CombinedGenerics<'a>);
 pub(crate) struct CombinedTypeGenerics<'a>(&'a CombinedGenerics<'a>);
+pub(crate) struct CombinedWhereClauses<'a>(&'a CombinedGenerics<'a>);
 
 impl CombinedGenerics<'_> {
     pub(crate) fn split_for_impl(
@@ -92,10 +94,13 @@ pub(crate) fn split_for_impl(
     ) -> (
         CombinedImplGenerics<'_>,
         CombinedTypeGenerics<'_>,
-        // A stub type so `split_for_impl` signature matches that of `syn`'s.
-        impl Sized,
+        CombinedWhereClauses<'_>,
     ) {
-        (CombinedImplGenerics(self), CombinedTypeGenerics(self), ())
+        (
+            CombinedImplGenerics(self),
+            CombinedTypeGenerics(self),
+            CombinedWhereClauses(self),
+        )
     }
 }
 
@@ -242,6 +247,31 @@ fn to_tokens(&self, tokens: &mut TokenStream) {
     }
 }
 
+impl ToTokens for CombinedWhereClauses<'_> {
+    fn to_tokens(&self, tokens: &mut TokenStream) {
+        self.0
+             .0
+            .iter()
+            .filter_map(|x| Some(x.where_clause.as_ref()?.where_token))
+            .next_back()
+            .unwrap_or_default()
+            .to_tokens(tokens);
+
+        let comma: Token![,] = Default::default();
+
+        for generics in self.0 .0.iter() {
+            let Some(where_clause) = &generics.where_clause else {
+                continue;
+            };
+
+            where_clause.predicates.to_tokens(tokens);
+            if !where_clause.predicates.empty_or_trailing() {
+                comma.to_tokens(tokens);
+            }
+        }
+    }
+}
+
 pub(crate) trait LifetimeExt {
     /// Get a visitor that call the provided function for all unbound lifetimes.
     fn visitor<'a>(f: impl FnMut(&'a Lifetime)) -> impl Visit<'a>;
@@ -329,3 +359,29 @@ fn visit_type_bare_fn(&mut self, bare_fn: &'a syn::TypeBareFn) {
         });
     }
 }
+
+pub(crate) trait GenericParamExt {
+    fn maybe_type_params_visitor<'a>(f: impl FnMut(&'a Ident)) -> impl Visit<'a>;
+}
+
+impl GenericParamExt for GenericParam {
+    fn maybe_type_params_visitor<'a>(f: impl FnMut(&'a Ident)) -> impl Visit<'a> {
+        struct TypeParamVisitor<F>(F);
+
+        impl<'a, F> Visit<'a> for TypeParamVisitor<F>
+        where
+            F: FnMut(&'a Ident),
+        {
+            fn visit_type_path(&mut self, ty: &'a TypePath) {
+                if ty.qself.is_none() {
+                    if let Some(ident) = ty.path.get_ident() {
+                        (self.0)(ident);
+                    }
+                }
+                syn::visit::visit_type_path(self, ty);
+            }
+        }
+
+        TypeParamVisitor(f)
+    }
+}

-- 
2.54.0