1
use anyhow::Result;
2
use std::sync::{Arc, Mutex};
3
use wasmtime::{Caller, Func, Memory, Store};
4

            
5
pub struct DataStore<Data> {
6
    pub data: Mutex<Data>,
7
    pub memory: Mutex<Option<Memory>>,
8
}
9

            
10
102
pub fn wrap_wasm_str_transform<Data, Transform>(
11
102
    store: &mut Store<Arc<DataStore<Data>>>,
12
102
    transform: Transform,
13
102
) -> Func
14
102
where
15
102
    Data: Send + 'static,
16
102
    Transform: Fn(&Data, String) -> Result<String, anyhow::Error> + Send + Sync + 'static,
17
{
18
102
    Func::wrap(
19
102
        store,
20
        move |mut caller: Caller<'_, Arc<DataStore<Data>>>,
21
              tag_ptr: i32,
22
              tag_len: i32,
23
              buf_ptr: i32,
24
              buf_len: i32|
25
51
              -> i32 {
26
51
            let data = Arc::clone(caller.data());
27

            
28
            // Define the inner logic
29
51
            let result = (|| -> Result<i32, anyhow::Error> {
30
                // Lock the WASM memory
31
51
                let memory = data
32
51
                    .memory
33
51
                    .lock()
34
51
                    .map_err(|_| anyhow::anyhow!("Failed to lock memory"))?;
35
51
                let memory = memory
36
51
                    .as_ref()
37
51
                    .ok_or_else(|| anyhow::anyhow!("Memory not initialized"))?;
38
51
                let mem = memory.data(&caller);
39

            
40
                // Validate memory bounds for the input key
41
51
                if (tag_ptr as usize + tag_len as usize) > mem.len() {
42
                    return Err(anyhow::anyhow!("Input memory bounds error"));
43
51
                }
44

            
45
                // Extract the key string from memory
46
51
                let key = String::from_utf8(
47
51
                    mem[tag_ptr as usize..tag_ptr as usize + tag_len as usize].to_vec(),
48
                )
49
51
                .map_err(|_| anyhow::anyhow!("Failed to decode input string"))?;
50

            
51
                // Lock the data and call the transform function
52
51
                let data_lock = data
53
51
                    .data
54
51
                    .lock()
55
51
                    .map_err(|_| anyhow::anyhow!("Failed to lock data"))?;
56
51
                let output = transform(&*data_lock, key)?;
57

            
58
                // Get output bytes
59
51
                let output_bytes = output.as_bytes();
60
51
                let output_len = output_bytes.len();
61

            
62
                // Validate that the output fits within the provided buffer length
63
51
                if output_len > buf_len as usize {
64
                    return Err(anyhow::anyhow!("Output exceeds buffer size"));
65
51
                }
66

            
67
                // Validate memory bounds for the output buffer
68
51
                if (buf_ptr as usize + output_len) > mem.len() {
69
                    return Err(anyhow::anyhow!("Output memory bounds error"));
70
51
                }
71

            
72
                // Write the result string into the output buffer
73
51
                memory.data_mut(&mut caller)[buf_ptr as usize..buf_ptr as usize + output_len]
74
51
                    .copy_from_slice(output_bytes);
75

            
76
51
                Ok(output_len as i32)
77
            })();
78

            
79
            // Handle the result or return an error code
80
51
            match result {
81
51
                Ok(len) => len,
82
                Err(e) => {
83
                    tracing::error!("Error in wrap_wasm_str_transform: {e:?}");
84
                    -1 // Indicate failure
85
                }
86
            }
87
51
        },
88
    )
89
102
}
90

            
91
102
pub fn wrap_wasm_str_list_get<Data, Getter>(
92
102
    store: &mut Store<Arc<DataStore<Data>>>,
93
102
    getter: Getter,
94
102
) -> Func
95
102
where
96
102
    Data: Send + 'static,
97
102
    Getter: Fn(&Data) -> Result<Vec<String>, anyhow::Error> + Send + Sync + 'static,
98
{
99
102
    Func::wrap(
100
102
        store,
101
102
        move |mut caller: Caller<'_, Arc<DataStore<Data>>>, loc: i32, max_len: i32| -> i32 {
102
102
            let data = Arc::clone(caller.data());
103

            
104
            // Try to execute the getter and handle potential errors
105
102
            let result = (|| -> Result<i32> {
106
102
                let builder = data
107
102
                    .data
108
102
                    .lock()
109
102
                    .map_err(|_| anyhow::anyhow!("Failed to lock data"))?;
110
102
                let list = getter(&builder)?; // Call getter, which may return an error
111
102
                let mut bytes = Vec::new();
112

            
113
153
                for s in list {
114
153
                    bytes.extend_from_slice(s.as_bytes());
115
153
                    bytes.push(0); // Null terminator
116
153
                }
117

            
118
102
                let len = bytes.len().min(max_len as usize);
119

            
120
102
                let memory = data
121
102
                    .memory
122
102
                    .lock()
123
102
                    .map_err(|_| anyhow::anyhow!("Failed to lock memory"))?;
124
102
                if let Some(memory) = &*memory {
125
102
                    let mem = memory.data_mut(&mut caller);
126

            
127
102
                    if (loc as usize + len) <= mem.len() {
128
102
                        mem[loc as usize..loc as usize + len].copy_from_slice(&bytes[..len]);
129
102
                        return Ok(len as i32);
130
                    }
131
                    return Err(anyhow::anyhow!("Memory bounds error"));
132
                }
133

            
134
                Err(anyhow::anyhow!("Memory not available"))
135
            })();
136

            
137
            // Return the result or indicate failure (-1)
138
102
            match result {
139
102
                Ok(len) => len,
140
                Err(e) => {
141
                    tracing::error!("Error in wrap_wasm_str_list_get: {e:?}");
142
                    -1
143
                }
144
            }
145
102
        },
146
    )
147
102
}
148

            
149
102
pub fn wrap_wasm_str_map<Data>(
150
102
    store: &mut Store<Arc<DataStore<Data>>>,
151
102
    updater: impl Fn(&mut Data, String, String) -> Result<(), anyhow::Error> + Send + Sync + 'static,
152
102
) -> Func
153
102
where
154
102
    Data: Send + 'static,
155
{
156
102
    Func::wrap(
157
102
        store,
158
        move |caller: Caller<'_, Arc<DataStore<Data>>>,
159
              key_loc: i32,
160
              key_len: i32,
161
              value_loc: i32,
162
              value_len: i32|
163
153
              -> i32 {
164
153
            let data = caller.data();
165

            
166
153
            let result = (|| -> Result<()> {
167
153
                let memory = data
168
153
                    .memory
169
153
                    .lock()
170
153
                    .map_err(|_| anyhow::anyhow!("Failed to lock memory"))?;
171
153
                if let Some(memory) = &*memory {
172
153
                    let mem = memory.data(&caller);
173

            
174
153
                    if (key_loc as usize + key_len as usize) > mem.len()
175
153
                        || (value_loc as usize + value_len as usize) > mem.len()
176
                    {
177
                        return Err(anyhow::anyhow!("Memory bounds error"));
178
153
                    }
179

            
180
153
                    let key = String::from_utf8(
181
153
                        mem[key_loc as usize..key_loc as usize + key_len as usize].to_vec(),
182
                    )?;
183
153
                    let value = String::from_utf8(
184
153
                        mem[value_loc as usize..value_loc as usize + value_len as usize].to_vec(),
185
                    )?;
186

            
187
153
                    let mut builder = data
188
153
                        .data
189
153
                        .lock()
190
153
                        .map_err(|_| anyhow::anyhow!("Failed to lock data"))?;
191
153
                    updater(&mut builder, key, value)?;
192

            
193
153
                    Ok(())
194
                } else {
195
                    Err(anyhow::anyhow!("Memory not available"))
196
                }
197
            })();
198

            
199
153
            match result {
200
153
                Ok(()) => 0,
201
                Err(e) => {
202
                    tracing::error!("Error in wrap_wasm_str_map: {e:?}");
203
                    -1
204
                }
205
            }
206
153
        },
207
    )
208
102
}