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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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