Reworked devserver so that update images generated are cached in directories named after

autoupdate.py

 # Copyright (c) 2009-2010 The Chromium OS Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 from buildutil import BuildObject
 from xml.dom import minidom
+
 import cherrypy
 import os
 import shutil
 import subprocess
-import tempfile
 import time
 
 
 def _LogMessage(message):
   cherrypy.log(message, 'UPDATE')
 
 
 class Autoupdate(BuildObject):
   """Class that contains functionality that handles Chrome OS update pings.
 
   Members:
     serve_only: Serve images from a pre-built image.zip file.  static_dir
       must be set to the location of the image.zip.
     factory_config: Path to the factory config file if handling factory
       requests.
     use_test_image: Use chromiumos_test_image.bin rather than the standard.
     static_url_base: base URL, other than devserver, for update images.
     client_prefix: The prefix for the update engine client.
     forced_image: Path to an image to use for all updates.
   """
 
   def __init__(self, serve_only=None, test_image=False, urlbase=None,
                factory_config_path=None, client_prefix=None, forced_image=None,
-               use_cached=False, port=8080, src_image='', vm=False, board=None,
+               port=8080, src_image='', vm=False, board=None,
                *args, **kwargs):
     super(Autoupdate, self).__init__(*args, **kwargs)
     self.serve_only = serve_only
     self.factory_config = factory_config_path
     self.use_test_image = test_image
     if urlbase:
       self.urlbase = urlbase
     else:
       self.urlbase = None
 
     self.client_prefix = client_prefix
     self.forced_image = forced_image
-    self.use_cached = use_cached
     self.src_image = src_image
     self.vm = vm
     self.board = board
     self.crosutils = os.path.join(os.path.dirname(__file__), '../../scripts')
 
   def _GetSecondsSinceMidnight(self):
     """Returns the seconds since midnight as a decimal value."""
     now = time.localtime()
     return now[3] * 3600 + now[4] * 60 + now[5]
 
   def _GetDefaultBoardID(self):
     """Returns the default board id stored in .default_board."""
     board_file = '%s/.default_board' % (self.scripts_dir)
     try:
       return open(board_file).read()
     except IOError:
       return 'x86-generic'
 
   def _GetLatestImageDir(self, board_id):
     """Returns the latest image dir based on shell script."""
     cmd = '%s/get_latest_image.sh --board %s' % (self.scripts_dir, board_id)
     return os.popen(cmd).read().strip()
 
   def _GetVersionFromDir(self, image_dir):
     """Returns the version of the image based on the name of the directory."""
     latest_version = os.path.basename(image_dir)
     return latest_version.split('-')[0]
 
   def _CanUpdate(self, client_version, latest_version):
-    """Returns true if the latest_version is greater than the client_version."""
+    """Returns true if the latest_version is greater than the client_version.
+    """
     client_tokens = client_version.replace('_', '').split('.')
     latest_tokens = latest_version.replace('_', '').split('.')
     _LogMessage('client version %s latest version %s'
-              % (client_version, latest_version))
+                % (client_version, latest_version))
     for i in range(4):
       if int(latest_tokens[i]) == int(client_tokens[i]):
         continue
       return int(latest_tokens[i]) > int(client_tokens[i])
     return False
 
   def _UnpackZip(self, image_dir):
     """Unpacks an image.zip into a given directory."""
     image = os.path.join(image_dir, self._GetImageName())
     if os.path.exists(image):
       return True
     else:
       # -n, never clobber an existing file, in case we get invoked
       # simultaneously by multiple request handlers. This means that
       # we're assuming each image.zip file lives in a versioned
       # directory (a la Buildbot).
       return os.system('cd %s && unzip -n image.zip' % image_dir) == 0
 
   def _GetImageName(self):
     """Returns the name of the image that should be used."""
     if self.use_test_image:
       image_name = 'chromiumos_test_image.bin'
     else:
       image_name = 'chromiumos_image.bin'
     return image_name
 
-  def _IsImageNewerThanCached(self, image_path, cached_file_path):
-    """Returns true if the image is newer than the cached image."""
-    if os.path.exists(cached_file_path) and os.path.exists(image_path):
-      _LogMessage('Usable cached image found at %s.' % cached_file_path)
-      return os.path.getmtime(image_path) > os.path.getmtime(cached_file_path)
-    elif not os.path.exists(cached_file_path) and not os.path.exists(image_path):
-      raise Exception('Image does not exist and cached image missing')
-    else:
-      # Only one is missing, figure out which one.
-      if os.path.exists(image_path):
-        _LogMessage('No cached image found - image generation required.')
-        return True
-      else:
-        _LogMessage('Cached image found to serve at %s.' % cached_file_path)
-        return False
-
   def _GetSize(self, update_path):
     """Returns the size of the file given."""
     return os.path.getsize(update_path)
 
   def _GetHash(self, update_path):
     """Returns the sha1 of the file given."""
     cmd = ('cat %s | openssl sha1 -binary | openssl base64 | tr \'\\n\' \' \';'
            % update_path)
     return os.popen(cmd).read().rstrip()
 
@@ skipping 10 matching lines @@
   # it takes advantage of reduced I/O and multiple processors. Something like:
   # % tee < FILE > /dev/null \
   #     >( openssl dgst -sha256 -binary | openssl base64 ) \
   #     >( openssl sha1 -binary | openssl base64 )
   def _GetSHA256(self, update_path):
     """Returns the sha256 of the file given."""
     cmd = ('cat %s | openssl dgst -sha256 -binary | openssl base64' %
            update_path)
     return os.popen(cmd).read().rstrip()
 
+  def _GetMd5(self, update_path):
+    """Returns the md5 checksum of the file given."""
+    cmd = ("md5sum %s | awk '{print $1}'" % update_path)
+    return os.popen(cmd).read().rstrip()
+
+  def _Symlink(self, source, dest):
+    """Creates a symlink at dest to source"""
+    if os.path.exists(dest):
+      os.remove(dest)
+    os.symlink(source, dest)
+
   def GetUpdatePayload(self, hash, sha256, size, url, is_delta_format):
     """Returns a payload to the client corresponding to a new update.
 
     Args:
       hash: hash of update blob
       sha256: SHA-256 hash of update blob
       size: size of update blob
       url: where to find update blob
     Returns:
       Xml string to be passed back to client.
@@ skipping 26 matching lines @@
       < gupdate xmlns = "http://www.google.com/update2/response" protocol = "2.0" >
         < daystart elapsed_seconds = "%s" />
         < app appid = "{%s}" status = "ok" >
           < ping status = "ok" />
           < updatecheck status = "noupdate" />
         </ app >
       </ gupdate >
     """
     return payload % (self._GetSecondsSinceMidnight(), self.app_id)
 
-  def GenerateUpdateFile(self, image_path):
+  def GenerateUpdateFile(self, src_image, image_path, output_dir):
     """Generates an update gz given a full path to an image.
 
     Args:
       image_path: Full path to image.
     Returns:
       Path to created update_payload or None on error.
     """
-    image_dir = os.path.dirname(image_path)
-    update_path = os.path.join(image_dir, 'update.gz')
+    update_path = os.path.join(output_dir, 'update.gz')
     patch_kernel_flag = '--patch_kernel'
     _LogMessage('Generating update image %s' % update_path)
 
     # Don't patch the kernel for vm images as they don't need the patch.
     if self.vm:
       patch_kernel_flag = ''
 
     mkupdate_command = (
         '%s/cros_generate_update_payload --image="%s" --output="%s" '
         '%s --noold_style --src_image="%s"' % (
             self.scripts_dir, image_path, update_path, patch_kernel_flag,
-            self.src_image))
+            src_image))
     _LogMessage(mkupdate_command)
     if os.system(mkupdate_command) != 0:
       _LogMessage('Failed to create base update file')
       return None
 
     return update_path
 
-  def GenerateStatefulFile(self, image_path):
-    """Generates a stateful update given a full path to an image.
+  def GenerateStatefulFile(self, image_path, output_dir):
+    """Generates a stateful update payload given a full path to an image.
 
     Args:
       image_path: Full path to image.
     Returns:
-      Path to created stateful update_payload.
+      Path to created stateful update_payload or None on error.
     Raises:
       A subprocess exception if the update generator fails to generate a
       stateful payload.
     """
-    work_dir = os.path.dirname(image_path)
-    output_gz = os.path.join(work_dir, 'stateful.tgz')
+    output_gz = os.path.join(output_dir, 'stateful.tgz')
     subprocess.check_call(
         ['%s/cros_generate_stateful_update_payload' % self.crosutils,
          '--image=%s' % image_path,
-         '--output_dir=%s' % work_dir,
+         '--output_dir=%s' % output_dir,
         ])
     return output_gz
 
-  def MoveImagesToStaticDir(self, update_path, stateful_update_path,
-                            static_image_dir):
-    """Moves gz files from their directories to serving directories.
+  def FindCachedUpdateImageSubDir(self, src_image, dest_image):
+    """Find directory to store a cached update.
+
+       Given one, or two images for an update, this finds which
+       cache directory should hold the update files, even if they don't exist
+       yet. The directory will be inside static_image_dir, and of the form:
+
+       Non-delta updates:
+         cache/12345678
+
+       Delta updates:
+         cache/12345678_12345678
+       """
+    # If there is no src, we only have an image file, check image for changes
+    if not src_image:
+      return os.path.join('cache', self._GetMd5(dest_image))
+
+    # If we have src and dest, we are a delta, and check both for changes
+    return os.path.join('cache',
+                        "%s_%s" % (self._GetMd5(src_image),
+                                   self._GetMd5(dest_image)))
+
+  def GenerateUpdateImage(self, src_image, image_path, output_dir):
+    """Force generates an update payload based on the given image_path.
 
     Args:
-      update_path: full path to main update gz.
-      stateful_update_path: full path to stateful partition gz.
-      static_image_dir: where to put files.
+      src_image: image we are updating from (Null/empty for non-delta)
+      image_path: full path to the image.
+      output_dir: the directory to write the update payloads in
     Returns:
-      Returns True if the files were moved over successfully.
+      update and stateful payload tuple with full file names
     """
-    try:
-      shutil.copy(update_path, static_image_dir)
-      shutil.copy(stateful_update_path, static_image_dir)
-      os.remove(update_path)
-      os.remove(stateful_update_path)
-    except Exception:
-      _LogMessage('Failed to move %s and %s to %s' % (update_path,
-                                                    stateful_update_path,
-                                                    static_image_dir))
-      return False
+    update_file = None
+    stateful_update_file = None
 
-    return True
+    # Actually do the generation
+    _LogMessage('Generating update for image %s' % image_path)
+    update_file = self.GenerateUpdateFile(src_image,
+                                          image_path,
+                                          output_dir)
 
-  def GenerateUpdateImage(self, image_path, move_to_static_dir=False,
-                          static_image_dir=None):
-    """Generates an update payload based on the given image_path.
+    if update_file:
+      stateful_update_file = self.GenerateStatefulFile(image_path,
+                                                       output_dir)
+
+    if update_file and stateful_update_file:
+      return update_file, stateful_update_file
+
+    _LogMessage('Failed to generate update')
+
+    # Cleanup incomplete files, if they exist
+    if update_file and os.path.exists(update_file):
+      os.remove(update_file)
+
+    return None
+
+  def GenerateUpdateImageWithCache(self, image_path, static_image_dir):
+    """Force generates an update payload based on the given image_path.
 
     Args:
       image_path: full path to the image.
-      move_to_static_dir: Moves the files from their dir to the static dir.
       static_image_dir: the directory to move images to after generating.
     Returns:
-      True if the update payload was created successfully.
+      update filename (not directory) relative to static_image_dir on success,
+        or None
     """
-    _LogMessage('Generating update for image %s' % image_path)
-    update_path = self.GenerateUpdateFile(image_path)
-    if update_path:
-      stateful_update_path = self.GenerateStatefulFile(image_path)
-      if move_to_static_dir:
-        return self.MoveImagesToStaticDir(update_path, stateful_update_path,
-                                          static_image_dir)
-        return True
+    _LogMessage('Generating update for src %s image %s' % (self.src_image,
+                                                           image_path))
 
-    _LogMessage('Failed to generate update')
-    return False
+    # Which sub_dir of static_image_dir should hold our cached update image
+    cache_sub_dir = self.FindCachedUpdateImageSubDir(self.src_image, image_path)
+    _LogMessage('Caching in sub_dir "%s"' % cache_sub_dir)
+
+    # The cached payloads exist in a cache dir
+    cache_update_payload = os.path.join(static_image_dir,
+                                        cache_sub_dir,
+                                        'update.gz')
+    cache_stateful_payload = os.path.join(static_image_dir,
+                                          cache_sub_dir,
+                                          'stateful.tgz')
+
+    # The final results (symlinks?) exist directly in static
+    update_payload = os.path.join(static_image_dir,
+                                  'update.gz')
+    stateful_payload = os.path.join(static_image_dir,
+                                    'stateful.tgz')
+
+    # If there isn't a cached payload, make one
+    if not os.path.exists(cache_update_payload):
+      full_cache_dir = os.path.join(static_image_dir, cache_sub_dir)
+
+      # Create the directory for the cache values
+      if not os.path.exists(full_cache_dir):
+        os.makedirs(full_cache_dir)
+
+      payloads = self.GenerateUpdateImage(self.src_image,
+                                          image_path,
+                                          full_cache_dir)
+
+      if not payloads:
+        # Clean up cache dir if it's not valid
+        os.system("rm -rf %s" % os.path.join(static_image_dir, cache_sub_dir))
+        return None
+
+      # Verify they were created with the expected names
+      new_update_payload, new_stateful_payload = payloads
+
+      _LogMessage('"%s" "%s"' % (new_update_payload, cache_update_payload))
+      assert new_update_payload == cache_update_payload
+      _LogMessage('"%s" "%s"' % (new_stateful_payload, cache_stateful_payload))
+      assert new_stateful_payload == cache_stateful_payload
+
+    # If the generation worked, create symlinks
+    self._Symlink(cache_update_payload, update_payload)
+    self._Symlink(cache_stateful_payload, stateful_payload)
+
+    # return just the filename which is symlink in static_image_dir
+    return 'update.gz'
 
   def GenerateLatestUpdateImage(self, board_id, client_version,
-                                static_image_dir=None):
+                                static_image_dir):
     """Generates an update using the latest image that has been built.
 
     This will only generate an update if the newest update is newer than that
     on the client or client_version is 'ForcedUpdate'.
 
     Args:
       board_id: Name of the board.
       client_version: Current version of the client or 'ForcedUpdate'
       static_image_dir: the directory to move images to after generating.
     Returns:
-      True if the update payload was created successfully.
+      Name of the update image relative to static_image_dir or None
     """
     latest_image_dir = self._GetLatestImageDir(board_id)
     latest_version = self._GetVersionFromDir(latest_image_dir)
     latest_image_path = os.path.join(latest_image_dir, self._GetImageName())
 
     _LogMessage('Preparing to generate update from latest built image %s.' %
               latest_image_path)
 
      # Check to see whether or not we should update.
     if client_version != 'ForcedUpdate' and not self._CanUpdate(
         client_version, latest_version):
       _LogMessage('no update')
-      return False
+      return None
 
-    cached_file_path = os.path.join(static_image_dir, 'update.gz')
-    if (os.path.exists(cached_file_path) and
-        not self._IsImageNewerThanCached(latest_image_path, cached_file_path)):
-      return True
-
-    return self.GenerateUpdateImage(latest_image_path, move_to_static_dir=True,
-                                    static_image_dir=static_image_dir)
+    return self.GenerateUpdateImageWithCache(latest_image_path,
+                                             static_image_dir=static_image_dir)
 
   def GenerateImageFromZip(self, static_image_dir):
     """Generates an update from an image zip file.
 
     This method assumes you have an image.zip in directory you are serving
     from.  If this file is newer than a previously cached file, it will unzip
     this file, create a payload and serve it.
 
     Args:
       static_image_dir: Directory where the zip file exists.
     Returns:
-      True if the update payload was created successfully.
+      Name of the update payload relative to static_image_dir if successful.
     """
-    _LogMessage('Preparing to generate update from zip in %s.' % static_image_dir)
+    _LogMessage('Preparing to generate update from zip in %s.' %
+                static_image_dir)
     image_path = os.path.join(static_image_dir, self._GetImageName())
-    cached_file_path = os.path.join(static_image_dir, 'update.gz')
     zip_file_path = os.path.join(static_image_dir, 'image.zip')
-    if not self._IsImageNewerThanCached(zip_file_path, cached_file_path):
-      return True
+
+    # TODO(dgarrett): Either work caching into this path before
+    #       we unpack, or remove zip support (sosa is considering).
+    #       It does currently cache, but after the unpack.
 
     if not self._UnpackZip(static_image_dir):
       _LogMessage('unzip image.zip failed.')
-      return False
+      return None
 
-    return self.GenerateUpdateImage(image_path, move_to_static_dir=False,
-                                    static_image_dir=None)
+    return self.GenerateUpdateImageWithCache(image_path,
+                                             static_image_dir=static_image_dir)
 
   def ImportFactoryConfigFile(self, filename, validate_checksums=False):
     """Imports a factory-floor server configuration file. The file should
     be in this format:
       config = [
         {
           'qual_ids': set([1, 2, 3, "x86-generic"]),
           'factory_image': 'generic-factory.gz',
           'factory_checksum': 'AtiI8B64agHVN+yeBAyiNMX3+HM=',
           'release_image': 'generic-release.gz',
@@ skipping 76 matching lines @@
       return self.GetNoUpdatePayload()
     url = '%s/static/%s' % (self.hostname, filename)
     is_delta_format = self._IsDeltaFormatFile(filename)
     _LogMessage('returning update payload ' + url)
     # Factory install is using memento updater which is using the sha-1 hash so
     # setting sha-256 to an empty string.
     return self.GetUpdatePayload(checksum, '', size, url, is_delta_format)
 
   def GenerateUpdatePayloadForNonFactory(self, board_id, client_version,
                                          static_image_dir):
-    """Generates an update for non-factory and returns True on success."""
-    if self.use_cached and os.path.exists(os.path.join(static_image_dir,
-                                                       'update.gz')):
-      _LogMessage('Using cached image regardless of timestamps.')
-      return True
+    """Generates an update for non-factory image.
+
+       Returns:
+         file name relative to static_image_dir on success.
+    """
+    if self.forced_image:
+      return self.GenerateUpdateImageWithCache(
+          self.forced_image,
+          static_image_dir=static_image_dir)
+    elif self.serve_only:
+      return self.GenerateImageFromZip(static_image_dir)
     else:
-      if self.forced_image:
-        has_built_image = self.GenerateUpdateImage(
-            self.forced_image, move_to_static_dir=True,
-            static_image_dir=static_image_dir)
-        return has_built_image
-      elif self.serve_only:
-        return self.GenerateImageFromZip(static_image_dir)
-      else:
-        if board_id:
-          return self.GenerateLatestUpdateImage(board_id,
-                                                client_version,
-                                                static_image_dir)
+      if board_id:
+        return self.GenerateLatestUpdateImage(board_id,
+                                              client_version,
+                                              static_image_dir)
 
-        _LogMessage('You must set --board for pre-generating latest update.')
-        return False
+      _LogMessage('You must set --board for pre-generating latest update.')
+      return None
 
   def PreGenerateUpdate(self):
-    """Pre-generates an update.  Returns True on success."""
+    """Pre-generates an update.  Returns True on success.
+    """
      # Does not work with factory config.
     assert(not self.factory_config)
     _LogMessage('Pre-generating the update payload.')
     # Does not work with labels so just use static dir.
     if self.GenerateUpdatePayloadForNonFactory(self.board, '0.0.0.0',
                                                self.static_dir):
-      # Force the devserver to use the pre-generated payload.
-      self.use_cached = True
       _LogMessage('Pre-generated update successfully.')
       return True
     else:
       _LogMessage('Failed to pre-generate update.')
       return False
 
   def HandleUpdatePing(self, data, label=None):
     """Handles an update ping from an update client.
 
     Args:
@@ skipping 41 matching lines @@
 
     # Separate logic as Factory requests have static url's that override
     # other options.
     if self.factory_config:
       return self.HandleFactoryRequest(board_id, channel)
     else:
       static_image_dir = self.static_dir
       if label:
         static_image_dir = os.path.join(static_image_dir, label)
 
-      if self.GenerateUpdatePayloadForNonFactory(board_id, client_version,
-                                                 static_image_dir):
-        filename = os.path.join(static_image_dir, 'update.gz')
+      payload_path = self.GenerateUpdatePayloadForNonFactory(board_id,
+                                                             client_version,
+                                                             static_image_dir)
+      if payload_path:
+        filename = os.path.join(static_image_dir, payload_path)
         hash = self._GetHash(filename)
         sha256 = self._GetSHA256(filename)
         size = self._GetSize(filename)
         is_delta_format = self._IsDeltaFormatFile(filename)
         if label:
-          url = '%s/%s/update.gz' % (static_urlbase, label)
+          url = '%s/%s/%s' % (static_urlbase, label, payload_path)
         else:
-          url = '%s/update.gz' % static_urlbase
+          url = '%s/%s' % (static_urlbase, payload_path)
 
         _LogMessage('Responding to client to use url %s to get image.' % url)
         return self.GetUpdatePayload(hash, sha256, size, url, is_delta_format)
       else:
         return self.GetNoUpdatePayload()