rust: pin-init: internal: pin_data: self-referential drop order checks

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.

Acked-by: Benno Lossin <lossin@kernel.org>
Link: https://patch.msgid.link/20261008-dev-selfref-v2-6-e280b3c8fba5@garyguo.net/
Signed-off-by: Gary Guo <gary@garyguo.net>
This commit is contained in:
Gary Guo 2026-05-01 20:50:50 +01:00
parent 29213c00cc
commit 71293ee9f3
2 changed files with 263 additions and 9 deletions

View File

@ -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,123 @@ 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 span = ty.span().resolved_at(Span::mixed_site());
let ident = f.member.as_ident();
let lt = f.borrowed.as_ref().map(|b| &b.lifetime);
quote_spanned!(span =>
let #ident: &#lt mut #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,

View File

@ -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,30 @@ 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()
&& ty.path.leading_colon.is_none()
&& ty.path.segments[0].arguments.is_none()
{
(self.0)(&ty.path.segments[0].ident);
}
syn::visit::visit_type_path(self, ty);
}
}
TypeParamVisitor(f)
}
}