From 5be3f2970c6156225b08794ee559d27dfed645bf Mon Sep 17 00:00:00 2001 From: luaneko Date: Sat, 21 Jun 2025 06:04:50 +1000 Subject: [PATCH] Add signature check to lua_pollable --- crates/luaffi/src/future.rs | 75 ++++++++++++++++++++++++------------- crates/luajit/src/lib.rs | 33 ++++++++-------- 2 files changed, 65 insertions(+), 43 deletions(-) diff --git a/crates/luaffi/src/future.rs b/crates/luaffi/src/future.rs index 85077bc..64c36f7 100644 --- a/crates/luaffi/src/future.rs +++ b/crates/luaffi/src/future.rs @@ -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> { // - // 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, take: unsafe extern "C" fn(&mut Self) -> ::To, drop: unsafe extern "C" fn(&mut Self), @@ -34,15 +50,17 @@ pub struct lua_future> { #[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` - // 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` 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` 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` 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 { @@ -54,6 +72,7 @@ enum State { impl> lua_future { pub fn new(fut: F) -> Self { Self { + sig: SIGNATURE, poll: Self::poll, state: State::Pending(fut), take: Self::take, @@ -61,23 +80,16 @@ impl> lua_future { } } - 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> lua_future { 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> lua_future { } } +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 { - // 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 + 'static> Type for lua_future { unsafe impl + 'static> CDef for lua_future { 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:: ::To>("__take") .field::("__drop"); } diff --git a/crates/luajit/src/lib.rs b/crates/luajit/src/lib.rs index 6c8c309..ea48090 100644 --- a/crates/luajit/src/lib.rs +++ b/crates/luajit/src/lib.rs @@ -11,6 +11,7 @@ use std::{ marker::PhantomData, ops::{Deref, DerefMut}, os::raw::{c_char, c_int, c_void}, + pin::Pin, process, ptr::{self, NonNull}, rc::Rc, @@ -981,26 +982,26 @@ impl Stack { loop { match self.resume(narg)? { ResumeStatus::Ok => { - let n = self.size() - base; - break Ok(if nret == LUA_MULTRET { - n + if nret == LUA_MULTRET { + break Ok(self.size() - base); } else { - self.resize(nret); - nret - }); + self.resize(base + nret); + break Ok(nret); + } } ResumeStatus::Suspended => { + narg = 0; self.resize(1); - let value = self.slot(1); - let ptr = value.cdata::().cast_mut(); - if ptr.is_null() { - return Err(Error::InvalidType( - "cdata", - value.type_of().name(), - )); - } else { - unsafe { (&mut *ptr).await } - narg = 0; + + let ptr = self.slot(1).cdata::().cast_mut(); + if !ptr.is_null() { + // SAFETY: rust futures boxed in cdata payloads are never moved by the GC so + // we can safely make a Pin out of this pointer. see also comments in + // `luaffi::future::lua_future`. + let fut = unsafe { Pin::new_unchecked(&mut *ptr) }; + if fut.is_valid() { + fut.await; + } } } }