axmap: ensure that we never return a value that is larger than the map size

Signed-off-by: Jens Axboe <axboe@kernel.dk>

diff --git a/lib/axmap.c b/lib/axmap.c
index df547f346fdca8045716161f8127d7e16ca6fbed..533a155c7243ea5195b19baae6d023cbc23335b2 100644 (file)
--- a/lib/axmap.c
+++ b/lib/axmap.c
@@ -48,6 +48,7 @@ struct axmap {
        unsigned int nr_levels;
        struct axmap_level *levels;
        uint64_t first_free;
+       uint64_t nr_bits;
 };
 
 static unsigned long ulog64(unsigned long val, unsigned int log)
@@ -103,6 +104,7 @@ struct axmap *axmap_new(unsigned long nr_bits)
 
        axmap->nr_levels = levels;
        axmap->levels = smalloc(axmap->nr_levels * sizeof(struct axmap_level));
+       axmap->nr_bits = nr_bits;
 
        for (i = 0; i < axmap->nr_levels; i++) {
                struct axmap_level *al = &axmap->levels[i];
@@ -261,6 +263,11 @@ static void __axmap_set(struct axmap *axmap, uint64_t bit_nr,
            axmap->first_free < bit_nr + data->nr_bits)
                axmap->first_free = -1ULL;
 
+       if (bit_nr > axmap->nr_bits)
+               return;
+       else if (bit_nr + nr_bits > axmap->nr_bits)
+               nr_bits = axmap->nr_bits - bit_nr;
+
        set_bits = 0;
        while (nr_bits) {
                axmap_handler(axmap, bit_nr, axmap_set_fn, data);
@@ -302,7 +309,10 @@ static int axmap_isset_fn(struct axmap_level *al, unsigned long offset,
 
 int axmap_isset(struct axmap *axmap, uint64_t bit_nr)
 {
-       return axmap_handler_topdown(axmap, bit_nr, axmap_isset_fn, NULL);
+       if (bit_nr <= axmap->nr_bits)
+               return axmap_handler_topdown(axmap, bit_nr, axmap_isset_fn, NULL);
+
+       return 0;
 }
 
 static uint64_t axmap_find_first_free(struct axmap *axmap, unsigned int level,
@@ -339,7 +349,10 @@ static uint64_t axmap_find_first_free(struct axmap *axmap, unsigned int level,
                }
        }
 
-       return ret;
+       if (ret < axmap->nr_bits)
+               return ret;
+
+       return (uint64_t) -1ULL;
 }
 
 uint64_t axmap_first_free(struct axmap *axmap)