Merge lp:~mvo/ubuntu-system-image/lp1271684 into lp:~registry/ubuntu-system-image/client

[Michael Vogt]

This branch makes system-image more memory efficient when calculating the file signatures. This fixes lp:1271684.

Its a bit of a drive-by branch, I stumbled over the bug and it looked easy enough to get familiar with the code.

Note that I added a README to help new developers getting into the code, I had trouble running the full testsuite, I'm happy to replace the "???" in the README with the proper command. I tried "python3 -m nose2 -v" but that fails early in find_dbus_process on my utopic system.

[Barry Warsaw (barry)]

Hi Michael - thanks for this branch! Somehow I missed this bug. The branch looks good, I'll merge this into trunk for s-i 2.3. Thanks!

FWIW, you can just run `tox` at the command line to run the test suite. That has the side-effect of building the environment to run the tests. `python3 -m nose2 -v` *can* work, but it may not have everything set up. I'll update the README with details.

review: Approve

=== added file 'README'
--- README	1970-01-01 00:00:00 +0000
+++ README	2014-06-02 08:52:41 +0000
@@ -0,0 +1,15 @@
+System Image Updater
+====================
+
+Tools for system image based upgrades. For more details, see:
+https://wiki.ubuntu.com/ImageBasedUpgrades
+
+
+Testing
+-------
+
+To test locally run:
+
+  python3 -m nose2 -v
+
+  
 
=== modified file 'systemimage/gpg.py'
--- systemimage/gpg.py	2014-02-13 21:44:16 +0000
+++ systemimage/gpg.py	2014-06-02 08:52:41 +0000
@@ -28,7 +28,10 @@
 
 from contextlib import ExitStack
 from systemimage.config import config
-from systemimage.helpers import temporary_directory
+from systemimage.helpers import (
+    calculate_signature,
+    temporary_directory
+)
 
 
 class SignatureError(Exception):
@@ -54,19 +57,19 @@
         # want to be able to quickly and easily compare the file on disk
         # against the file on the server.
         with open(self.signature_path, 'rb') as fp:
-            self.signature_checksum =  hashlib.md5(fp.read()).hexdigest()
+            self.signature_checksum =  calculate_signature(fp, hashlib.md5)
         with open(self.data_path, 'rb') as fp:
-            self.data_checksum = hashlib.md5(fp.read()).hexdigest()
+            self.data_checksum = calculate_signature(fp, hashlib.md5)
         self.keyring_checksums = []
         for path in self.keyrings:
             with open(path, 'rb') as fp:
-                checksum = hashlib.md5(fp.read()).hexdigest()
+                checksum = calculate_signature(fp, hashlib.md5)
                 self.keyring_checksums.append(checksum)
         if self.blacklist is None:
             self.blacklist_checksum = None
         else:
             with open(self.blacklist, 'rb') as fp:
-                self.blacklist_checksum = hashlib.md5(fp.read()).hexdigest()
+                self.blacklist_checksum = calculate_signature(fp, hashlib.md5)
 
     def __str__(self):
         if self.blacklist is None:
 
=== modified file 'systemimage/helpers.py'
--- systemimage/helpers.py	2014-02-13 18:16:17 +0000
+++ systemimage/helpers.py	2014-06-02 08:52:41 +0000
@@ -49,6 +49,21 @@
 LAST_UPDATE_FILE = '/userdata/.last_update'
 UNIQUE_MACHINE_ID_FILE = '/var/lib/dbus/machine-id'
 DEFAULT_DIRMODE = 0o02700
+SIGNATURE_CHUNK_SIZE = 1*1024*1024
+
+
+def calculate_signature(fp, hash_class):
+    """
+    Calculate the hexdigest hash signature for the open file 'fp' using
+    the given hash_class in a memory efficient way
+    """
+    hash = hash_class()
+    while True:
+        chunk = fp.read(SIGNATURE_CHUNK_SIZE)
+        if not chunk:
+            break
+        hash.update(chunk)
+    return hash.hexdigest()
 
 
 def safe_remove(path):
 
=== modified file 'systemimage/state.py'
--- systemimage/state.py	2014-05-28 20:18:42 +0000
+++ systemimage/state.py	2014-06-02 08:52:41 +0000
@@ -39,7 +39,7 @@
 from systemimage.download import DBusDownloadManager, Record
 from systemimage.gpg import Context, SignatureError
 from systemimage.helpers import (
-    atomic, makedirs, safe_remove, temporary_directory)
+    atomic, calculate_signature, makedirs, safe_remove, temporary_directory)
 from systemimage.index import Index
 from systemimage.keyring import KeyringError, get_keyring
 from urllib.parse import urljoin
@@ -69,7 +69,7 @@
     if checksum is None:
         return True
     with open(txt, 'rb') as fp:
-        got = hashlib.sha256(fp.read()).hexdigest()
+        got = calculate_signature(fp, hashlib.sha256)
         return got == checksum
 
 
@@ -508,7 +508,7 @@
             # Verify the checksums.
             for dst, checksum in checksums:
                 with open(dst, 'rb') as fp:
-                    got = hashlib.sha256(fp.read()).hexdigest()
+                    got = calculate_signature(fp, hashlib.sha256)
                     if got != checksum:
                         raise ChecksumError(dst, got, checksum)
             # Everything is fine so nothing needs to be cleared.
 
=== modified file 'systemimage/tests/test_helpers.py'
--- systemimage/tests/test_helpers.py	2014-02-13 18:16:17 +0000
+++ systemimage/tests/test_helpers.py	2014-06-02 08:52:41 +0000
@@ -24,15 +24,18 @@
 
 
 import os
+import hashlib
 import logging
+import tempfile
 import unittest
 
 from contextlib import ExitStack
 from datetime import datetime, timedelta
 from systemimage.config import Configuration, config
 from systemimage.helpers import (
-    Bag, as_loglevel, as_object, as_timedelta, last_update_date,
-    phased_percentage, temporary_directory, version_detail)
+    Bag, as_loglevel, as_object, as_timedelta, calculate_signature,
+    last_update_date, phased_percentage, SIGNATURE_CHUNK_SIZE,
+    temporary_directory, version_detail)
 from systemimage.testing.helpers import configuration, data_path, touch_build
 from unittest.mock import patch
 
@@ -255,3 +258,15 @@
             self.assertEqual(phased_percentage(reset=True), 81)
             # The next one will have a different value.
             self.assertEqual(phased_percentage(), 17)
+
+
+class TestSignature(unittest.TestCase):
+    def test_calculcate_signature(self):
+        with tempfile.TemporaryFile() as fp:
+            # ensure the file is bigger than chunk size
+            fp.write(b"\0" * (SIGNATURE_CHUNK_SIZE+1))
+            fp.seek(0)
+            hash1 = calculate_signature(fp, hashlib.sha256)
+            fp.seek(0)
+            hash2 = hashlib.sha256(fp.read()).hexdigest()
+            self.assertEqual(hash1, hash2)