Added algorithm to auto-parse .bmap 2.0 files and write accordingly

In `client.py` a bmap file in the same folder as the image supplied with `mtda-cli storage write image.wic` is searched.
If there is no .bmap  or the parsing of the .bmap file goes wrong,
the image is written just like before. The algorithm was tested with 2 different images with holes.
Each currently supported format (raw, bz, gz, xz and zst) was tested for both of them and successfully flashed to a storage.

Signed-off-by: Manuel Josef Matzinger <[email protected]>

mtda/client.py

@@ -18,6 +18,7 @@
 
 from mtda.main import MultiTenantDeviceAccess
 import mtda.constants as CONSTS
+import xml.etree.ElementTree as ET
 
 
 class Client:
@@ -260,6 +261,24 @@ def storage_write_image(self, path, callback=None):
         except FileNotFoundError:
             return False
 
+        # Automatically discover the bmap file
+        image_path = path
+        bmapDict = None
+        while True:
+            bmap_path = image_path + ".bmap"
+            try:
+                bmap = ET.parse(bmap_path)
+                print("discovered bmap file '%s'" % bmap_path)
+                bmapDict = self.parseBmap(bmap, bmap_path)
+                self._impl.storage_bmap_dict(bmapDict, self._session)
+                break
+            except Exception:
+                pass
+            image_path, ext = os.path.splitext(image_path)
+            if ext == "":
+                print("No bmap file found at location of image")
+                break
+
         # Open the shared storage device
         status = self.storage_open()
         if status is False:
@@ -268,6 +287,32 @@ def storage_write_image(self, path, callback=None):
 
         return self._storage_write(image, imgname, imgsize, callback)
 
+    def parseBmap(self, bmap, bmap_path):
+        try:
+            bmapDict = {}
+            broot = bmap.getroot()
+            bmapDict["BlockSize"] = int(
+                broot.find("BlockSize").text.strip())
+            bmapDict["BlocksCount"] = int(
+                bmap.find("BlocksCount").text.strip())
+            bmapDict["MappedBlocksCount"] = int(
+                bmap.find("MappedBlocksCount").text.strip())
+            bmapDict["ImageSize"] = int(
+                bmap.find("ImageSize").text.strip())
+            bmapDict["BlockMap"] = []
+            for child in broot.find("BlockMap").findall("Range"):
+                first, last = child.text.strip().split("-")
+                bmapDict["BlockMap"].append({
+                    "first": int(first),
+                    "last": int(last),
+                    "chksum": child.attrib["chksum"]
+                })
+        except Exception:
+            print("Error parsing '%s', probably not a bmap 2.0 file"
+                  % bmap_path)
+            return None
+        return bmapDict
+
     def storage_to_host(self):
         return self._impl.storage_to_host(self._session)
 

mtda/main.py

@@ -568,6 +568,17 @@ def storage_compression(self, compression, session=None):
         self.mtda.debug(3, "main.storage_compression(): %s" % str(result))
         return result
 
+    def storage_bmap_dict(self, bmapDict, session=None):
+        self.mtda.debug(3, "main.storage_bmap_dict()")
+
+        self._session_check(session)
+        if self.storage_controller is None:
+            result = None
+        else:
+            self.storage_controller.setBmap(bmapDict)
+            result = True
+        self.mtda.debug(3, "main.storage_bmap_dict()(): %s" % str(result))
+
     def storage_close(self, session=None):
         self.mtda.debug(3, "main.storage_close()")
 

mtda/storage/controller.py

@@ -41,6 +41,11 @@ def probe(self):
         """ Check presence of the shared storage device"""
         return False
 
+    @abc.abstractmethod
+    def setBmap(self, bmapDict):
+        """ set up the bmapDict for writing the image faster"""
+        return False
+
     @abc.abstractmethod
     def supports_hotplug(self):
         """ Whether the shared storage device may be hot-plugged"""

mtda/storage/helpers/image.py

@@ -16,6 +16,7 @@
 import psutil
 import subprocess
 import threading
+import io
 
 # Local imports
 import mtda.constants as CONSTS
@@ -29,6 +30,10 @@ def __init__(self, mtda):
         self.handle = None
         self.isfuse = False
         self.isloop = False
+        self.bmapDict = None
+        self.lastBlockIdx = 0
+        self.crtBlockRange = 0
+        self.overlap = 0
         self.lock = threading.Lock()
         atexit.register(self._umount)
 
@@ -269,6 +274,12 @@ def status(self):
         self.lock.release()
         return result
 
+    def setBmap(self, bmapDict):
+        self.bmapDict = bmapDict
+        self.crtBlockRange = 0
+        self.lastBlockIdx = 0
+        self.overlap = 0
+
     def supports_hotplug(self):
         return False
 
@@ -330,8 +341,94 @@ def write(self, data):
 
         result = None
         if self.handle is not None:
-            result = self.handle.write(data)
-
+            # Check if there is a valid bmapDict, write all data otherwise
+            if self.bmapDict is not None:
+                result = 0
+                datalen = len(data)
+                blksz = self.bmapDict["BlockSize"]
+                cRange = self.bmapDict["BlockMap"][self.crtBlockRange]
+                # This happens at xz,bz2 and gz compression
+                # since the decompressed chunk can be smaller than a block
+                if datalen < blksz:
+                    # If there is overlap and the new data is enough
+                    # to form a block, write it
+                    if self.overlap + datalen >= blksz:
+                        dataBlockWritten = \
+                            self.handle.write(data[0:(blksz-self.overlap)])
+                        result += dataBlockWritten
+                        self.lastBlockIdx += 1
+                        # Write the rest of the data
+                        rest = self.handle.write(data[dataBlockWritten:])
+                        result += rest
+                        self.overlap = rest
+                        self.mtda.debug(3, "storage.helpers.image.write(): %s"
+                                        % str(result))
+                        self.lock.release()
+                        return result
+                    else:
+                        # Not enough data to even write a block
+                        # write it and add to the overlap
+                        subBlockWritten = self.handle.write(data)
+                        result += subBlockWritten
+                        self.overlap += subBlockWritten
+                        self.mtda.debug(3, "storage.helpers.image.write(): %s"
+                                        % str(result))
+                        self.lock.release()
+                        return result
+                # Complete partially written last block,
+                # start blockpointer depending on this
+                if self.overlap:
+                    result += self.handle.write(data[0:(blksz-self.overlap)])
+                    self.lastBlockIdx += 1
+                    b = blksz-self.overlap
+                else:
+                    b = 0
+
+                while b < datalen:
+                    if self.lastBlockIdx > cRange['last']:
+                        self.crtBlockRange += 1
+                        cRange = \
+                            self.bmapDict["BlockMap"][self.crtBlockRange]
+                    flBlks2Write = 1
+                    advanceBlocks = 1
+                    if self.lastBlockIdx >= cRange['first']:
+                        flBlks2Write = int((datalen-b)/blksz)
+                        # Write multiple blocks if possible in that range
+                        if (self.lastBlockIdx + flBlks2Write) < cRange['last']:
+                            # Check if it is a full Block
+                            if flBlks2Write:
+                                result += self.handle.write(
+                                    data[b:b+flBlks2Write*blksz])
+                                advanceBlocks = flBlks2Write
+                            else:
+                                result += self.handle.write(data[b:b+blksz])
+                                # Needed for incrementing dynamic blockpointer
+                                flBlks2Write = 1
+                        else:
+                            result += self.handle.write(data[b:b+blksz])
+                            flBlks2Write = 1
+                    else:
+                        # If there is a block to skip use the opportunity
+                        # to also drop the overlap, skip the subBlock
+                        # This allows for faster write-times
+                        # since the time consuming self.handle.write()
+                        # does not need to be called at the start
+                        # when there is no overlap anymore
+                        if self.overlap:
+                            b += self.overlap
+                            self.handle.seek(self.overlap, io.SEEK_CUR)
+                            self.overlap = 0
+                        self.handle.seek(min(datalen-b, blksz), io.SEEK_CUR)
+                    # Only increment block index in case we wrote full blocks
+                    if b + blksz*flBlks2Write <= datalen:
+                        self.lastBlockIdx += advanceBlocks
+                    else:
+                        self.overlap = datalen - b
+                    # Increase blockpointer
+                    b += flBlks2Write * blksz
+            else:
+                # No bmap
+                result = self.handle.write(data)
         self.mtda.debug(3, "storage.helpers.image.write(): %s" % str(result))
         self.lock.release()
         return result