diff --git a/pyroaring/abstract_bitmap.pxi b/pyroaring/abstract_bitmap.pxi index ab30abc..c29926d 100644 --- a/pyroaring/abstract_bitmap.pxi +++ b/pyroaring/abstract_bitmap.pxi @@ -12,9 +12,18 @@ try: except NameError: # python 3 pass -cdef croaring.roaring_bitmap_t *deserialize_ptr(char *buff): +cdef croaring.roaring_bitmap_t *deserialize_ptr(bytes buff): cdef croaring.roaring_bitmap_t *ptr - ptr = croaring.roaring_bitmap_portable_deserialize(buff) + cdef const char *reason_failure = NULL + buff_size = len(buff) + ptr = croaring.roaring_bitmap_portable_deserialize_safe(buff, buff_size) + if ptr == NULL: + raise ValueError("Could not deserialize bitmap!!") + # Validate the bitmap + if not croaring.roaring_bitmap_internal_validate(ptr, &reason_failure): + # If validation fails, free the bitmap and raise an exception + croaring.roaring_bitmap_free(ptr) + raise ValueError(f"Invalid bitmap: {reason_failure.decode('utf-8')}") return ptr cdef croaring.roaring64_bitmap_t *deserialize64_ptr(bytes buff): @@ -744,7 +753,7 @@ cdef class AbstractBitMap: @classmethod - def deserialize(cls, char *buff): + def deserialize(cls, bytes buff): """ Generate a bitmap from the given serialization. See AbstractBitMap.serialize for the reverse operation. diff --git a/pyroaring/croaring.pxd b/pyroaring/croaring.pxd index ed5a33b..4192733 100644 --- a/pyroaring/croaring.pxd +++ b/pyroaring/croaring.pxd @@ -100,6 +100,8 @@ cdef extern from "roaring.h": size_t roaring_bitmap_portable_size_in_bytes(const roaring_bitmap_t *ra) size_t roaring_bitmap_portable_serialize(const roaring_bitmap_t *ra, char *buf) roaring_bitmap_t *roaring_bitmap_portable_deserialize(const char *buf) + roaring_bitmap_t *roaring_bitmap_portable_deserialize_safe(const char *buf, size_t maxbytes) + bool roaring_bitmap_internal_validate(const roaring_bitmap_t *r, const char **reason) roaring_uint32_iterator_t *roaring_iterator_create(const roaring_bitmap_t *ra) bool roaring_uint32_iterator_advance(roaring_uint32_iterator_t *it) uint32_t roaring_uint32_iterator_read(roaring_uint32_iterator_t *it, uint32_t* buf, uint32_t count)