massa_ledger_exports/
key.rs

1use massa_db_exports::LEDGER_PREFIX;
2use massa_models::{
3    address::{Address, AddressDeserializer, AddressSerializer},
4    serialization::{VecU8Deserializer, VecU8Serializer},
5};
6use massa_serialization::{
7    Deserializer, SerializeError, Serializer, U64VarIntDeserializer, U64VarIntSerializer,
8};
9use nom::error::{ContextError, ParseError};
10use num_enum::{IntoPrimitive, TryFromPrimitive};
11use std::ops::Bound::Included;
12
13pub const VERSION_IDENT: u8 = 0u8;
14pub const BALANCE_IDENT: u8 = 1u8;
15pub const BYTECODE_IDENT: u8 = 2u8;
16pub const DATASTORE_IDENT: u8 = 3u8;
17pub const KEY_VERSION: u64 = 0;
18
19#[derive(PartialEq, Eq, Clone, IntoPrimitive, TryFromPrimitive, Debug)]
20#[repr(u8)]
21enum KeyTypeId {
22    Version = 0,
23    Balance = 1,
24    Bytecode = 2,
25    Datastore = 3,
26}
27
28#[derive(PartialEq, Eq, Clone, Debug)]
29pub enum KeyType {
30    VERSION,
31    BALANCE,
32    BYTECODE,
33    DATASTORE(Vec<u8>),
34}
35
36#[derive(Default, Clone)]
37pub struct KeyTypeSerializer {
38    vec_u8_serializer: VecU8Serializer,
39    // Whether is deserialized with VecU8Deserializer or not.
40    // If true, we use the VecU8Serializer to serialize the key which will add the length at the beginning.
41    // If false, we just serialize the key as is.
42    // This allows us to store the datastore key length at the beginning of the key or not.
43    // The datastore key length is useful when transferring multiple keys, like in packets,
44    // but isn't when storing a datastore key in the ledger.
45    with_datastore_key_length: bool,
46}
47
48impl KeyTypeSerializer {
49    /// Creates a new KeyTypeSerializer.
50    /// `with_datastore_key_length` if true, the datastore key is serialized with its length.
51    pub fn new(with_datastore_key_length: bool) -> Self {
52        Self {
53            vec_u8_serializer: VecU8Serializer::new(),
54            with_datastore_key_length,
55        }
56    }
57}
58
59impl Serializer<KeyType> for KeyTypeSerializer {
60    fn serialize(&self, value: &KeyType, buffer: &mut Vec<u8>) -> Result<(), SerializeError> {
61        match value {
62            KeyType::VERSION => buffer.extend(&[u8::from(KeyTypeId::Version)]),
63            KeyType::BALANCE => buffer.extend(&[u8::from(KeyTypeId::Balance)]),
64            KeyType::BYTECODE => buffer.extend(&[u8::from(KeyTypeId::Bytecode)]),
65            KeyType::DATASTORE(data) => {
66                buffer.extend(&[u8::from(KeyTypeId::Datastore)]);
67                if self.with_datastore_key_length {
68                    self.vec_u8_serializer.serialize(data, buffer)?;
69                } else {
70                    buffer.extend(data);
71                }
72            }
73        }
74        Ok(())
75    }
76}
77
78#[derive(Clone)]
79pub struct KeyTypeDeserializer {
80    vec_u8_deserializer: VecU8Deserializer,
81    // Same as in KeyTypeSerializer but for deserialization.
82    with_datastore_key_length: bool,
83}
84
85impl KeyTypeDeserializer {
86    /// Creates a new KeyTypeDeserializer.
87    /// `max_datastore_key_length` is the maximum length of a datastore key.
88    /// `with_datastore_key_length` if true, the datastore key is deserialized with its length.
89    pub fn new(max_datastore_key_length: u8, with_datastore_key_length: bool) -> Self {
90        Self {
91            vec_u8_deserializer: VecU8Deserializer::new(
92                Included(u64::MIN),
93                Included(max_datastore_key_length as u64),
94            ),
95            with_datastore_key_length,
96        }
97    }
98}
99
100impl Deserializer<KeyType> for KeyTypeDeserializer {
101    fn deserialize<'a, E: ParseError<&'a [u8]> + ContextError<&'a [u8]>>(
102        &self,
103        input: &'a [u8],
104    ) -> nom::IResult<&'a [u8], KeyType, E> {
105        let (rest, key_type) = nom::number::complete::le_u8(input)?;
106        match KeyTypeId::try_from(key_type) {
107            Ok(KeyTypeId::Balance) => Ok((rest, KeyType::BALANCE)),
108            Ok(KeyTypeId::Bytecode) => Ok((rest, KeyType::BYTECODE)),
109            Ok(KeyTypeId::Datastore) => {
110                if self.with_datastore_key_length {
111                    let (rest, data) = self.vec_u8_deserializer.deserialize(rest)?;
112                    Ok((rest, KeyType::DATASTORE(data)))
113                } else {
114                    Ok((&[], KeyType::DATASTORE(rest.to_vec())))
115                }
116            }
117            Ok(KeyTypeId::Version) => Ok((rest, KeyType::VERSION)),
118            Err(_) => Err(nom::Err::Error(E::from_error_kind(
119                rest,
120                nom::error::ErrorKind::Tag,
121            ))),
122        }
123    }
124}
125
126/// Disk ledger keys representation
127#[derive(PartialEq, Eq, Clone, Debug)]
128pub struct Key {
129    pub key_type: KeyType,
130    pub address: Address,
131}
132
133impl Key {
134    pub fn new(address: &Address, key_type: KeyType) -> Self {
135        Self {
136            key_type,
137            address: *address,
138        }
139    }
140}
141
142/// Gives the general prefix of the datastore of an address while respecting a provided key prefix
143pub fn datastore_prefix_from_address(address: &Address, prefix: &[u8]) -> Vec<u8> {
144    let mut res_prefix = LEDGER_PREFIX.as_bytes().to_vec();
145    U64VarIntSerializer::new()
146        .serialize(&KEY_VERSION, &mut res_prefix)
147        .unwrap();
148    AddressSerializer::new()
149        .serialize(address, &mut res_prefix)
150        .unwrap();
151    res_prefix.push(DATASTORE_IDENT);
152    res_prefix.extend(prefix);
153    res_prefix
154}
155
156/// Basic key serializer
157#[derive(Default, Clone)]
158pub struct KeySerializer {
159    address_serializer: AddressSerializer,
160    key_type_serializer: KeyTypeSerializer,
161    version_byte_serializer: U64VarIntSerializer,
162}
163
164impl KeySerializer {
165    /// Creates a new `KeySerializer`
166    /// `with_datastore_key_length` if true, the datastore key is serialized with its length.
167    pub fn new(with_datastore_key_length: bool) -> Self {
168        Self {
169            address_serializer: AddressSerializer::new(),
170            key_type_serializer: KeyTypeSerializer::new(with_datastore_key_length),
171            version_byte_serializer: U64VarIntSerializer::new(),
172        }
173    }
174}
175
176impl Serializer<Key> for KeySerializer {
177    /// ```
178    /// use massa_models::address::Address;
179    /// use massa_ledger_exports::{KeySerializer, KeyType, Key};
180    /// use massa_serialization::Serializer;
181    /// use massa_hash::Hash;
182    /// use std::str::FromStr;
183    ///
184    /// let mut serialized = Vec::new();
185    /// let address = Address::from_str("AU12dG5xP1RDEB5ocdHkymNVvvSJmUL9BgHwCksDowqmGWxfpm93x").unwrap();
186    /// let store_key = Hash::compute_from(b"test");
187    /// let mut key = Key::new(&address, KeyType::DATASTORE(store_key.into_bytes().to_vec()));
188    /// KeySerializer::new(true).serialize(&key, &mut serialized).unwrap();
189    /// ```
190    fn serialize(&self, value: &Key, buffer: &mut Vec<u8>) -> Result<(), SerializeError> {
191        buffer.extend(LEDGER_PREFIX.as_bytes());
192
193        self.version_byte_serializer
194            .serialize(&KEY_VERSION, buffer)?;
195        self.address_serializer.serialize(&value.address, buffer)?;
196        self.key_type_serializer
197            .serialize(&value.key_type, buffer)?;
198
199        Ok(())
200    }
201}
202
203/// Basic key deserializer
204#[derive(Clone)]
205pub struct KeyDeserializer {
206    address_deserializer: AddressDeserializer,
207    key_type_deserializer: KeyTypeDeserializer,
208    version_byte_deserializer: U64VarIntDeserializer,
209}
210
211impl KeyDeserializer {
212    /// Creates a new `KeyDeserializer`
213    /// `max_datastore_key_length` is the maximum length of a datastore key.
214    /// `with_datastore_key_length` if true, the datastore key is deserialized with its length.
215    pub fn new(max_datastore_key_length: u8, with_datastore_key_length: bool) -> Self {
216        Self {
217            address_deserializer: AddressDeserializer::new(),
218            key_type_deserializer: KeyTypeDeserializer::new(
219                max_datastore_key_length,
220                with_datastore_key_length,
221            ),
222            version_byte_deserializer: U64VarIntDeserializer::new(Included(0), Included(u64::MAX)),
223        }
224    }
225}
226
227impl Deserializer<Key> for KeyDeserializer {
228    /// ## Example
229    /// ```
230    /// use massa_models::address::Address;
231    /// use massa_ledger_exports::{KeyDeserializer, KeySerializer, DATASTORE_IDENT, BALANCE_IDENT, KeyType, Key};
232    /// use massa_serialization::{Deserializer, Serializer, DeserializeError};
233    /// use massa_hash::Hash;
234    /// use std::str::FromStr;
235    ///
236    /// let address = Address::from_str("AU12dG5xP1RDEB5ocdHkymNVvvSJmUL9BgHwCksDowqmGWxfpm93x").unwrap();
237    /// let store_key = Hash::compute_from(b"test");
238    ///
239    /// let mut key = Key::new(&address, KeyType::DATASTORE(store_key.into_bytes().to_vec()));
240    /// let mut serialized = Vec::new();
241    /// KeySerializer::new(true).serialize(&key, &mut serialized).unwrap();
242    /// let (rest, key_deser) = KeyDeserializer::new(255, true).deserialize::<DeserializeError>(&serialized).unwrap();
243    /// assert!(rest.is_empty());
244    /// assert_eq!(key_deser, key);
245    ///
246    /// let mut key = Key::new(&address, KeyType::BALANCE);
247    /// let mut serialized = Vec::new();
248    /// KeySerializer::new(true).serialize(&key, &mut serialized).unwrap();
249    /// let (rest, key_deser) = KeyDeserializer::new(255, true).deserialize::<DeserializeError>(&serialized).unwrap();
250    /// assert!(rest.is_empty());
251    /// assert_eq!(key_deser, key);
252    /// ```
253    fn deserialize<'a, E: ParseError<&'a [u8]> + ContextError<&'a [u8]>>(
254        &self,
255        buffer: &'a [u8],
256    ) -> nom::IResult<&'a [u8], Key, E> {
257        // Skip the ledger prefix without slicing past the end of the buffer:
258        // an undersized input must yield a clean parse error rather than an
259        // out-of-bounds panic.
260        let prefix_len = LEDGER_PREFIX.len();
261        let after_prefix = buffer.get(prefix_len..).ok_or_else(|| {
262            nom::Err::Error(E::from_error_kind(buffer, nom::error::ErrorKind::Eof))
263        })?;
264
265        let (rest, version) = self.version_byte_deserializer.deserialize(after_prefix)?;
266        // Enforce the canonical key version. `KeySerializer` always writes
267        // `KEY_VERSION`, so any other value is a non-canonical / malformed key.
268        // Rejecting it prevents version malleability where distinct encodings
269        // deserialize into the same logical `Key`.
270        if version != KEY_VERSION {
271            return Err(nom::Err::Error(E::from_error_kind(
272                after_prefix,
273                nom::error::ErrorKind::Verify,
274            )));
275        }
276        let (rest, address) = self.address_deserializer.deserialize(rest)?;
277        let (rest, key_type) = self.key_type_deserializer.deserialize(rest)?;
278
279        Ok((rest, Key { address, key_type }))
280    }
281}
282
283#[cfg(test)]
284mod tests {
285    use super::{Key, KeyDeserializer, KeySerializer, KeyType};
286    use massa_db_exports::LEDGER_PREFIX;
287    use massa_models::address::Address;
288    use massa_serialization::{DeserializeError, Deserializer, Serializer};
289    use std::str::FromStr;
290
291    fn sample_address() -> Address {
292        Address::from_str("AU12dG5xP1RDEB5ocdHkymNVvvSJmUL9BgHwCksDowqmGWxfpm93x").unwrap()
293    }
294
295    #[test]
296    fn canonical_key_still_round_trips() {
297        let key = Key::new(&sample_address(), KeyType::BALANCE);
298        let mut serialized = Vec::new();
299        KeySerializer::new(true)
300            .serialize(&key, &mut serialized)
301            .unwrap();
302        let (rest, deser) = KeyDeserializer::new(255, true)
303            .deserialize::<DeserializeError>(&serialized)
304            .unwrap();
305        assert!(rest.is_empty());
306        assert_eq!(deser, key);
307    }
308
309    #[test]
310    fn undersized_buffer_errors_instead_of_panicking() {
311        // A buffer shorter than the ledger prefix must not trigger an
312        // out-of-bounds slice panic.
313        let short = vec![0u8; LEDGER_PREFIX.len().saturating_sub(1)];
314        let res = KeyDeserializer::new(255, true).deserialize::<DeserializeError>(&short);
315        assert!(res.is_err(), "undersized buffer must produce a parse error");
316    }
317
318    #[test]
319    fn non_canonical_version_is_rejected() {
320        let key = Key::new(&sample_address(), KeyType::BALANCE);
321        let mut serialized = Vec::new();
322        KeySerializer::new(true)
323            .serialize(&key, &mut serialized)
324            .unwrap();
325
326        // The version is a u64 varint written immediately after the prefix;
327        // `KEY_VERSION == 0` encodes as a single `0x00` byte. Flip it to a
328        // non-canonical value.
329        let version_idx = LEDGER_PREFIX.len();
330        assert_eq!(serialized[version_idx], 0u8);
331        serialized[version_idx] = 1u8;
332
333        let res = KeyDeserializer::new(255, true).deserialize::<DeserializeError>(&serialized);
334        assert!(res.is_err(), "a non-canonical key version must be rejected");
335    }
336}
massa_ledger_exports/key.rs

massa_ledger_exports/
key.rs
