near_vm_vm/
sig_registry.rs

1// This file contains code from external sources.
2// Attributions: https://github.com/wasmerio/wasmer/blob/2.3.0/ATTRIBUTIONS.md
3
4//! Implement a registry of function signatures, for fast indirect call
5//! signature checking.
6
7use near_vm_types::FunctionType;
8use std::collections::{HashMap, hash_map};
9use std::convert::TryFrom;
10
11/// An index into the shared signature registry, usable for checking signatures
12/// at indirect calls.
13#[repr(C)]
14#[derive(Debug, Eq, PartialEq, Clone, Copy, Hash)]
15pub struct VMSharedSignatureIndex(u32);
16
17impl VMSharedSignatureIndex {
18    /// Create a new `VMSharedSignatureIndex`.
19    pub fn new(value: u32) -> Self {
20        Self(value)
21    }
22}
23
24/// WebAssembly requires that the caller and callee signatures in an indirect
25/// call must match. To implement this efficiently, keep a registry of all
26/// signatures, shared by all instances, so that call sites can just do an
27/// index comparison.
28#[derive(Debug)]
29pub struct SignatureRegistry {
30    type_to_index: HashMap<FunctionType, VMSharedSignatureIndex>,
31    index_to_data: Vec<FunctionType>,
32}
33
34impl SignatureRegistry {
35    /// Create a new `SignatureRegistry`.
36    pub fn new() -> Self {
37        Self { type_to_index: HashMap::new(), index_to_data: Vec::new() }
38    }
39
40    /// Register a signature and return its unique index.
41    pub fn register(&mut self, sig: FunctionType) -> VMSharedSignatureIndex {
42        let len = self.index_to_data.len();
43        // TODO(0-copy): this. should. not. allocate. (and take FunctionTypeRef as a parameter)
44        //
45        // This is pretty hard to avoid, however. In order to implement bijective map, we'd want
46        // a `Rc<FunctionType>`, but indexing into a map keyed by `Rc<FunctionType>` with
47        // `FunctionTypeRef` is… not possible given the current API either.
48        //
49        // Consider `transmute` or `hashbrown`'s raw_entry.
50        match self.type_to_index.entry(sig.clone()) {
51            hash_map::Entry::Occupied(entry) => *entry.get(),
52            hash_map::Entry::Vacant(entry) => {
53                debug_assert!(
54                    u32::try_from(len).is_ok(),
55                    "invariant: can't have more than 2³²-1 signatures!"
56                );
57                let sig_id = VMSharedSignatureIndex::new(u32::try_from(len).unwrap());
58                entry.insert(sig_id);
59                self.index_to_data.push(sig);
60                sig_id
61            }
62        }
63    }
64
65    /// Looks up a shared signature index within this registry.
66    ///
67    /// Note that for this operation to be semantically correct the `idx` must
68    /// have previously come from a call to `register` of this same object.
69    pub fn lookup(&self, idx: VMSharedSignatureIndex) -> Option<&FunctionType> {
70        self.index_to_data.get(idx.0 as usize)
71    }
72}