Add signature check to lua_pollable

This commit is contained in:
2025-06-21 06:04:50 +10:00
parent 728ee58e0d
commit 5be3f2970c
2 changed files with 65 additions and 43 deletions

View File

@@ -5,26 +5,42 @@ use crate::{
use luaify::luaify;
use std::{
fmt::Display,
marker::PhantomPinned,
mem,
pin::Pin,
ptr,
task::{Context, Poll},
};
// guardrail bytes prepended to all lua_future values in the header part to hopefully try to prevent
// misinterpreting other kinds of cdata as a lua_future if the lua user code yields a cdata that is
// not a lua_future for some odd reason.
type Signature = u64;
const SIGNATURE: Signature = Signature::from_ne_bytes(*b"\x00lb_poll");
#[repr(C)]
#[allow(non_camel_case_types)]
pub struct lua_future<F: Future<Output: ToFfi>> {
//
// SAFETY: .poll MUST be the first field. It is only to be called by the Rust async runtime to
// advance the future without knowing its type (see `lua_pollable` below).
// SAFETY: LuaJIT guarantees that cdata payloads, which are GC-managed, are never relocated
// (i.e. pinned). We can safely assume that we are pinned and poll the future inside this
// wrapper. We use this to our advantage by storing the future value directly inside the cdata
// payload instead of boxing the future and introducing indirection.
//
// https://github.com/LuaJIT/LuaJIT/issues/1167#issuecomment-1968047229
//
// .sig and .poll fields MUST come first. .poll is only to be called by the Rust async runtime
// to advance the future without knowing its type (see `lua_pollable` below).
//
// This assumes that the crate containing the async runtime and the crate containing the future
// type are ABI-compatible (compiled by the same compiler with the same target into the same binary).
// This is always the case for luby because all modules are statically linked into one binary.
// type are ABI-compatible (compiled by the same compiler with the same target into the same
// binary). This is always the case for luby because all modules are statically linked into one
// binary.
//
// .poll and .state are opaque to Lua itself.
// only .take and .drop are visible to Lua itself; other fields are opaque.
//
poll: fn(&mut Self, cx: &mut Context) -> Poll<()>,
sig: Signature,
poll: fn(Pin<&mut Self>, cx: &mut Context) -> Poll<()>,
state: State<F>,
take: unsafe extern "C" fn(&mut Self) -> <F::Output as ToFfi>::To,
drop: unsafe extern "C" fn(&mut Self),
@@ -34,15 +50,17 @@ pub struct lua_future<F: Future<Output: ToFfi>> {
#[allow(non_camel_case_types)]
pub struct lua_pollable {
//
// SAFETY: The only way to obtain a reference to a `lua_pollable` is by returning a `lua_future<T>`
// from Rust to Lua, which LuaJIT boxes into cdata and `coroutine.yield`'s back to Rust, then
// casting the yielded pointer value to `*mut lua_pollable`.
// SAFETY: The only way to obtain a reference to a `lua_pollable` is by returning a
// `lua_future<T>` from Rust to Lua, which LuaJIT boxes into cdata and `coroutine.yield`'s back
// to Rust, then casting the yielded pointer value to `*mut lua_pollable`.
//
// This is the type-erased "header" part of a `lua_future<T>` which allows the async runtime to
// poll the future without knowing its concrete type (essentially dynamic dispatch). It has the
// same layout as `lua_future<T>` without the state part.
//
poll: fn(&mut Self, cx: &mut Context) -> Poll<()>,
sig: Signature,
poll: fn(Pin<&mut Self>, cx: &mut Context) -> Poll<()>,
_phantom: PhantomPinned,
}
enum State<F: Future> {
@@ -54,6 +72,7 @@ enum State<F: Future> {
impl<F: Future<Output: ToFfi>> lua_future<F> {
pub fn new(fut: F) -> Self {
Self {
sig: SIGNATURE,
poll: Self::poll,
state: State::Pending(fut),
take: Self::take,
@@ -61,23 +80,16 @@ impl<F: Future<Output: ToFfi>> lua_future<F> {
}
}
fn poll(&mut self, cx: &mut Context) -> Poll<()> {
//
// SAFETY: LuaJIT guarantees that cdata payloads, which are GC-managed, are never
// relocated (i.e. pinned). We can safely assume that we are pinned and poll the future.
//
// We use this to our advantage by storing the future value directly inside the cdata
// payload instead of boxing the future and introducing indirection.
//
// https://github.com/LuaJIT/LuaJIT/issues/1167#issuecomment-1968047229
//
match self.state {
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> {
// SAFETY: we do not ever move the future here, so this is safe
let this = unsafe { Pin::into_inner_unchecked(self) };
match this.state {
State::Pending(ref mut fut) => match unsafe { Pin::new_unchecked(fut) }.poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(value) => Poll::Ready(self.state = State::Fulfilled(value)),
Poll::Ready(value) => Poll::Ready(this.state = State::Fulfilled(value)), // drop the future in-place
},
State::Fulfilled(_) => Poll::Ready(()),
State::Complete => unreachable!("lua_future::poll() called on completed future"),
State::Complete => unreachable!("lua_future::poll() called on a completed future"),
}
}
@@ -93,7 +105,7 @@ impl<F: Future<Output: ToFfi>> lua_future<F> {
State::Fulfilled(value) => value.convert(),
_ => unreachable!(),
},
State::Pending(_) => panic!("lua_future::take() called on pending future"),
State::Pending(_) => panic!("lua_future::take() called on a pending future"),
State::Complete => panic!("lua_future::take() called twice"),
}
}
@@ -103,12 +115,21 @@ impl<F: Future<Output: ToFfi>> lua_future<F> {
}
}
impl lua_pollable {
pub fn is_valid(&self) -> bool {
// TODO: signature check can currently read out-of-bounds if lua code for some reason yields
// a cdata of size less than 8 bytes that is not a lua_future. there is no easy way to fix
// afaik this because there is no way to find the size of a cdata payload using the C API.
self.sig == SIGNATURE
}
}
impl Future for lua_pollable {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
// SAFETY: see comment above in `lua_future::poll()`
(self.poll)(Pin::into_inner(self), cx)
assert!(self.is_valid(), "invalid lua_pollable value");
(self.poll)(self, cx)
}
}
@@ -128,7 +149,7 @@ unsafe impl<F: Future<Output: ToFfi> + 'static> Type for lua_future<F> {
unsafe impl<F: Future<Output: ToFfi> + 'static> CDef for lua_future<F> {
fn build(s: &mut CDefBuilder) {
s.field_opaque(mem::offset_of!(Self, take))
s.field_opaque(mem::offset_of!(Self, take)) // opaque .sig, .poll and .state
.field::<unsafe extern "C" fn(*mut Self) -> <F::Output as ToFfi>::To>("__take")
.field::<unsafe extern "C" fn(*mut Self)>("__drop");
}