@@ -6,7 +6,7 @@ use std::{
66} ;
77
88use bytes:: { Buf , Bytes } ;
9- use loro_common:: LoroResult ;
9+ use loro_common:: { LoroError , LoroResult } ;
1010use once_cell:: sync:: OnceCell ;
1111
1212use crate :: {
@@ -17,6 +17,9 @@ use crate::{
1717
1818use super :: sstable:: { SIZE_OF_U16 , SIZE_OF_U8 } ;
1919
20+ const MAX_NORMAL_BLOCK_DATA_LEN : usize = u16:: MAX as usize ;
21+ const MAX_NORMAL_BLOCK_ENTRIES : usize = u16:: MAX as usize ;
22+
2023#[ derive( Debug , Clone ) ]
2124pub struct LargeValueBlock {
2225 // without checksum
@@ -118,23 +121,108 @@ impl NormalBlock {
118121 first_key : Bytes ,
119122 compression_type : CompressionType ,
120123 ) -> LoroResult < NormalBlock > {
124+ if raw_block_and_check. len ( ) < SIZE_OF_U32 {
125+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
126+ }
127+
121128 let buf = raw_block_and_check. slice ( ..raw_block_and_check. len ( ) - SIZE_OF_U32 ) ;
122129 let mut data = vec ! [ ] ;
123130 decompress ( & mut data, buf, compression_type) ?;
131+ if data. len ( ) < SIZE_OF_U16 {
132+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
133+ }
134+
124135 let offsets_len = ( & data[ data. len ( ) - SIZE_OF_U16 ..] ) . get_u16_le ( ) as usize ;
125- let data_end = data. len ( ) - SIZE_OF_U16 * ( offsets_len + 1 ) ;
136+ if offsets_len == 0 {
137+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
138+ }
139+
140+ let offsets_bytes_len = SIZE_OF_U16
141+ . checked_mul ( offsets_len + 1 )
142+ . ok_or_else ( || LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ?;
143+ if data. len ( ) < offsets_bytes_len {
144+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
145+ }
146+
147+ let data_end = data. len ( ) - offsets_bytes_len;
148+ if data_end > u16:: MAX as usize {
149+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
150+ }
151+
126152 let offsets = & data[ data_end..data. len ( ) - SIZE_OF_U16 ] ;
127- let offsets = offsets
153+ let offsets: Vec < u16 > = offsets
128154 . chunks ( SIZE_OF_U16 )
129155 . map ( |mut chunk| chunk. get_u16_le ( ) )
130156 . collect ( ) ;
157+ Self :: validate_decoded_data ( & data[ ..data_end] , & offsets, & first_key) ?;
131158 Ok ( NormalBlock {
132159 data : Bytes :: copy_from_slice ( & data[ ..data_end] ) ,
133160 encoded_data : OnceCell :: with_value ( ( raw_block_and_check, compression_type) ) ,
134161 offsets,
135162 first_key,
136163 } )
137164 }
165+
166+ fn validate_decoded_data ( data : & [ u8 ] , offsets : & [ u16 ] , first_key : & [ u8 ] ) -> LoroResult < ( ) > {
167+ if offsets. first ( ) . copied ( ) != Some ( 0 ) {
168+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
169+ }
170+
171+ let mut prev_key: Option < Vec < u8 > > = None ;
172+ let mut prev_offset = 0usize ;
173+ for ( idx, offset) in offsets. iter ( ) . map ( |x| * x as usize ) . enumerate ( ) {
174+ let offset_end = offsets
175+ . get ( idx + 1 )
176+ . map_or ( data. len ( ) , |next| * next as usize ) ;
177+ if offset < prev_offset || offset > offset_end || offset_end > data. len ( ) {
178+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
179+ }
180+
181+ let key = if idx == 0 {
182+ first_key. to_vec ( )
183+ } else {
184+ let header_end = offset
185+ . checked_add ( SIZE_OF_U8 + SIZE_OF_U16 )
186+ . ok_or_else ( || LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ?;
187+ if header_end > offset_end {
188+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
189+ }
190+
191+ let common_prefix_len = data[ offset] as usize ;
192+ if common_prefix_len > first_key. len ( ) {
193+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
194+ }
195+
196+ let key_suffix_len =
197+ u16:: from_le_bytes ( data[ offset + SIZE_OF_U8 ..header_end] . try_into ( ) . unwrap ( ) )
198+ as usize ;
199+ let key_end = header_end
200+ . checked_add ( key_suffix_len)
201+ . ok_or_else ( || LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ?;
202+ if key_end > offset_end {
203+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
204+ }
205+
206+ let mut key = Vec :: with_capacity ( common_prefix_len + key_suffix_len) ;
207+ key. extend_from_slice ( & first_key[ ..common_prefix_len] ) ;
208+ key. extend_from_slice ( & data[ header_end..key_end] ) ;
209+ key
210+ } ;
211+
212+ if key. is_empty ( )
213+ || prev_key
214+ . as_ref ( )
215+ . is_some_and ( |prev_key| prev_key. as_slice ( ) >= key. as_slice ( ) )
216+ {
217+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
218+ }
219+
220+ prev_offset = offset;
221+ prev_key = Some ( key) ;
222+ }
223+
224+ Ok ( ( ) )
225+ }
138226}
139227
140228#[ derive( Debug , Clone ) ]
@@ -189,20 +277,31 @@ impl Block {
189277 }
190278 }
191279
192- pub fn decode (
280+ pub ( crate ) fn try_decode (
193281 raw_block_and_check : Bytes ,
194282 is_large : bool ,
195283 key : Bytes ,
196284 compression_type : CompressionType ,
197- ) -> Self {
198- // The caller is responsible for validating SSTable integrity before lazy block reads.
285+ ) -> LoroResult < Self > {
286+ if key. is_empty ( ) {
287+ return Err ( LoroError :: DecodeError ( "Invalid bytes" . into ( ) ) ) ;
288+ }
289+
199290 if is_large {
200291 return LargeValueBlock :: decode ( raw_block_and_check, key, compression_type)
201- . map ( Block :: Large )
202- . expect ( "validated SSTable block should decode" ) ;
292+ . map ( Block :: Large ) ;
203293 }
204- NormalBlock :: decode ( raw_block_and_check, key, compression_type)
205- . map ( Block :: Normal )
294+ NormalBlock :: decode ( raw_block_and_check, key, compression_type) . map ( Block :: Normal )
295+ }
296+
297+ pub fn decode (
298+ raw_block_and_check : Bytes ,
299+ is_large : bool ,
300+ key : Bytes ,
301+ compression_type : CompressionType ,
302+ ) -> Self {
303+ // The caller is responsible for validating SSTable integrity before lazy block reads.
304+ Self :: try_decode ( raw_block_and_check, is_large, key, compression_type)
206305 . expect ( "validated SSTable block should decode" )
207306 }
208307
@@ -273,9 +372,13 @@ impl BlockBuilder {
273372/// └─────────────────────────────────────────────────────┘
274373///
275374pub fn add ( & mut self , key : & [ u8 ] , value : & [ u8 ] ) -> bool {
375+ if key. is_empty ( ) {
376+ return false ;
377+ }
378+
276379 debug_assert ! ( !key. is_empty( ) , "key cannot be empty" ) ;
277380 if self . first_key . is_empty ( ) {
278- if value. len ( ) > self . block_size {
381+ if value. len ( ) > self . block_size || value . len ( ) > MAX_NORMAL_BLOCK_DATA_LEN {
279382 self . data . extend_from_slice ( value) ;
280383 self . is_large = true ;
281384 self . first_key = Bytes :: copy_from_slice ( key) ;
@@ -288,16 +391,39 @@ impl BlockBuilder {
288391 return true ;
289392 }
290393
291- // whether the block is full
292- if self . estimated_size ( ) + key. len ( ) + value. len ( ) + SIZE_OF_U8 + SIZE_OF_U16
293- > self . block_size
294- {
394+ if self . offsets . len ( ) >= MAX_NORMAL_BLOCK_ENTRIES {
295395 return false ;
296396 }
297397
298- self . offsets . push ( self . data . len ( ) as u16 ) ;
299398 let ( common, suffix) = get_common_prefix_len_and_strip ( key, & self . first_key ) ;
300399 let key_len = suffix. len ( ) ;
400+ let Some ( next_data_len) = self
401+ . data
402+ . len ( )
403+ . checked_add ( SIZE_OF_U8 + SIZE_OF_U16 )
404+ . and_then ( |len| len. checked_add ( key_len) )
405+ . and_then ( |len| len. checked_add ( value. len ( ) ) )
406+ else {
407+ return false ;
408+ } ;
409+ if next_data_len > MAX_NORMAL_BLOCK_DATA_LEN {
410+ return false ;
411+ }
412+
413+ // whether the block is full
414+ let Some ( estimated_size) = self
415+ . estimated_size ( )
416+ . checked_add ( key_len)
417+ . and_then ( |len| len. checked_add ( value. len ( ) ) )
418+ . and_then ( |len| len. checked_add ( SIZE_OF_U8 + SIZE_OF_U16 ) )
419+ else {
420+ return false ;
421+ } ;
422+ if estimated_size > self . block_size {
423+ return false ;
424+ }
425+
426+ self . offsets . push ( self . data . len ( ) as u16 ) ;
301427 self . data . push ( common) ;
302428 self . data . extend_from_slice ( & ( key_len as u16 ) . to_le_bytes ( ) ) ;
303429 self . data . extend_from_slice ( suffix) ;
![lodyai[bot]](https://avatars.githubusercontent.com/in/2535747?v=4&size=40){kind=avatar}
0 commit comments