1- use std:: { path:: PathBuf , process:: ExitCode } ;
1+ use std:: { collections :: HashMap , path:: PathBuf , process:: ExitCode } ;
22
33use wasmer:: {
4- Function , FunctionEnv , FunctionEnvMut , Instance , Memory , Module , RuntimeError , Store , imports,
4+ AsStoreRef , Function , FunctionEnv , FunctionEnvMut , Instance , Memory , Module , RuntimeError ,
5+ Store , imports,
56} ;
67
78const USAGE : & str = "usage: wasm-runtime MODULE.wasm (-c CODE | FILE.py)" ;
89
910struct Context {
1011 memory : Option < Memory > ,
12+ kv : HashMap < Vec < u8 > , Vec < u8 > > ,
1113}
1214
13- fn print ( mut ctx : FunctionEnvMut < Context > , ptr : u32 , len : u32 ) -> Result < i32 , RuntimeError > {
14- let ( context, store) = ctx. data_and_store_mut ( ) ;
15- let memory = context
16- . memory
17- . as_ref ( )
18- . ok_or_else ( || RuntimeError :: new ( "guest memory is not initialized" ) ) ?
19- . view ( & store) ;
15+ impl Context {
16+ fn memory ( & self ) -> Result < & Memory , RuntimeError > {
17+ self . memory
18+ . as_ref ( )
19+ . ok_or_else ( || RuntimeError :: new ( "guest memory is not initialized" ) )
20+ }
21+ }
22+
23+ fn read_memory (
24+ context : & Context ,
25+ store : & impl AsStoreRef ,
26+ ptr : u32 ,
27+ len : u32 ,
28+ ) -> Result < Vec < u8 > , RuntimeError > {
2029 let mut bytes = vec ! [ 0 ; len as usize ] ;
21- memory
30+ context
31+ . memory ( ) ?
32+ . view ( store)
2233 . read ( u64:: from ( ptr) , & mut bytes)
2334 . map_err ( |err| RuntimeError :: new ( err. to_string ( ) ) ) ?;
35+ Ok ( bytes)
36+ }
37+
38+ fn print ( mut ctx : FunctionEnvMut < Context > , ptr : u32 , len : u32 ) -> Result < i32 , RuntimeError > {
39+ let ( context, store) = ctx. data_and_store_mut ( ) ;
40+ let bytes = read_memory ( context, & store, ptr, len) ?;
2441 let message = std:: str:: from_utf8 ( & bytes) . map_err ( |err| RuntimeError :: new ( err. to_string ( ) ) ) ?;
2542 println ! ( "{message}" ) ;
2643 Ok ( 0 )
2744}
2845
46+ fn kv_put (
47+ mut ctx : FunctionEnvMut < Context > ,
48+ key_ptr : u32 ,
49+ key_len : u32 ,
50+ value_ptr : u32 ,
51+ value_len : u32 ,
52+ ) -> Result < i32 , RuntimeError > {
53+ let ( context, store) = ctx. data_and_store_mut ( ) ;
54+ let key = read_memory ( context, & store, key_ptr, key_len) ?;
55+ let value = read_memory ( context, & store, value_ptr, value_len) ?;
56+ context. kv . insert ( key, value) ;
57+ Ok ( 0 )
58+ }
59+
60+ fn kv_get (
61+ mut ctx : FunctionEnvMut < Context > ,
62+ key_ptr : u32 ,
63+ key_len : u32 ,
64+ buf_ptr : u32 ,
65+ buf_len : u32 ,
66+ ) -> Result < i32 , RuntimeError > {
67+ let ( context, store) = ctx. data_and_store_mut ( ) ;
68+ let key = read_memory ( context, & store, key_ptr, key_len) ?;
69+ let Some ( value) = context. kv . get ( & key) else {
70+ return Ok ( -1 ) ;
71+ } ;
72+ let len = i32:: try_from ( value. len ( ) ) . map_err ( |err| RuntimeError :: new ( err. to_string ( ) ) ) ?;
73+ if value. len ( ) <= buf_len as usize {
74+ context
75+ . memory ( ) ?
76+ . view ( & store)
77+ . write ( u64:: from ( buf_ptr) , value)
78+ . map_err ( |err| RuntimeError :: new ( err. to_string ( ) ) ) ?;
79+ }
80+ Ok ( len)
81+ }
82+
2983fn run ( ) -> Result < i32 , Box < dyn std:: error:: Error > > {
3084 let mut args = std:: env:: args_os ( ) . skip ( 1 ) ;
3185 let module_path = PathBuf :: from ( args. next ( ) . ok_or ( USAGE ) ?) ;
@@ -48,14 +102,22 @@ fn run() -> Result<i32, Box<dyn std::error::Error>> {
48102 let module = Module :: new ( & store, std:: fs:: read ( & module_path) ?) ?;
49103 // Fail closed if a future build accidentally acquires OS or JS imports.
50104 for import in module. imports ( ) {
51- if import. module ( ) != "env" || import. name ( ) != "print" {
105+ if import. module ( ) != "env" || ! matches ! ( import. name( ) , "print" | "kv_get" | "kv_put" ) {
52106 return Err ( format ! ( "unexpected import: {}.{}" , import. module( ) , import. name( ) ) . into ( ) ) ;
53107 }
54108 }
55- let env = FunctionEnv :: new ( & mut store, Context { memory : None } ) ;
109+ let env = FunctionEnv :: new (
110+ & mut store,
111+ Context {
112+ memory : None ,
113+ kv : HashMap :: new ( ) ,
114+ } ,
115+ ) ;
56116 let imports = imports ! {
57117 "env" => {
58118 "print" => Function :: new_typed_with_env( & mut store, & env, print) ,
119+ "kv_get" => Function :: new_typed_with_env( & mut store, & env, kv_get) ,
120+ "kv_put" => Function :: new_typed_with_env( & mut store, & env, kv_put) ,
59121 }
60122 } ;
61123 let instance = Instance :: new ( & mut store, & module, & imports) ?;
0 commit comments