1212// See the License for the specific language governing permissions and
1313// limitations under the License.
1414
15- use crate :: { Error , Decode , Input } ;
15+ use crate :: { Decode , Error , Input } ;
1616
1717/// The error message returned when depth limit is reached.
1818const DECODE_MAX_DEPTH_MSG : & str = "Maximum recursion depth reached when decoding" ;
1919
2020/// Extension trait to [`Decode`] for decoding with a maximum recursion depth.
2121pub trait DecodeLimit : Sized {
22- /// Decode `Self` with the given maximum recursion depth.
22+ /// Decode `Self` with the given maximum recursion depth and advance `input` by the number of
23+ /// bytes consumed.
2324///
2425/// If `limit` is hit, an error is returned.
25- fn decode_with_depth_limit ( limit : u32 , input : & [ u8 ] ) -> Result < Self , Error > ;
26-
27- /// Decode `Self` and advance `input` by the number of bytes consumed.
28- ///
29- /// If `limit` is hit, an error is returned.
30- fn decode_and_advance_with_depth_limit < I : Input > ( limit : u32 , input : & mut I ) -> Result < Self , Error > ;
26+ fn decode_with_depth_limit < I : Input > ( limit : u32 , input : & mut I ) -> Result < Self , Error > ;
3127
3228 /// Decode `Self` and consume all of the given input data.
3329///
3430/// If not all data is consumed or `limit` is hit, an error is returned.
35- fn decode_all_with_depth_limit ( limit : u32 , input : & [ u8 ] ) -> Result < Self , Error > ;
31+ fn decode_all_with_depth_limit ( limit : u32 , input : & mut & [ u8 ] ) -> Result < Self , Error > ;
3632}
3733
38-
3934struct DepthTrackingInput < ' a , I > {
4035 input : & ' a mut I ,
4136 depth : u32 ,
4237 max_depth : u32 ,
4338}
4439
45- impl < ' a , I : Input > Input for DepthTrackingInput < ' a , I > {
40+ impl < ' a , I : Input > Input for DepthTrackingInput < ' a , I > {
4641 fn remaining_len ( & mut self ) -> Result < Option < usize > , Error > {
4742 self . input . remaining_len ( )
4843 }
@@ -72,36 +67,18 @@ impl<'a, I:Input> Input for DepthTrackingInput<'a, I> {
7267}
7368
7469impl < T : Decode > DecodeLimit for T {
75- fn decode_all_with_depth_limit ( limit : u32 , input : & [ u8 ] ) -> Result < Self , Error > {
76- let mut input = DepthTrackingInput {
77- input : & mut & input[ ..] ,
78- depth : 0 ,
79- max_depth : limit,
80- } ;
81- let res = T :: decode ( & mut input) ?;
82-
83- if input. input . is_empty ( ) {
84- Ok ( res)
70+ fn decode_all_with_depth_limit ( limit : u32 , input : & mut & [ u8 ] ) -> Result < Self , Error > {
71+ let t = <Self as DecodeLimit >:: decode_with_depth_limit ( limit, input) ?;
72+
73+ if input. is_empty ( ) {
74+ Ok ( t)
8575 } else {
8676 Err ( crate :: decode_all:: DECODE_ALL_ERR_MSG . into ( ) )
8777 }
8878 }
8979
90- fn decode_and_advance_with_depth_limit < I : Input > ( limit : u32 , input : & mut I ) -> Result < Self , Error > {
91- let mut input = DepthTrackingInput {
92- input,
93- depth : 0 ,
94- max_depth : limit,
95- } ;
96- T :: decode ( & mut input)
97- }
98-
99- fn decode_with_depth_limit ( limit : u32 , input : & [ u8 ] ) -> Result < Self , Error > {
100- let mut input = DepthTrackingInput {
101- input : & mut & input[ ..] ,
102- depth : 0 ,
103- max_depth : limit,
104- } ;
80+ fn decode_with_depth_limit < I : Input > ( limit : u32 , input : & mut I ) -> Result < Self , Error > {
81+ let mut input = DepthTrackingInput { input, depth : 0 , max_depth : limit } ;
10582 T :: decode ( & mut input)
10683 }
10784}
@@ -117,19 +94,38 @@ mod tests {
11794 let nested: NestedVec = vec ! [ vec![ vec![ vec![ 1 ] ] ] ] ;
11895 let encoded = nested. encode ( ) ;
11996
120- let decoded = NestedVec :: decode_with_depth_limit ( 3 , & encoded) . unwrap ( ) ;
97+ let decoded = NestedVec :: decode_with_depth_limit ( 3 , & mut encoded. as_slice ( ) ) . unwrap ( ) ;
12198 assert_eq ! ( decoded, nested) ;
122- assert ! ( NestedVec :: decode_with_depth_limit( 2 , & encoded) . is_err( ) ) ;
99+ assert ! ( NestedVec :: decode_with_depth_limit( 2 , & mut encoded. as_slice ( ) ) . is_err( ) ) ;
123100 }
124101
125102 #[ test]
126- fn decode_and_advance_works ( ) {
103+ fn decode_limit_advances_input ( ) {
127104 type NestedVec = Vec < Vec < Vec < Vec < u8 > > > > ;
128105 let nested: NestedVec = vec ! [ vec![ vec![ vec![ 1 ] ] ] ] ;
129- let encoded = & mut & nested. encode ( ) [ ..] ;
106+ let encoded = nested. encode ( ) ;
107+ let encoded_slice = & mut encoded. as_slice ( ) ;
130108
131- let decoded = Vec :: < u8 > :: decode_and_advance_with_depth_limit ( 1 , encoded ) . unwrap ( ) ;
109+ let decoded = Vec :: < u8 > :: decode_with_depth_limit ( 1 , encoded_slice ) . unwrap ( ) ;
132110 assert_eq ! ( decoded, vec![ 4 ] ) ;
133- assert ! ( NestedVec :: decode_with_depth_limit( 3 , encoded) . is_err( ) ) ;
111+ assert ! ( NestedVec :: decode_with_depth_limit( 3 , encoded_slice) . is_err( ) ) ;
112+ }
113+
114+ #[ test]
115+ fn decode_all_with_limit_advances_input ( ) {
116+ type NestedVec = Vec < Vec < Vec < Vec < u8 > > > > ;
117+ let nested: NestedVec = vec ! [ vec![ vec![ vec![ 1 ] ] ] ] ;
118+ let mut encoded = NestedVec :: encode ( & nested) ;
119+
120+ let decoded = NestedVec :: decode_all_with_depth_limit ( 3 , & mut encoded. as_slice ( ) ) . unwrap ( ) ;
121+ assert_eq ! ( decoded, nested) ;
122+
123+ encoded. extend ( & [ 1 , 2 , 3 , 4 , 5 , 6 ] ) ;
124+ assert_eq ! (
125+ NestedVec :: decode_all_with_depth_limit( 3 , & mut encoded. as_slice( ) )
126+ . unwrap_err( )
127+ . to_string( ) ,
128+ "Input buffer has still data left after decoding!" ,
129+ ) ;
134130 }
135131}
0 commit comments