From f8ce367fd5d6861114a2ec445037d0a9423f05ac Mon Sep 17 00:00:00 2001
From: Mike Auty <mike.auty@gmail.com>
Date: Sun, 26 Sep 2021 10:02:54 +0100
Subject: [PATCH] Layers: Fix segmentation for non-linear and linear

---
 volatility3/framework/layers/qemu.py      | 16 +++++++++++++---
 volatility3/framework/layers/segmented.py | 21 +++++++++++++++++----
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/volatility3/framework/layers/qemu.py b/volatility3/framework/layers/qemu.py
index ff0644dbc..4ce17bb63 100644
--- a/volatility3/framework/layers/qemu.py
+++ b/volatility3/framework/layers/qemu.py
@@ -1,6 +1,7 @@
 # This file is Copyright 2020 Volatility Foundation and licensed under the Volatility Software License 1.0
 # which is available at https://www.volatilityfoundation.org/license/vsl-v1.0
 #
+import bisect
 import functools
 import json
 import math
@@ -211,9 +212,18 @@ class QemuSuspendLayer(segmented.NonLinearlySegmentedLayer):
         return index
 
     def _decode_data(self, data: bytes, mapped_offset: int, offset: int, output_length: int) -> bytes:
-        if mapped_offset in self._compressed:
-            return (data * 0x1000)[:output_length]
-        return data
+        """Takes the full segment from the base_layer that the data occurs in, checks whether it's compressed
+        (by locating it in the segment list and verifying if that address is compressed), then reading/expanding the
+        data, and finally cutting it to the right size.  Offset may be the address requested rather than the location
+        of the starting data.  It is the responsibility of the layer to turn the provided data chunk into the right
+        portion of data necessary.
+        """
+        start_offset, _, start_mapped_offset, _ = self._segments[
+            bisect.bisect_right(self._segments, (offset, 0xffffffffffffff,)) - 1]
+        if start_mapped_offset in self._compressed:
+            data = (data * 0x1000)
+        result = data[offset - start_offset:output_length + offset - start_offset]
+        return result
 
     @functools.lru_cache(maxsize = 512)
     def read(self, offset: int, length: int, pad: bool = False) -> bytes:
diff --git a/volatility3/framework/layers/segmented.py b/volatility3/framework/layers/segmented.py
index 80c89723a..05fc01b97 100644
--- a/volatility3/framework/layers/segmented.py
+++ b/volatility3/framework/layers/segmented.py
@@ -35,7 +35,7 @@ class NonLinearlySegmentedLayer(interfaces.layers.TranslationLayerInterface, met
     def _load_segments(self) -> None:
         """Populates the _segments variable.
 
-        Segments must be (address, mapped address, length) and must be
+        Segments must be (address, mapped address, length, mapped_length) and must be
         sorted by address when this method exits
         """
 
@@ -69,6 +69,10 @@ class NonLinearlySegmentedLayer(interfaces.layers.TranslationLayerInterface, met
                 return self._segments[i]
         raise exceptions.InvalidAddressException(self.name, offset, f"Invalid address at {offset:0x}")
 
+    # Determines whether larger segments are in use and the offsets within them should be tracked linearly
+    # When no decoding of the data occurs, this should be set to true
+    _track_offset = False
+
     def mapping(self,
                 offset: int,
                 length: int,
@@ -85,7 +89,8 @@ class NonLinearlySegmentedLayer(interfaces.layers.TranslationLayerInterface, met
                 if current_offset > logical_offset:
                     difference = current_offset - logical_offset
                     logical_offset += difference
-                    mapped_offset += difference
+                    if self._track_offset:
+                        mapped_offset += difference
                     size -= difference
             except exceptions.InvalidAddressException:
                 if not ignore_errors:
@@ -103,7 +108,7 @@ class NonLinearlySegmentedLayer(interfaces.layers.TranslationLayerInterface, met
                     return
             # Crop it to the amount we need left
             chunk_size = min(size, length + offset - logical_offset)
-            yield logical_offset, chunk_size, mapped_offset, chunk_size, self._base_layer
+            yield logical_offset, chunk_size, mapped_offset, mapped_size, self._base_layer
             current_offset += chunk_size
             # Terminate if we've gone (or reached) our required limit
             if current_offset >= offset + length:
@@ -139,4 +144,12 @@ class NonLinearlySegmentedLayer(interfaces.layers.TranslationLayerInterface, met
 
 
 class SegmentedLayer(NonLinearlySegmentedLayer, linear.LinearlyMappedLayer, metaclass = ABCMeta):
-    pass
+    _track_offset = True
+
+    def mapping(self,
+                offset: int,
+                length: int,
+                ignore_errors: bool = False) -> Iterable[Tuple[int, int, int, int, str]]:
+        # Linear mappings must return the same length of segment as that requested
+        for offset, length, mapped_offset, mapped_length, layer in super().mapping(offset, length, ignore_errors):
+            yield offset, length, mapped_offset, length, layer