cleanup reading kcore blocks (#102)

Author: bmc-msft

src/image.rs

@@ -58,6 +58,7 @@ pub struct Header {
     pub version: u32,
 }
 
+#[derive(PartialEq, Eq, Debug, Clone)]
 pub struct Block {
     pub offset: u64,
     pub range: Range<u64>,

src/iomem.rs

@@ -48,7 +48,7 @@ fn parse_file(path: &Path) -> Result<Vec<Range<u64>>, Error> {
         ranges.push(start..end);
     }
 
-    Ok(ranges)
+    Ok(merge_ranges(ranges))
 }
 
 #[must_use]
@@ -122,43 +122,71 @@ mod tests {
 
     #[test]
     fn test_parse_iomem() -> Result<(), Error> {
-        let ranges = parse_file(Path::new("test/iomem.txt"))?;
-        let expected = [
-            4096..654_335,
-            1_048_576..1_073_676_287,
-            4_294_967_296..6_979_321_855,
-        ];
-        assert_eq!(ranges, expected);
-
-        let ranges = parse_file(Path::new("test/iomem-2.txt"))?;
-        let expected = [
-            4096..655_359,
-            1_048_576..1_055_838_207,
-            1_056_026_624..1_073_328_127,
-            1_073_737_728..1_073_741_823,
-            4_294_967_296..6_979_321_855,
-        ];
-        assert_eq!(ranges, expected);
-
-        let ranges = parse_file(Path::new("test/iomem-3.txt"))?;
-        let expected = [
-            65_536..649_215,
-            1_048_576..2_146_303_999,
-            2_146_435_072..2_147_483_647,
-        ];
-        assert_eq!(ranges, expected);
-
-        let ranges = parse_file(Path::new("test/iomem-4.txt"))?;
-        let expected = [
-            4_096..655_359,
-            1_048_576..1_423_523_839,
-            1_423_585_280..1_511_186_431,
-            1_780_150_272..1_818_623_999,
-            1_818_828_800..1_843_613_695,
-            2_071_535_616..2_071_986_175,
-            4_294_967_296..414_464_344_063,
-        ];
-        assert_eq!(ranges, expected);
+        for (filename, expected) in [
+            (
+                Path::new("test/iomem.txt"),
+                vec![
+                    4096..654_335,
+                    1_048_576..1_073_676_287,
+                    4_294_967_296..6_979_321_855,
+                ],
+            ),
+            (
+                Path::new("test/iomem-2.txt"),
+                vec![
+                    4096..655_359,
+                    1_048_576..1_055_838_207,
+                    1_056_026_624..1_073_328_127,
+                    1_073_737_728..1_073_741_823,
+                    4_294_967_296..6_979_321_855,
+                ],
+            ),
+            (
+                Path::new("test/iomem-3.txt"),
+                vec![
+                    65_536..649_215,
+                    1_048_576..2_146_303_999,
+                    2_146_435_072..2_147_483_647,
+                ],
+            ),
+            (
+                Path::new("test/iomem-4.txt"),
+                vec![
+                    4_096..655_359,
+                    1_048_576..1_423_523_839,
+                    1_423_585_280..1_511_186_431,
+                    1_780_150_272..1_818_623_999,
+                    1_818_828_800..1_843_613_695,
+                    2_071_535_616..2_071_986_175,
+                    4_294_967_296..414_464_344_063,
+                ],
+            ),
+            (
+                Path::new("test/iomem-5.txt"),
+                vec![
+                    4_096..655_359,
+                    1_048_576..175_058_967,
+                    175_058_968..175_067_223,
+                    175_067_224..175_071_255,
+                    175_071_256..175_077_463,
+                    175_077_464..175_079_447,
+                    175_079_448..175_087_703,
+                    175_087_704..175_091_735,
+                    175_091_736..175_124_567,
+                    175_124_568..241_524_735,
+                    241_643_520..251_310_079,
+                    251_326_464..251_383_807,
+                    251_424_768..264_671_231,
+                    264_675_328..267_280_383,
+                    267_739_136..267_866_111,
+                    267_870_208..3_221_225_471,
+                    4_294_967296..13_958_643_711,
+                ],
+            ),
+        ] {
+            let ranges = parse_file(filename)?;
+            assert_eq!(ranges, expected);
+        }
 
         Ok(())
     }

src/snapshot.rs

@@ -20,9 +20,6 @@ pub enum Error {
     #[error("locked down /proc/kcore")]
     LockedDownKcore,
 
-    #[error("unable to find memory range: {0:?}")]
-    UnableToFindMemoryRange(Range<u64>),
-
     #[error("unable to create memory snapshot")]
     UnableToCreateMemorySnapshot(#[from] crate::image::Error),
 
@@ -194,6 +191,50 @@ impl<'a, 'b> Snapshot<'a, 'b> {
         Ok(())
     }
 
+    // given a set of ranges from iomem and a set of Blocks derived from the
+    // psuedo-elf phys section headers, derive a set of ranges that can be used
+    // to create a snapshot.
+    fn find_kcore_blocks(ranges: &[Range<u64>], headers: &[Block]) -> Vec<Block> {
+        let mut result = vec![];
+
+        'outer: for range in ranges {
+            let mut range = range.clone();
+
+            'inner: for header in headers {
+                match (
+                    header.range.contains(&range.start),
+                    // TODO: ranges is currently inclusive, but not a
+                    // RangeInclusive.  this should be adjusted.
+                    header.range.contains(&(range.end - 1)),
+                ) {
+                    (true, true) => {
+                        let block = Block {
+                            offset: header.offset + range.start - header.range.start,
+                            range: range.clone(),
+                        };
+
+                        result.push(block);
+                        continue 'outer;
+                    }
+                    (true, false) => {
+                        let block = Block {
+                            offset: header.offset + range.start - header.range.start,
+                            range: range.start..header.range.end,
+                        };
+
+                        result.push(block);
+                        range.start = header.range.end;
+                    }
+                    _ => {
+                        continue 'inner;
+                    }
+                };
+            }
+        }
+
+        result
+    }
+
     fn kcore(&self) -> Result<()> {
         if !is_kcore_ok() {
             return Err(Error::LockedDownKcore);
@@ -206,21 +247,19 @@ impl<'a, 'b> Snapshot<'a, 'b> {
         file.phdrs.sort_by(|a, b| a.vaddr.cmp(&b.vaddr));
         let start = file.phdrs[0].vaddr - self.memory_ranges[0].start;
 
-        let mut blocks = vec![];
-        'outer: for range in &self.memory_ranges {
-            for phdr in &file.phdrs {
-                if range.start == phdr.vaddr - start {
-                    let size = u64::min(phdr.memsz, range.end - range.start);
-                    blocks.push(Block {
-                        offset: phdr.offset,
-                        range: range.start..range.start + size,
-                    });
-                    continue 'outer;
-                }
-            }
-            return Err(Error::UnableToFindMemoryRange(range.clone()));
+        let mut physical_ranges = vec![];
+
+        for phdr in file.phdrs {
+            let entry_start = phdr.vaddr - start;
+            let entry_end = entry_start + phdr.memsz;
+
+            physical_ranges.push(Block {
+                range: entry_start..entry_end,
+                offset: phdr.offset,
+            });
         }
 
+        let blocks = Self::find_kcore_blocks(&self.memory_ranges, &physical_ranges);
         image.write_blocks(&blocks)?;
         Ok(())
     }
@@ -247,3 +286,55 @@ impl<'a, 'b> Snapshot<'a, 'b> {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn translate_ranges() {
+        let ranges = [10..20, 30..35, 45..55];
+
+        let core_ranges = [
+            Block {
+                range: 10..20,
+                offset: 0,
+            },
+            Block {
+                range: 25..35,
+                offset: 10,
+            },
+            Block {
+                range: 40..50,
+                offset: 20,
+            },
+            Block {
+                range: 50..55,
+                offset: 35,
+            },
+        ];
+
+        let expected = vec![
+            Block {
+                offset: 0,
+                range: 10..20,
+            },
+            Block {
+                offset: 10 + 5,
+                range: 30..35,
+            },
+            Block {
+                offset: 25,
+                range: 45..50,
+            },
+            Block {
+                offset: 35,
+                range: 50..55,
+            },
+        ];
+
+        let result = Snapshot::find_kcore_blocks(&ranges, &core_ranges);
+
+        assert_eq!(result, expected);
+    }
+}