From 1bbb63f06ba2a0d64cb0b81445ef0a7959b5dbe9 Mon Sep 17 00:00:00 2001
From: Ming-Hung Tsai <mtsai@redhat.com>
Date: Wed, 12 May 2021 15:50:14 +0800
Subject: [PATCH] [cache_check (rust)] Fix discard bitset size checking

---
 src/cache/check.rs  |  2 +-
 src/pdata/bitset.rs | 30 +++++++++++++++---------------
 2 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/src/cache/check.rs b/src/cache/check.rs
index 15b129a..26cda61 100644
--- a/src/cache/check.rs
+++ b/src/cache/check.rs
@@ -337,7 +337,7 @@ pub fn check(opts: CacheCheckOptions) -> anyhow::Result<()> {
         let (_discard_bits, err) = read_bitset_with_sm(
             engine.clone(),
             sb.discard_root,
-            sb.cache_blocks as usize,
+            sb.discard_nr_blocks as usize,
             metadata_sm.clone(),
             opts.ignore_non_fatal,
         )?;
diff --git a/src/pdata/bitset.rs b/src/pdata/bitset.rs
index ef23ea8..2586e74 100644
--- a/src/pdata/bitset.rs
+++ b/src/pdata/bitset.rs
@@ -31,15 +31,15 @@ impl CheckedBitSet {
 }
 
 struct BitsetVisitor {
-    nr_entries: usize,
+    nr_bits: usize,
     bits: Mutex<CheckedBitSet>,
 }
 
 impl BitsetVisitor {
-    pub fn new(nr_entries: usize) -> Self {
+    pub fn new(nr_bits: usize) -> Self {
         BitsetVisitor {
-            nr_entries,
-            bits: Mutex::new(CheckedBitSet::with_capacity(nr_entries)),
+            nr_bits,
+            bits: Mutex::new(CheckedBitSet::with_capacity(nr_bits)),
         }
     }
 
@@ -51,15 +51,15 @@ impl BitsetVisitor {
 impl ArrayVisitor<u64> for BitsetVisitor {
     fn visit(&self, index: u64, b: ArrayBlock<u64>) -> array::Result<()> {
         let mut begin = (index as usize * (b.header.max_entries as usize)) << 6;
+        if begin >= self.nr_bits as usize {
+            return Err(array::value_err(format!(
+                "bitset size exceeds limit: {} bits",
+                self.nr_bits
+            )));
+        }
 
         for i in 0..b.header.nr_entries as usize {
-            if begin > self.nr_entries as usize {
-                return Err(array::value_err(
-                    "bitset size exceeds expectation".to_string(),
-                ));
-            }
-
-            let end: usize = std::cmp::min(begin + 64, self.nr_entries as usize);
+            let end: usize = std::cmp::min(begin + 64, self.nr_bits as usize);
             let mut mask = 1;
             let bits = b.values[i];
 
@@ -77,11 +77,11 @@ impl ArrayVisitor<u64> for BitsetVisitor {
 pub fn read_bitset(
     engine: Arc<dyn IoEngine + Send + Sync>,
     root: u64,
-    nr_entries: usize,
+    nr_bits: usize,
     ignore_none_fatal: bool,
 ) -> (CheckedBitSet, Option<array::ArrayError>) {
     let w = ArrayWalker::new(engine, ignore_none_fatal);
-    let mut v = BitsetVisitor::new(nr_entries);
+    let mut v = BitsetVisitor::new(nr_bits);
     let err = w.walk(&mut v, root);
     let e = match err {
         Ok(()) => None,
@@ -94,12 +94,12 @@ pub fn read_bitset(
 pub fn read_bitset_with_sm(
     engine: Arc<dyn IoEngine + Send + Sync>,
     root: u64,
-    nr_entries: usize,
+    nr_bits: usize,
     sm: Arc<Mutex<dyn SpaceMap + Send + Sync>>,
     ignore_none_fatal: bool,
 ) -> array::Result<(CheckedBitSet, Option<array::ArrayError>)> {
     let w = ArrayWalker::new_with_sm(engine, sm, ignore_none_fatal)?;
-    let mut v = BitsetVisitor::new(nr_entries);
+    let mut v = BitsetVisitor::new(nr_bits);
     let err = w.walk(&mut v, root);
     let e = match err {
         Ok(()) => None,