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

155

                // Memory bounds are invalid; indicate failure

156

                return -1;

157

158

159

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

160

-1

161

},

162

163

164

165

16

pub fn wrap_wasm_str_transform<Data, Transform>(

166

16

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

167

16

    transform: Transform,

168

16

) -> Func

169

16

where

170

16

    Data: Send + 'static,

171

16

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

172

173

16

    Func::wrap(

174

16

        store,

175

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

176

              tag_ptr: i32,

177

              tag_len: i32,

178

              buf_ptr: i32,

179

              buf_len: i32|

180

8

              -> i32 {

181

8

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

182

183

            // Define the inner logic

184

8

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

185

                // Lock the WASM memory

186

8

                let memory = data

187

8

                    .memory

188

8

                    .lock()

189

8

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

190

8

                let memory = memory

191

8

                    .as_ref()

192

8

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

193

8

                let mem = memory.data(&caller);

194

195

                // Validate memory bounds for the input key

196

8

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

197

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

198

8

199

200

                // Extract the key string from memory

201

8

                let key = String::from_utf8(

202

8

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

203

204

8

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

205

206

                // Lock the data and call the transform function

207

8

                let data_lock = data

208

8

                    .data

209

8

                    .lock()

210

8

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

211

8

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

212

213

                // Get output bytes

214

8

                let output_bytes = output.as_bytes();

215

8

                let output_len = output_bytes.len();

216

217

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

218

8

                if output_len > buf_len as usize {

219

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

220

8

221

222

                // Validate memory bounds for the output buffer

223

8

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

224

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

225

8

226

227

                // Write the result string into the output buffer

228

8

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

229

8

                    .copy_from_slice(output_bytes);

230

231

8

                Ok(output_len as i32)

232

            })();

233

234

            // Handle the result or return an error code

235

8

            match result {

236

8

                Ok(len) => len,

237

                Err(e) => {

238

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

239

                    -1 // Indicate failure

240

241

242

8

},

243

244

16

245

246

16

pub fn wrap_wasm_str_list_get<Data, Getter>(

247

16

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

248

16

    getter: Getter,

249

16

) -> Func

250

16

where

251

16

    Data: Send + 'static,

252

16

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

253

254

16

    Func::wrap(

255

16

        store,

256

16

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

257

16

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

258

259

            // Try to execute the getter and handle potential errors

260

16

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

261

16

                let builder = data

262

16

                    .data

263

16

                    .lock()

264

16

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

265

16

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

266

16

                let mut bytes = Vec::new();

267

268

24

                for s in list {

269

24

                    bytes.extend_from_slice(s.as_bytes());

270

24

                    bytes.push(0); // Null terminator

271

24

272

273

16

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

274

275

16

                let memory = data

276

16

                    .memory

277

16

                    .lock()

278

16

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

279

16

                if let Some(memory) = &*memory {

280

16

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

281

282

16

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

283

16

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

284

16

                        return Ok(len as i32);

285

286

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

287

288

289

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

290

            })();

291

292

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

293

16

            match result {

294

16

                Ok(len) => len,

295

                Err(e) => {

296

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

297

-1

298

299

300

16

},

301

302

16

303

304

16

pub fn wrap_wasm_str_map<Data>(

305

16

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

306

16

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

307

16

) -> Func

308

16

where

309

16

    Data: Send + 'static,

310

311

16

    Func::wrap(

312

16

        store,

313

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

314

              key_loc: i32,

315

              key_len: i32,

316

              value_loc: i32,

317

              value_len: i32|

318

24

              -> i32 {

319

24

            let data = caller.data();

320

321

24

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

322

24

                let memory = data

323

24

                    .memory

324

24

                    .lock()

325

24

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

326

24

                if let Some(memory) = &*memory {

327

24

                    let mem = memory.data(&caller);

328

329

24

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

330

24

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

331

332

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

333

24

334

335

24

                    let key = String::from_utf8(

336

24

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

337

)?;

338

24

                    let value = String::from_utf8(

339

24

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

340

)?;

341

342

24

                    let mut builder = data

343

24

                        .data

344

24

                        .lock()

345

24

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

346

24

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

347

348

24

                    Ok(())

349

                } else {

350

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

351

352

            })();

353

354

24

            match result {

355

24

                Ok(()) => 0,

356

                Err(e) => {

357

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

358

-1

359

360

361

24

},

362

363

16