Grcov report - wrappers.rs

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

// Helper function for builder operations, allowing any type

11

16

fn with_builder<F, Builder>(caller: &Caller<'_, Arc<DataStore<Builder>>>, action: F) -> Result<()>

12

16

where

13

16

    F: FnOnce(&mut Builder) -> Result<()>,

14

15

16

    let data = caller.data();

16

17

    // Lock the builder

18

16

    let mut builder = match data.data.lock() {

19

16

        Ok(builder) => builder,

20

        Err(_) => anyhow::bail!(t!("Can't lock the builder")),

21

};

22

23

    // Execute the action on the builder

24

16

    action(&mut builder)

25

16

26

27

// Helper function for setting a str from memory with two numbers

28

fn with_data_set_str<F, Data>(

29

    caller: &Caller<'_, Arc<DataStore<Data>>>,

30

    loc: i32,

31

    len: i32,

32

    set_field: F,

33

) -> Result<()>

34

where

35

    F: FnOnce(&mut Data, String),

36

37

    with_builder(caller, |builder| {

38

        let data = caller.data();

39

40

        // Lock the memory

41

        let m = match data.memory.lock() {

42

            Ok(m) => m,

43

            Err(_) => anyhow::bail!(t!("Can't lock the memory")),

44

};

45

46

        // Access memory and set the field if data is found

47

        if let Some(memory) = &*m

48

            && let Some(bytes) = memory

49

                .data(caller)

50

                .get(loc as usize..loc as usize + len as usize)

51

52

            let string_value = String::from_utf8_lossy(bytes).to_string();

53

            set_field(builder, string_value);

54

55

56

        Ok(())

57

})

58

59

60

// Wrapper for WASM numbers

61

16

pub fn wrap_wasm_i64<Data, Setter>(store: &mut Store<Arc<DataStore<Data>>>, setter: Setter) -> Func

62

16

where

63

16

    Data: Send + 'static,

64

16

    Setter: Fn(&mut Data, i64) + Send + Sync + 'static,

65

66

16

    Func::wrap(

67

16

        store,

68

16

        move |caller: Caller<'_, Arc<DataStore<Data>>>, value: i64| {

69

16

            with_builder(&caller, |builder| {

70

16

                setter(builder, value);

71

16

                Ok(())

72

16

})

73

16

},

74

75

16

76

77

// Generic wrapper for creating `Func` to set string fields in the builder

78

pub fn wrap_wasm_str<Data, Setter>(store: &mut Store<Arc<DataStore<Data>>>, setter: Setter) -> Func

79

where

80

    Data: Send + 'static,

81

    Setter: Fn(&mut Data, String) + Send + Sync + 'static,

82

83

    Func::wrap(

84

        store,

85

        move |caller: Caller<'_, Arc<DataStore<Data>>>, loc: i32, len: i32| {

86

            with_data_set_str(&caller, loc, len, |builder, value| {

87

                setter(builder, value);

88

})

89

},

90

91

92

93

// Wrapper for WASM i64 getter

94

16

pub fn wrap_wasm_i64_get<Data, Getter>(

95

16

    store: &mut Store<Arc<DataStore<Data>>>,

96

16

    getter: Getter,

97

16

) -> Func

98

16

where

99

16

    Data: Send + 'static,

100

16

    Getter: Fn(&Data) -> i64 + Send + Sync + 'static,

101

102

16

    Func::wrap(

103

16

        store,

104

16

        move |caller: Caller<'_, Arc<DataStore<Data>>>| -> i64 {

105

16

            let data = caller.data();

106

16

            let builder = data.data.lock().expect("Failed to lock data");

107

108

            // Call the getter to retrieve the i64 value

109

16

            getter(&builder)

110

16

},

111

112

16

113

114

// Wrapper for WASM string getter that writes the string to memory and returns its actual length

115

pub fn wrap_wasm_str_get<Data, Getter>(

116

    store: &mut Store<Arc<DataStore<Data>>>,

117

    getter: Getter,

118

) -> Func

119

where

120

    Data: Send + 'static,

121

    Getter: Fn(&Data) -> String + Send + Sync + 'static,

122

123

    Func::wrap(

124

        store,

125

        move |mut caller: Caller<'_, Arc<DataStore<Data>>>, loc: i32, max_len: i32| -> i32 {

126

            // Clone the Arc to avoid holding an immutable borrow on `caller`

127

            let data = Arc::clone(caller.data());

128

129

            // Extract the value from the builder in a separate scope

130

            let value = {

131

                let builder = data.data.lock().expect("Failed to lock data");

132

                getter(&builder)

133

            }; // `builder` is dropped here, ending the immutable borrow

134

135

            let bytes = value.as_bytes();

136

137

            // Lock the memory

138

            let m = data.memory.lock().expect("Failed to lock memory");

139

140

            if let Some(memory) = &*m {

141

                // Obtain a mutable context from the caller

142

                let mem = memory.data_mut(&mut caller);

143

144

                // Calculate the length to copy, limited by `max_len`

145

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

146

147

                // Ensure the memory bounds are valid

148

                if (loc as usize + len) <= mem.len() {

149

                    // Copy bytes to the specified location in WASM memory

150

                    mem[loc as usize..loc as usize + len].copy_from_slice(&bytes[..len]);

151

152

                    // Return the actual length of the copied string

153

                    return len as i32;

154

                } else {

155

                    // Memory bounds are invalid; indicate failure

156

                    return -1;

157

158

159

160

            // Return -1 to indicate failure if memory is unavailable

161

-1

162

},

163

164

165

166

16

pub fn wrap_wasm_str_transform<Data, Transform>(

167

16

    store: &mut Store<Arc<DataStore<Data>>>,

168

16

    transform: Transform,

169

16

) -> Func

170

16

where

171

16

    Data: Send + 'static,

172

16

    Transform: Fn(&Data, String) -> Result<String, anyhow::Error> + Send + Sync + 'static,

173

174

16

    Func::wrap(

175

16

        store,

176

        move |mut caller: Caller<'_, Arc<DataStore<Data>>>,

177

              tag_ptr: i32,

178

              tag_len: i32,

179

              buf_ptr: i32,

180

              buf_len: i32|

181

8

              -> i32 {

182

8

            let data = Arc::clone(caller.data());

183

184

            // Define the inner logic

185

8

            let result = (|| -> Result<i32, anyhow::Error> {

186

                // Lock the WASM memory

187

8

                let memory = data

188

8

                    .memory

189

8

                    .lock()

190

8

                    .map_err(|_| anyhow::anyhow!("Failed to lock memory"))?;

191

8

                let memory = memory

192

8

                    .as_ref()

193

8

                    .ok_or_else(|| anyhow::anyhow!("Memory not initialized"))?;

194

8

                let mem = memory.data(&caller);

195

196

                // Validate memory bounds for the input key

197

8

                if (tag_ptr as usize + tag_len as usize) > mem.len() {

198

                    return Err(anyhow::anyhow!("Input memory bounds error"));

199

8

200

201

                // Extract the key string from memory

202

8

                let key = String::from_utf8(

203

8

                    mem[tag_ptr as usize..tag_ptr as usize + tag_len as usize].to_vec(),

204

205

8

                .map_err(|_| anyhow::anyhow!("Failed to decode input string"))?;

206

207

                // Lock the data and call the transform function

208

8

                let data_lock = data

209

8

                    .data

210

8

                    .lock()

211

8

                    .map_err(|_| anyhow::anyhow!("Failed to lock data"))?;

212

8

                let output = transform(&*data_lock, key)?;

213

214

                // Get output bytes

215

8

                let output_bytes = output.as_bytes();

216

8

                let output_len = output_bytes.len();

217

218

                // Validate that the output fits within the provided buffer length

219

8

                if output_len > buf_len as usize {

220

                    return Err(anyhow::anyhow!("Output exceeds buffer size"));

221

8

222

223

                // Validate memory bounds for the output buffer

224

8

                if (buf_ptr as usize + output_len) > mem.len() {

225

                    return Err(anyhow::anyhow!("Output memory bounds error"));

226

8

227

228

                // Write the result string into the output buffer

229

8

                memory.data_mut(&mut caller)[buf_ptr as usize..buf_ptr as usize + output_len]

230

8

                    .copy_from_slice(output_bytes);

231

232

8

                Ok(output_len as i32)

233

            })();

234

235

            // Handle the result or return an error code

236

8

            match result {

237

8

                Ok(len) => len,

238

                Err(e) => {

239

                    log::error!("Error in wrap_wasm_str_transform: {e:?}");

240

                    -1 // Indicate failure

241

242

243

8

},

244

245

16

246

247

16

pub fn wrap_wasm_str_list_get<Data, Getter>(

248

16

    store: &mut Store<Arc<DataStore<Data>>>,

249

16

    getter: Getter,

250

16

) -> Func

251

16

where

252

16

    Data: Send + 'static,

253

16

    Getter: Fn(&Data) -> Result<Vec<String>, anyhow::Error> + Send + Sync + 'static,

254

255

16

    Func::wrap(

256

16

        store,

257

16

        move |mut caller: Caller<'_, Arc<DataStore<Data>>>, loc: i32, max_len: i32| -> i32 {

258

16

            let data = Arc::clone(caller.data());

259

260

            // Try to execute the getter and handle potential errors

261

16

            let result = (|| -> Result<i32> {

262

16

                let builder = data

263

16

                    .data

264

16

                    .lock()

265

16

                    .map_err(|_| anyhow::anyhow!("Failed to lock data"))?;

266

16

                let list = getter(&builder)?; // Call getter, which may return an error

267

16

                let mut bytes = Vec::new();

268

269

40

                for s in list {

270

24

                    bytes.extend_from_slice(s.as_bytes());

271

24

                    bytes.push(0); // Null terminator

272

24

273

274

16

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

275

276

16

                let memory = data

277

16

                    .memory

278

16

                    .lock()

279

16

                    .map_err(|_| anyhow::anyhow!("Failed to lock memory"))?;

280

16

                if let Some(memory) = &*memory {

281

16

                    let mem = memory.data_mut(&mut caller);

282

283

16

                    if (loc as usize + len) <= mem.len() {

284

16

                        mem[loc as usize..loc as usize + len].copy_from_slice(&bytes[..len]);

285

16

                        return Ok(len as i32);

286

                    } else {

287

                        return Err(anyhow::anyhow!("Memory bounds error"));

288

289

290

291

                Err(anyhow::anyhow!("Memory not available"))

292

            })();

293

294

            // Return the result or indicate failure (-1)

295

16

            match result {

296

16

                Ok(len) => len,

297

                Err(e) => {

298

                    log::error!("Error in wrap_wasm_str_list_get: {e:?}");

299

-1

300

301

302

16

},

303

304

16

305

306

16

pub fn wrap_wasm_str_map<Data>(

307

16

    store: &mut Store<Arc<DataStore<Data>>>,

308

16

    updater: impl Fn(&mut Data, String, String) -> Result<(), anyhow::Error> + Send + Sync + 'static,

309

16

) -> Func

310

16

where

311

16

    Data: Send + 'static,

312

313

16

    Func::wrap(

314

16

        store,

315

        move |caller: Caller<'_, Arc<DataStore<Data>>>,

316

              key_loc: i32,

317

              key_len: i32,

318

              value_loc: i32,

319

              value_len: i32|

320

24

              -> i32 {

321

24

            let data = caller.data();

322

323

24

            let result = (|| -> Result<()> {

324

24

                let memory = data

325

24

                    .memory

326

24

                    .lock()

327

24

                    .map_err(|_| anyhow::anyhow!("Failed to lock memory"))?;

328

24

                if let Some(memory) = &*memory {

329

24

                    let mem = memory.data(&caller);

330

331

24

                    if (key_loc as usize + key_len as usize) > mem.len()

332

24

                        || (value_loc as usize + value_len as usize) > mem.len()

333

334

                        return Err(anyhow::anyhow!("Memory bounds error"));

335

24

336

337

24

                    let key = String::from_utf8(

338

24

                        mem[key_loc as usize..key_loc as usize + key_len as usize].to_vec(),

339

)?;

340

24

                    let value = String::from_utf8(

341

24

                        mem[value_loc as usize..value_loc as usize + value_len as usize].to_vec(),

342

)?;

343

344

24

                    let mut builder = data

345

24

                        .data

346

24

                        .lock()

347

24

                        .map_err(|_| anyhow::anyhow!("Failed to lock data"))?;

348

24

                    updater(&mut builder, key, value)?;

349

350

24

                    Ok(())

351

                } else {

352

                    Err(anyhow::anyhow!("Memory not available"))

353

354

            })();

355

356

24

            match result {

357

24

                Ok(()) => 0,

358

                Err(e) => {

359

                    log::error!("Error in wrap_wasm_str_map: {e:?}");

360

-1

361

362

363

24

},

364

365

16