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}