@@ -308,14 +308,31 @@ fn declare_import(module: &mut JITModule, name: &str, n_args: usize) -> Result<F
308308 . map_err ( |e| err ( "declare import" , e) )
309309}
310310
311+ /// A compiled function plus the metadata `call` needs to invoke it safely: the
312+ /// param count, so `call` can reject an argument slice of the wrong length (the
313+ /// generated entry block reads exactly `n_params` words from `args_ptr` and does
314+ /// not bound-check against `n_args`).
315+ struct CompiledFunc {
316+ f : CompiledAbi ,
317+ n_params : usize ,
318+ }
319+
320+ /// Process-wide source of per-module identities, so a [`CompiledId`] minted by one
321+ /// [`NativeModule`] is rejected by another (it would otherwise index a different
322+ /// module's function table). Monotonic; wraparound is not a practical concern.
323+ static NEXT_MODULE_ID : std:: sync:: atomic:: AtomicU64 = std:: sync:: atomic:: AtomicU64 :: new ( 0 ) ;
324+
311325/// Owns the JIT-compiled machine code. Compiled functions live as long as the
312326/// module, so callers keep this alive and invoke by [`CompiledId`].
313327pub struct NativeModule {
314328 module : JITModule ,
315329 ctx : Context ,
316330 fbctx : FunctionBuilderContext ,
317- funcs : Vec < CompiledAbi > ,
331+ funcs : Vec < CompiledFunc > ,
318332 counter : u32 ,
333+ /// Identity stamped into every [`CompiledId`] this module mints (see
334+ /// [`NEXT_MODULE_ID`]).
335+ id : u64 ,
319336 /// Declared host-helper imports (see [`HostHelpers`]).
320337 imports : HostFuncs ,
321338}
@@ -329,9 +346,14 @@ struct HostFuncs {
329346 list_get_int : FuncId ,
330347}
331348
332- /// Handle to a function compiled into a [`NativeModule`].
349+ /// Handle to a function compiled into a [`NativeModule`]. Carries the minting
350+ /// module's identity so it can't be used against a different module (which would
351+ /// index the wrong function table).
333352#[ derive( Debug , Clone , Copy , PartialEq , Eq ) ]
334- pub struct CompiledId ( usize ) ;
353+ pub struct CompiledId {
354+ module_id : u64 ,
355+ index : usize ,
356+ }
335357
336358impl NativeModule {
337359 pub fn new ( helpers : HostHelpers ) -> Result < Self , JitError > {
@@ -374,6 +396,7 @@ impl NativeModule {
374396 fbctx : FunctionBuilderContext :: new ( ) ,
375397 funcs : Vec :: new ( ) ,
376398 counter : 0 ,
399+ id : NEXT_MODULE_ID . fetch_add ( 1 , std:: sync:: atomic:: Ordering :: Relaxed ) ,
377400 imports,
378401 } )
379402 }
@@ -421,8 +444,14 @@ impl NativeModule {
421444 // SAFETY: `code` points at the machine code we just emitted with exactly
422445 // the `CompiledAbi` signature declared above.
423446 let f: CompiledAbi = unsafe { std:: mem:: transmute :: < * const u8 , CompiledAbi > ( code) } ;
424- let handle = CompiledId ( self . funcs . len ( ) ) ;
425- self . funcs . push ( f) ;
447+ let handle = CompiledId {
448+ module_id : self . id ,
449+ index : self . funcs . len ( ) ,
450+ } ;
451+ self . funcs . push ( CompiledFunc {
452+ f,
453+ n_params : function. n_params as usize ,
454+ } ) ;
426455 Ok ( handle)
427456 }
428457
@@ -437,7 +466,19 @@ impl NativeModule {
437466 /// only `unsafe` is the indirect call through a pointer this module emitted with
438467 /// the matching ABI, with every pointer it passes derived from owned locals.
439468 pub fn call ( & self , id : CompiledId , args : & [ i64 ] ) -> Option < i64 > {
440- let f = self . funcs [ id. 0 ] ;
469+ // Reject an id from a different module and an out-of-range index: either
470+ // would invoke the wrong (or no) function. Falling back is always safe.
471+ if id. module_id != self . id {
472+ return None ;
473+ }
474+ let func = self . funcs . get ( id. index ) ?;
475+ // The generated entry block reads exactly `n_params` words from `args_ptr`
476+ // without consulting `n_args`, so an args slice shorter than `n_params`
477+ // would read out of bounds. Reject any length mismatch and fall back.
478+ if args. len ( ) != func. n_params {
479+ return None ;
480+ }
481+ let f = func. f ;
441482 let mut out: i64 = 0 ;
442483 BAIL_FLAG . with ( |bail| {
443484 bail. set ( 0 ) ;
@@ -1384,4 +1425,45 @@ mod tests {
13841425 ) )
13851426 . expect ( "well-formed heap read should validate" ) ;
13861427 }
1428+
1429+ // fn(a, b) { return a + b } — a 2-param function for the call-guard tests.
1430+ fn two_param_add ( ) -> JitFunction {
1431+ f (
1432+ 2 ,
1433+ 3 ,
1434+ vec ! [
1435+ JitInstr :: Add {
1436+ dst: 2 ,
1437+ lhs: 0 ,
1438+ rhs: 1 ,
1439+ } ,
1440+ JitInstr :: Return { src: 2 } ,
1441+ ] ,
1442+ )
1443+ }
1444+
1445+ #[ test]
1446+ fn call_rejects_wrong_arg_count ( ) {
1447+ // The generated entry block reads exactly `n_params` words from `args_ptr`,
1448+ // so a short slice must be rejected by `call` (otherwise: out-of-bounds
1449+ // read). Both too-few and too-many fall back rather than misread memory.
1450+ let mut m = module ( ) ;
1451+ let id = m. compile ( & two_param_add ( ) ) . unwrap ( ) ;
1452+ assert_eq ! ( m. call( id, & [ 2 , 3 ] ) , Some ( 5 ) ) ;
1453+ assert_eq ! ( m. call( id, & [ 2 ] ) , None ) ; // too few — must not read past the slice
1454+ assert_eq ! ( m. call( id, & [ ] ) , None ) ;
1455+ assert_eq ! ( m. call( id, & [ 2 , 3 , 4 ] ) , None ) ; // too many
1456+ }
1457+
1458+ #[ test]
1459+ fn call_rejects_id_from_another_module ( ) {
1460+ // A `CompiledId` minted by one module indexes that module's table; using it
1461+ // against another module must be rejected, not silently mis-dispatched.
1462+ let mut m1 = module ( ) ;
1463+ let mut m2 = module ( ) ;
1464+ let id1 = m1. compile ( & two_param_add ( ) ) . unwrap ( ) ;
1465+ let _id2 = m2. compile ( & two_param_add ( ) ) . unwrap ( ) ;
1466+ assert_eq ! ( m1. call( id1, & [ 2 , 3 ] ) , Some ( 5 ) ) ;
1467+ assert_eq ! ( m2. call( id1, & [ 2 , 3 ] ) , None ) ; // foreign id → fallback, no panic
1468+ }
13871469}
0 commit comments