From ec44d043e80741b86eb07da6278547357535e4bd Mon Sep 17 00:00:00 2001
From: Ming-Hung Tsai <mtsai@redhat.com>
Date: Thu, 5 Aug 2021 12:28:10 +0800
Subject: [PATCH] [space_map (rust)] Restrict metadata space map size

---
 src/cache/repair.rs             |  4 ++--
 src/cache/restore.rs            |  3 +--
 src/pdata/space_map_metadata.rs | 12 ++++++++++++
 src/thin/repair.rs              |  4 ++--
 src/thin/restore.rs             |  2 +-
 5 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/src/cache/repair.rs b/src/cache/repair.rs
index 1273da7..30650b4 100644
--- a/src/cache/repair.rs
+++ b/src/cache/repair.rs
@@ -6,7 +6,7 @@ use crate::cache::dump::*;
 use crate::cache::restore::*;
 use crate::cache::superblock::*;
 use crate::io_engine::*;
-use crate::pdata::space_map::*;
+use crate::pdata::space_map_metadata::*;
 use crate::report::*;
 use crate::write_batcher::*;
 
@@ -54,7 +54,7 @@ pub fn repair(opts: CacheRepairOptions) -> Result<()> {
 
     let sb = read_superblock(ctx.engine_in.as_ref(), SUPERBLOCK_LOCATION)?;
 
-    let sm = core_sm(ctx.engine_out.get_nr_blocks(), u32::MAX);
+    let sm = core_metadata_sm(ctx.engine_out.get_nr_blocks(), u32::MAX);
     let mut w = WriteBatcher::new(
         ctx.engine_out.clone(),
         sm.clone(),
diff --git a/src/cache/restore.rs b/src/cache/restore.rs
index 067462a..7997e94 100644
--- a/src/cache/restore.rs
+++ b/src/cache/restore.rs
@@ -14,7 +14,6 @@ use crate::cache::xml;
 use crate::io_engine::*;
 use crate::math::*;
 use crate::pdata::array_builder::*;
-use crate::pdata::space_map::*;
 use crate::pdata::space_map_metadata::*;
 use crate::pdata::unpack::Pack;
 use crate::report::*;
@@ -259,7 +258,7 @@ pub fn restore(opts: CacheRestoreOptions) -> Result<()> {
 
     let ctx = mk_context(&opts)?;
 
-    let sm = core_sm(ctx.engine.get_nr_blocks(), u32::MAX);
+    let sm = core_metadata_sm(ctx.engine.get_nr_blocks(), u32::MAX);
     let mut w = WriteBatcher::new(ctx.engine.clone(), sm.clone(), ctx.engine.get_batch_size());
 
     // build cache mappings
diff --git a/src/pdata/space_map_metadata.rs b/src/pdata/space_map_metadata.rs
index 9002ca0..39728b5 100644
--- a/src/pdata/space_map_metadata.rs
+++ b/src/pdata/space_map_metadata.rs
@@ -2,9 +2,11 @@ use anyhow::{anyhow, Result};
 use byteorder::{LittleEndian, WriteBytesExt};
 use nom::{number::complete::*, IResult};
 use std::io::Cursor;
+use std::sync::{Arc, Mutex};
 
 use crate::checksum;
 use crate::io_engine::*;
+use crate::pdata::space_map::*;
 use crate::pdata::space_map_common::*;
 use crate::pdata::unpack::*;
 use crate::write_batcher::*;
@@ -12,6 +14,7 @@ use crate::write_batcher::*;
 //------------------------------------------
 
 const MAX_METADATA_BITMAPS: usize = 255;
+const MAX_METADATA_BLOCKS: usize = MAX_METADATA_BITMAPS * ENTRIES_PER_BITMAP;
 
 //------------------------------------------
 
@@ -102,6 +105,15 @@ fn adjust_counts(
     })
 }
 
+//------------------------------------------
+
+pub fn core_metadata_sm(nr_blocks: u64, max_count: u32) -> Arc<Mutex<dyn SpaceMap + Send + Sync>> {
+    core_sm(
+        std::cmp::min(nr_blocks, MAX_METADATA_BLOCKS as u64),
+        max_count,
+    )
+}
+
 pub fn write_metadata_sm(w: &mut WriteBatcher) -> Result<SMRoot> {
     let r1 = w.get_reserved_range();
 
diff --git a/src/thin/repair.rs b/src/thin/repair.rs
index d4ad6ff..4f74371 100644
--- a/src/thin/repair.rs
+++ b/src/thin/repair.rs
@@ -3,7 +3,7 @@ use std::path::Path;
 use std::sync::Arc;
 
 use crate::io_engine::*;
-use crate::pdata::space_map::*;
+use crate::pdata::space_map_metadata::*;
 use crate::report::*;
 use crate::thin::dump::*;
 use crate::thin::metadata::*;
@@ -57,7 +57,7 @@ pub fn repair(opts: ThinRepairOptions) -> Result<()> {
     let md = build_metadata(ctx.engine_in.clone(), &sb)?;
     let md = optimise_metadata(md)?;
 
-    let sm = core_sm(ctx.engine_out.get_nr_blocks(), u32::MAX);
+    let sm = core_metadata_sm(ctx.engine_out.get_nr_blocks(), u32::MAX);
     let mut w = WriteBatcher::new(
         ctx.engine_out.clone(),
         sm.clone(),
diff --git a/src/thin/restore.rs b/src/thin/restore.rs
index a6738bb..a1fbb6a 100644
--- a/src/thin/restore.rs
+++ b/src/thin/restore.rs
@@ -351,7 +351,7 @@ pub fn restore(opts: ThinRestoreOptions) -> Result<()> {
     let ctx = new_context(&opts)?;
     let max_count = u32::MAX;
 
-    let sm = core_sm(ctx.engine.get_nr_blocks(), max_count);
+    let sm = core_metadata_sm(ctx.engine.get_nr_blocks(), max_count);
     let mut w = WriteBatcher::new(ctx.engine.clone(), sm.clone(), ctx.engine.get_batch_size());
     let mut restorer = Restorer::new(&mut w, ctx.report);
     xml::read(input, &mut restorer)?;