diff --git a/mtda/client.py b/mtda/client.py index 634b239b..6381e670 100644 --- a/mtda/client.py +++ b/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) diff --git a/mtda/main.py b/mtda/main.py index 87e4c9a8..743cb38d 100644 --- a/mtda/main.py +++ b/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()") diff --git a/mtda/storage/controller.py b/mtda/storage/controller.py index e6503ff5..66805925 100644 --- a/mtda/storage/controller.py +++ b/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""" diff --git a/mtda/storage/helpers/image.py b/mtda/storage/helpers/image.py index d53e75d2..9d4f0174 100644 --- a/mtda/storage/helpers/image.py +++ b/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