1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

// This file is part of Substrate.

// Copyright (C) 2020-2021 Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! A set of utilities for resetting a wasm instance to its initial state.

use crate::error::{self, Error};
use std::mem;
use parity_wasm::elements::{deserialize_buffer, DataSegment, Instruction, Module as RawModule};

/// A bunch of information collected from a WebAssembly module.
pub struct WasmModuleInfo {
	raw_module: RawModule,
}

impl WasmModuleInfo {
	/// Create `WasmModuleInfo` from the given wasm code.
	///
	/// Returns `None` if the wasm code cannot be deserialized.
	pub fn new(wasm_code: &[u8]) -> Option<Self> {
		let raw_module: RawModule = deserialize_buffer(wasm_code).ok()?;
		Some(Self { raw_module })
	}

	/// Extract the data segments from the given wasm code.
	///
	/// Returns `Err` if the given wasm code cannot be deserialized.
	fn data_segments(&self) -> Vec<DataSegment> {
		self.raw_module
			.data_section()
			.map(|ds| ds.entries())
			.unwrap_or(&[])
			.to_vec()
	}

	/// The number of globals defined in locally in this module.
	pub fn declared_globals_count(&self) -> u32 {
		self.raw_module
			.global_section()
			.map(|gs| gs.entries().len() as u32)
			.unwrap_or(0)
	}

	/// The number of imports of globals.
	pub fn imported_globals_count(&self) -> u32 {
		self.raw_module
			.import_section()
			.map(|is| is.globals() as u32)
			.unwrap_or(0)
	}
}

/// This is a snapshot of data segments specialzied for a particular instantiation.
///
/// Note that this assumes that no mutable globals are used.
#[derive(Clone)]
pub struct DataSegmentsSnapshot {
	/// The list of data segments represented by (offset, contents).
	data_segments: Vec<(u32, Vec<u8>)>,
}

impl DataSegmentsSnapshot {
	/// Create a snapshot from the data segments from the module.
	pub fn take(module: &WasmModuleInfo) -> error::Result<Self> {
		let data_segments = module
			.data_segments()
			.into_iter()
			.map(|mut segment| {
				// Just replace contents of the segment since the segments will be discarded later
				// anyway.
				let contents = mem::replace(segment.value_mut(), vec![]);

				let init_expr = match segment.offset() {
					Some(offset) => offset.code(),
					// Return if the segment is passive
					None => return Err(Error::SharedMemUnsupported),
				};

				// [op, End]
				if init_expr.len() != 2 {
					return Err(Error::InitializerHasTooManyExpressions);
				}
				let offset = match &init_expr[0] {
					Instruction::I32Const(v) => *v as u32,
					Instruction::GetGlobal(_) => {
						// In a valid wasm file, initializer expressions can only refer imported
						// globals.
						//
						// At the moment of writing the Substrate Runtime Interface does not provide
						// any globals. There is nothing that prevents us from supporting this
						// if/when we gain those.
						return Err(Error::ImportedGlobalsUnsupported);
					}
					insn => {
						return Err(Error::InvalidInitializerExpression(format!("{:?}", insn)))
					}
				};

				Ok((offset, contents))
			})
			.collect::<error::Result<Vec<_>>>()?;

		Ok(Self { data_segments })
	}

	/// Apply the given snapshot to a linear memory.
	///
	/// Linear memory interface is represented by a closure `memory_set`.
	pub fn apply<E>(
		&self,
		mut memory_set: impl FnMut(u32, &[u8]) -> Result<(), E>,
	) -> Result<(), E> {
		for (offset, contents) in &self.data_segments {
			memory_set(*offset, contents)?;
		}
		Ok(())
	}
}