paritytech / parity-scale-codec

Lightweight, efficient, binary serialization and deserialization codec
Apache License 2.0
243 stars 95 forks source link

Decode with mem limit #602

Closed ggwpez closed 1 month ago

ggwpez commented 3 months ago

Adds a try_alloc function to the Input trait to track allocations. This makes it necessary to implement it in the decode any type, but is not type checked. Currently it is only implemented for vectors.

Some tests are in place to showcase the usage.

This MR adds functionality to restrict the memory usage while decoding a type.
Some tests show how to use it with nested types or tuples, but mostly a PoC to evaluate the approach.
We could also do it like the DecodeLimit and track the memory consumption in the Input stream, but it cannot be enforced by the type system.

PS: I will change it to track the limit in the Input, as it allows to also track the other limits (like recursion depth) at the expense of not being type checked.

Changes Adds a generic DecodeWithContext that allows a decode_with_context call to pass down some context when decoding.

This new trait is then used together with MemLimit such that Vec implements DecodeWithContext<MemLimit>.
There is a blanket in place that converts DecodeWithContext<MemLimit> to DecodeWithMemLimit, and therebyte, Vec is now DecodeWithMemLimit.

serban300 commented 1 month ago

I've been thinking about this some more and I'm not sure if this solution is feasible. My main concern is related to enums. For example, for the following enum:

enum TestEnum {
    A {field: String},
    B {field: u32},
    C {field: [u8; 10000]},
}

I don't know if we can check the type of each field in an enum variant and call try_alloc() just for some certain types. And even if we could, this process would be manual and seems very error-prone.

Another concern is that for structures like BTreeMap, BTreeSet, etc, we'll only record the memory of the elements, but we won't account for the memory of the internal nodes of the map/set. Altough this is not significant, it's still not ideal.

Another problem would be if we have skipped fields that consume significant memory. For example a [u8; 10000].

Another issue is that here we do impl<T: Decode> DecodeMemLimit for Vec<T> {}. This will track memory only for the structures for which we implement DecodeMemLimit. For custom structures memory tracking will work just to the extent to which they actually call decode() for the primitive types inside them. I think the right approach would be impl<T: DecodeMemLimit> DecodeMemLimit for Vec<T> {}, but then the problem is that there will be a lot of custom structures for which DecodeMemLimit is not implemented and it would be a burden to implement it for all.

Also, since I was looking on these, I think there are some issues with the decode_with_depth_limit as well.

  1. it also does impl<T: Decode> DecodeLimit for T { . This might miss structures.
  2. for example should't we do input.descend_ref() and input.ascend_ref() when decoding an enum ?

I'm not sure. Maybe I'm missing something. @ggwpez , @bkchr what do you think ?

serban300 commented 1 month ago

I've been thinking about this some more and I'm not sure if this solution is feasible. My main concern is related to enums. For example, for the following enum:

enum TestEnum {
    A {field: String},
    B {field: u32},
    C {field: [u8; 10000]},
}
* for `TestEnum::A` we need to account for `mem::size_of<TestEnum>` + the size of the decoded String.

* for `TestEnum::B` and `TestEnum::C` we need to account just for `mem::size_of<TestEnum>`

I don't know if we can check the type of each field in an enum variant and call try_alloc() just for some certain types. And even if we could, this process would be manual and seems very error-prone.

Another concern is that for structures like BTreeMap, BTreeSet, etc, we'll only record the memory of the elements, but we won't account for the memory of the internal nodes of the map/set. Altough this is not significant, it's still not ideal.

Another problem would be if we have skipped fields that consume significant memory. For example a [u8; 10000].

Another issue is that here we do impl<T: Decode> DecodeMemLimit for Vec<T> {}. This will track memory only for the structures for which we implement DecodeMemLimit. For custom structures memory tracking will work just to the extent to which they actually call decode() for the primitive types inside them. I think the right approach would be impl<T: DecodeMemLimit> DecodeMemLimit for Vec<T> {}, but then the problem is that there will be a lot of custom structures for which DecodeMemLimit is not implemented and it would be a burden to implement it for all.

Also, since I was looking on these, I think there are some issues with the decode_with_depth_limit as well.

1. it also does `impl<T: Decode> DecodeLimit for T {` . This might miss structures.

2. for example should't we do `input.descend_ref()` and `input.ascend_ref()` when decoding an enum ?

I'm not sure. Maybe I'm missing something. @ggwpez , @bkchr what do you think ?

We'll limit this to tracking only heap size, but we want to make it type safe.

Tracked as part of: https://github.com/paritytech/parity-scale-codec/issues/609

Will close this draft for the moment and will open a separate PR later