| #!/usr/bin/python |
| |
| # Copyright (c) 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. |
| |
| # This is a fully automated VM test of autoupdate. Here's what it does: |
| # - Downloads the latest dev channel image |
| # - Creates a VMware image based on that image |
| # - Creates an update for the image under test |
| # - Creates a copy of that dev channel image that's one number higher than |
| # the image under test (called the rollback image) |
| # - Creates an update for the rollback image |
| # - Launches a local HTTP server to pretend to be the AU server |
| # - Fires up the VM and waits for it to contact the AU server |
| # - AU server gives the VM the image under test update |
| # - Waits for the image to be installed, then reboots the VM |
| # - AU server gives the VM the rollback image |
| # - Waits for the image to be installed, then reboots the VM |
| # - Waits for the image to ping the AU server with the rollback image |
| # - Done! |
| |
| # Run this program by passing it a path to the build directory you want |
| # to test (this directory should contain rootfs.image). |
| |
| from xml.dom import minidom |
| from BaseHTTPServer import BaseHTTPRequestHandler |
| from BaseHTTPServer import HTTPServer |
| |
| import os |
| import re |
| import signal |
| import socket |
| import string |
| import subprocess |
| import sys |
| import tempfile |
| import threading |
| import time |
| |
| tmp_dir = '/tmp/au_vm_test' |
| scripts_dir = '../../scripts/' |
| |
| original_version = '0.0.0.0' |
| test_version = '0.0.0.0' |
| |
| # This class stores the state that the server is in and serves as glue |
| # between the AU server the control of the VMware process |
| class TestState(object): |
| # States we can be in |
| INIT, \ |
| INITIAL_WAIT, \ |
| TEST_DOWNLOAD, \ |
| TEST_WAIT, \ |
| ROLLBACK_DOWNLOAD, \ |
| ROLLBACK_WAIT = \ |
| xrange(6) |
| |
| def __init__(self, orig_vers, test_vers, rollback_vers, vm): |
| self.reboot_timeout_seconds = 60 * 4 |
| self.orig_vers = orig_vers |
| self.test_vers = test_vers |
| self.rollback_vers = rollback_vers |
| self.SetState(TestState.INIT) |
| self.vm = vm |
| |
| def Die(self, message): |
| print message |
| self.vm.Destroy() |
| # TODO(adlr): exit the entire process, not just this tread |
| sys.exit(1) |
| |
| def SetState(self, state): |
| self.state = state |
| |
| # Should be called to start the VM initially |
| def Start(self): |
| if self.state != TestState.INIT: |
| self.Die('Start called while in bad state') |
| self.SetState(TestState.INITIAL_WAIT) |
| self.vm.Launch() |
| # Kick off timer to wait for the AU ping |
| self.timer = threading.Timer(self.reboot_timeout_seconds, |
| self.StartupTimeout) |
| self.timer.start() |
| |
| def StartupTimeout(self): |
| self.Die('VM Failed to start and ping' + str(id(self))) |
| # TODO(adlr): exit the entire process, not just this tread |
| sys.exit(1) |
| |
| def FinishInstallTimeout(self): |
| self.vm.Shutdown() |
| self.vm.Launch() |
| self.timer = threading.Timer(60 * 5, # Sometimes VMWare is very slow |
| self.StartupTimeout) |
| self.timer.start() |
| |
| # Called by AU server when an update request comes in. Should return |
| # the version that the server should return to the AU client, or |
| # None if no update. |
| def HandleUpdateRequest(self, from_version): |
| print 'HandleUpdateRequest(%s) id:%s state:%s' % \ |
| (from_version, id(self), self.state) |
| ret = None |
| # Only cancel timer if we're waiting for the machine to startup |
| if self.timer != None and self.state == TestState.INITIAL_WAIT and \ |
| from_version == self.orig_vers: |
| print 'Successfully booted initial' |
| self.timer.cancel() |
| self.timer = None |
| elif self.timer != None and self.state == TestState.TEST_WAIT and \ |
| from_version == self.test_vers: |
| print 'Successfully booted test' |
| self.timer.cancel() |
| self.timer = None |
| elif self.timer != None and self.state == TestState.ROLLBACK_WAIT and \ |
| from_version == self.rollback_vers: |
| print 'Successfully booted rollback' |
| self.timer.cancel() |
| self.timer = None |
| print 'All done!' |
| # TODO(adlr): exit the entire process, not just this tread |
| sys.exit(0) |
| |
| # Pick the version to return |
| if from_version == self.orig_vers: |
| ret = self.test_vers |
| elif from_version == self.test_vers: |
| ret = self.rollback_vers |
| |
| # Checks to make sure we move through states correctly |
| if from_version == self.orig_vers: |
| if self.state != TestState.INITIAL_WAIT and \ |
| self.state != TestState.TEST_DOWNLOAD and \ |
| self.state != TestState.INITIAL_WAIT: |
| self.Die('Error: Request from %s while in state %s' % |
| (from_version, self.state)) |
| elif from_version == self.test_vers: |
| if self.state != TestState.TEST_WAIT and \ |
| self.state != TestState.ROLLBACK_DOWNLOAD and \ |
| self.state != TestState.ROLLBACK_WAIT: |
| self.Die('Error: Request from %s while in state %s' % |
| (from_version, self.state)) |
| else: |
| print 'odd version to be pinged from: %s' % from_version |
| print 'state is %s' % self.state |
| |
| # Update state if needed |
| if self.state == TestState.INITIAL_WAIT: |
| self.SetState(TestState.TEST_DOWNLOAD) |
| elif self.state == TestState.TEST_WAIT: |
| self.SetState(TestState.ROLLBACK_DOWNLOAD) |
| |
| if ret is not None: |
| return ret |
| print 'Ignoring update request while in state %s' % self.state |
| return '' |
| |
| # Called by AU server when the AU client has finished downloading an image |
| def ImageDownloadComplete(self): |
| print 'ImageDownloadComplete()' |
| valid_state = False |
| if self.state == TestState.TEST_DOWNLOAD: |
| valid_state = True |
| self.SetState(TestState.TEST_WAIT) |
| if self.state == TestState.ROLLBACK_DOWNLOAD: |
| valid_state = True |
| self.SetState(TestState.ROLLBACK_WAIT) |
| if not valid_state: |
| print 'Image download done called at invalid state' |
| # TODO(adlr): exit the entire process, not just this tread |
| sys.exit(1) |
| # Put a timer to reboot the VM |
| if self.timer is not None: |
| self.timer.cancel() |
| self.timer = None |
| self.timer = threading.Timer(self.reboot_timeout_seconds, |
| self.FinishInstallTimeout) |
| self.timer.start() |
| return |
| |
| # This subclass of HTTPServer contains info about the versions of |
| # software that the AU server should know about. The AUServerHandler |
| # object(s) will access this data. |
| class AUHTTPServer(HTTPServer): |
| def __init__(self, ip_port, klass): |
| HTTPServer.__init__(self, ip_port, klass) |
| self.update_info = {} |
| self.files = {} |
| |
| # For a given version of the software, the URL, size, and hash of the update |
| # that gives the user that version of the software. |
| def AddUpdateInfo(self, version, url, size, the_hash): |
| self.update_info[version] = (url, the_hash, size) |
| return |
| |
| # For a given path part of a url, return to the client the file at file_path |
| def AddServedFile(self, url_path, file_path): |
| self.files[url_path] = file_path |
| |
| def SetTestState(self, test_state): |
| self.test_state = test_state |
| |
| # This class handles HTTP requests. POST requests are when the client |
| # is pinging to see if there's an update. GET requests are to download |
| # an update. |
| class AUServerHandler(BaseHTTPRequestHandler): |
| def do_GET(self): |
| self.send_response(200) |
| self.end_headers() |
| print 'GET: %s' % self.path |
| |
| if self.server.files[self.path] != None: |
| print 'GET returning path %s' % self.server.files[self.path] |
| f = open(self.server.files[self.path]) |
| while True: |
| data = f.read(1024 * 1024 * 8) |
| if not data: |
| break |
| self.wfile.write(data) |
| self.wfile.flush() |
| f.close() |
| self.server.test_state.ImageDownloadComplete() |
| else: |
| print 'GET returning no path' |
| self.wfile.write(self.path + '\n') |
| return |
| |
| def do_POST(self): |
| # Parse the form data posted |
| post_length = int(self.headers.getheader('content-length')) |
| post_data = self.rfile.read(post_length) |
| |
| update_dom = minidom.parseString(post_data) |
| root = update_dom.firstChild |
| query = root.getElementsByTagName('o:app')[0] |
| client_version = query.getAttribute('version') |
| print 'Got update request from %s' % client_version |
| |
| # Send response |
| self.send_response(200) |
| self.end_headers() |
| |
| new_version = self.server.test_state.HandleUpdateRequest(client_version) |
| print 'Appropriate new version is: %s' % new_version |
| |
| if self.server.update_info[new_version] == None: |
| print 'Not sure how to serve reply for %s' % new_version |
| return |
| |
| payload = """<?xml version="1.0" encoding="UTF-8"?> |
| <gupdate xmlns="http://www.google.com/update2/response" protocol="2.0"> |
| <app appid="{87efface-864d-49a5-9bb3-4b050a7c227a}" status="ok"> |
| <ping status="ok"/> |
| <updatecheck |
| codebase="%s" |
| hash="%s" |
| needsadmin="false" |
| size="%s" |
| status="ok"/> |
| </app> |
| </gupdate> |
| """ % self.server.update_info[new_version] |
| |
| self.wfile.write(payload) |
| return |
| |
| # A wrapper for the vmplayer process. Can Launch/Shutdown a vm. |
| class vmplayer(object): |
| def __init__(self, filename): |
| self.filename = filename |
| |
| # Launch may (read: probably will) return before the OS has booted |
| def Launch(self): |
| self.process = subprocess.Popen(['/usr/bin/vmplayer', self.filename]) |
| self.running = True |
| |
| def Destroy(self): |
| if self.running: |
| self.Shutdown() |
| |
| # Shutdown will not return until the vmplayer process has fully terminated |
| # and any cleanup is done. |
| def Shutdown(self): |
| # Pretend user sent Ctrl-C to the vmplayer process |
| os.kill(self.process.pid, signal.SIGINT) |
| |
| # Wait while vmplayer saves the vm state... |
| self.process.wait() |
| |
| # Delete the saved vm state |
| # TODO(adlr): remove the state file from the disk |
| self.process = None |
| subprocess.check_call(['/bin/sed', '-i', '/checkpoint.vmState/d', |
| self.filename]) |
| self.running = False |
| |
| def MakePath(path): |
| subprocess.check_call(['/bin/mkdir', '-p', path]) |
| |
| def DownloadLatestImage(out_path): |
| url = 'http://codf30.jail.google.com/internal/archive/' + \ |
| 'x86-image-official/LATEST-dev-channel/image.zip' + \ |
| '.NOT_SAFE_FOR_USB_INSTALL' |
| url = 'http://www.corp.google.com/~adlr/adlr_test_orig.zip' |
| subprocess.check_call(['/usr/bin/wget', '-O', out_path, url]) |
| |
| def UnzipImage(path, directory): |
| subprocess.check_call(['/usr/bin/unzip', path, '-d', directory]) |
| |
| # Create a stateful partition with a text file that points the AU client |
| # to localhost at local_port (which is this very script). |
| def CreateDefaultStatefulPartition(local_ip, local_port, out_dir, out_file): |
| # Create sparse file for partition |
| part_size = 512 * 1024 * 1024 # bytes, so 500 MiB |
| |
| subprocess.check_call(["""#!/bin/bash |
| set -ex |
| OUT_FILE="%s" |
| STATE_DIR="%s" |
| SIZE="%s" |
| dd if=/dev/zero of="$OUT_FILE" bs=1 count=1 seek=$(($SIZE - 1)) |
| mkfs.ext4 -F "$OUT_FILE" |
| mkdir -p "$STATE_DIR" |
| sudo mount -o loop "$OUT_FILE" "$STATE_DIR" |
| sudo mkdir -p "$STATE_DIR/etc" |
| cat <<EOF |sudo dd of="$STATE_DIR/etc/lsb-release" |
| CHROMEOS_AUSERVER=http://%s:%s/update |
| HTTP_SERVER_OVERRIDE=true |
| EOF |
| for i in "$STATE_DIR/etc/lsb-release" "$STATE_DIR/etc" "$STATE_DIR"; do |
| sudo chown root:root "$i" |
| if [ -d "$i" ]; then |
| sudo chmod 0755 "$i" |
| else |
| sudo chmod 0644 "$i" |
| fi |
| done |
| sudo umount -d "$STATE_DIR" |
| """ % (out_file, out_dir + '/state', part_size, local_ip, local_port)], |
| shell=True) |
| return |
| |
| |
| def CreateVMForImage(image_dir): |
| subprocess.check_call([scripts_dir + 'image_to_vmware.sh', '--from', |
| image_dir, '--to', image_dir, '--state_image', |
| image_dir + '/state.image']) |
| return |
| |
| # Returns (size, hash, path) of the generated image.gz for the given rootfs |
| def CreateUpdateForImage(rootfs_image): |
| output = subprocess.Popen([scripts_dir + 'mk_memento_images.sh', |
| rootfs_image], |
| stdout=subprocess.PIPE).communicate()[0] |
| matches = re.search('Success. hash is ([^\n]+)', output) |
| the_hash = matches.group(1) |
| path = os.path.dirname(rootfs_image) + '/update.gz' |
| size = os.path.getsize(path) |
| return (size, the_hash, path) |
| |
| # Modify rootfs 'root_img' to have a new version new_version |
| def ModifyImageForRollback(root_img, new_version): |
| subprocess.check_call(["""#!/bin/bash |
| set -ex |
| DIR=$(mktemp -d) |
| sudo mount -o loop "%s" $DIR |
| # update versions in lsb-release |
| sudo sed -i \\ |
| -e 's/\\(^GOOGLE_RELEASE=\\|CHROMEOS_RELEASE_VERSION=\\).*/\\1%s/' \\ |
| -e 's/^\\(CHROMEOS_RELEASE_DESCRIPTION=.*\\)/\\1-ROLLBACK/' \\ |
| "$DIR"/etc/lsb-release |
| sudo umount -d $DIR |
| """ % (root_img, new_version)], shell=True) |
| |
| # Returns the release version of a rootfs (e.g. 0.6.39.201003241739-a1) |
| def GetVersionForRootfs(rootfs_image): |
| mount_dir = tempfile.mkdtemp() |
| subprocess.check_call(['sudo', 'mount', '-o', 'loop,ro', |
| rootfs_image, mount_dir]) |
| version = subprocess.Popen(['awk', '-F', '=', |
| '/GOOGLE_RELEASE=/{print $2}', |
| mount_dir + '/etc/lsb-release'], |
| stdout=subprocess.PIPE).communicate()[0].rstrip() |
| subprocess.check_call(['sudo', 'umount', '-d', mount_dir]) |
| subprocess.check_call(['sudo', 'rm', '-rf', mount_dir]) |
| return version |
| |
| # For a given version, increment the last number by 1. E.g.: |
| # IncrementVersionNumber('0.23.144.842') = '0.23.144.843' |
| def IncrementVersionNumber(version): |
| parts = version.split('.') |
| parts[-1] = str(int(parts[-1]) + 1) |
| return string.join(parts, '.') |
| |
| def UnpackRootfs(directory): |
| subprocess.check_call(["""#!/bin/bash -x |
| cd "%s" |
| ./unpack_partitions.sh chromiumos_image.bin |
| mv -f part_3 rootfs.image |
| """ % directory], shell=True) |
| |
| def main(): |
| if len(sys.argv) != 2: |
| print 'usage: %s path/to/new/image/dir' % sys.argv[0] |
| sys.exit(1) |
| orig_dir = tmp_dir + '/orig' |
| new_dir = sys.argv[1] |
| rollback_dir = tmp_dir + '/rollback' |
| |
| state_image = orig_dir + '/state.image' |
| port = 8080 |
| |
| MakePath(tmp_dir) |
| |
| # Download latest dev channel release |
| orig_zip = tmp_dir + '/orig.zip' |
| DownloadLatestImage(orig_zip) |
| UnzipImage(orig_zip, orig_dir) |
| UnpackRootfs(orig_dir) |
| orig_version = GetVersionForRootfs(orig_dir + '/rootfs.image') |
| print 'Have original image at version: %s' % orig_version |
| |
| # Create new AU image |
| print 'Creating update.gz for test image' |
| UnpackRootfs(new_dir) |
| new_update_details = CreateUpdateForImage(new_dir + '/rootfs.image') |
| new_version = GetVersionForRootfs(new_dir + '/rootfs.image') |
| print 'Have test image at version: %s' % new_version |
| |
| # Create rollback image |
| rollback_version = IncrementVersionNumber(new_version) |
| print 'Creating rollback image' |
| UnzipImage(orig_zip, rollback_dir) |
| UnpackRootfs(rollback_dir) |
| ModifyImageForRollback(rollback_dir + '/rootfs.image', rollback_version) |
| print 'Creating update.gz for rollback image' |
| rollback_update_details = CreateUpdateForImage(rollback_dir + '/rootfs.image') |
| print 'Have rollback image at version: %s' % rollback_version |
| |
| CreateDefaultStatefulPartition(socket.gethostname(), port, orig_dir, |
| state_image) |
| CreateVMForImage(orig_dir) |
| |
| player = vmplayer(orig_dir + '/chromeos.vmx') |
| |
| test_state = TestState(orig_version, new_version, rollback_version, player) |
| |
| server = AUHTTPServer((socket.gethostname(), port), AUServerHandler) |
| |
| base_url = 'http://%s:%s' % (socket.gethostname(), port) |
| |
| server.SetTestState(test_state) |
| server.AddUpdateInfo(new_version, base_url + '/' + new_version, |
| new_update_details[0], new_update_details[1]) |
| server.AddUpdateInfo(rollback_version, base_url + '/' + rollback_version, |
| rollback_update_details[0], rollback_update_details[1]) |
| server.AddServedFile('/' + new_version, new_update_details[2]) |
| server.AddServedFile('/' + rollback_version, rollback_update_details[2]) |
| |
| test_state.Start() |
| print 'Starting server, use <Ctrl-C> to stop' |
| server.serve_forever() |
| |
| if __name__ == '__main__': |
| main() |