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
122
pub fn wrap_wasm_str_transform<Data, Transform>(
11
122
    store: &mut Store<Arc<DataStore<Data>>>,
12
122
    transform: Transform,
13
122
) -> Func
14
122
where
15
122
    Data: Send + 'static,
16
122
    Transform: Fn(&Data, String) -> Result<String, anyhow::Error> + Send + Sync + 'static,
17
{
18
122
    Func::wrap(
19
122
        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
61
              -> i32 {
26
61
            let data = Arc::clone(caller.data());
27

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
127
122
                    if (loc as usize + len) <= mem.len() {
128
122
                        mem[loc as usize..loc as usize + len].copy_from_slice(&bytes[..len]);
129
122
                        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
122
            match result {
139
122
                Ok(len) => len,
140
                Err(e) => {
141
                    tracing::error!("Error in wrap_wasm_str_list_get: {e:?}");
142
                    -1
143
                }
144
            }
145
122
        },
146
    )
147
122
}
148

            
149
122
pub fn wrap_wasm_str_map<Data>(
150
122
    store: &mut Store<Arc<DataStore<Data>>>,
151
122
    updater: impl Fn(&mut Data, String, String) -> Result<(), anyhow::Error> + Send + Sync + 'static,
152
122
) -> Func
153
122
where
154
122
    Data: Send + 'static,
155
{
156
122
    Func::wrap(
157
122
        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
183
              -> i32 {
164
183
            let data = caller.data();
165

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

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

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

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

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

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