A FAFT test case for kernel A corruption.

This test corrupts kernel A and checks for kernel B on the next boot.
It will fail if kernel verification mis-behaved.

BUG=chromium-os:19710
TEST=run_remote_tests.sh --remote=$REMOTE_IP -a \
     "servo_vid=0x18d1 servo_pid=0x5001" firmware_CorruptKernelA

Change-Id: Ie41c1cb51216c9d830478335abd47c193a8f9dcc
Reviewed-on: https://gerrit.chromium.org/gerrit/10841
Commit-Ready: Tom Wai-Hong Tam <waihong@chromium.org>
Reviewed-by: Tom Wai-Hong Tam <waihong@chromium.org>
Tested-by: Tom Wai-Hong Tam <waihong@chromium.org>
diff --git a/client/cros/faft_client.py b/client/cros/faft_client.py
index e17f166..09a87ab 100644
--- a/client/cros/faft_client.py
+++ b/client/cros/faft_client.py
@@ -124,6 +124,26 @@
                 'crossystem %s' % key)[0]
 
 
+    def get_root_dev(self):
+        """Get the name of root device without partition number.
+
+        Returns:
+            A string of the root device without partition number.
+        """
+        self._chromeos_interface.log('Requesting get root device')
+        return self._chromeos_interface.get_root_dev()
+
+
+    def get_root_part(self):
+        """Get the name of root device with partition number.
+
+        Returns:
+            A string of the root device with partition number.
+        """
+        self._chromeos_interface.log('Requesting get root part')
+        return self._chromeos_interface.get_root_part()
+
+
     def set_try_fw_b(self):
         """Set 'Try Frimware B' flag in crossystem."""
         self._chromeos_interface.log('Requesting restart with firmware B')
@@ -158,6 +178,28 @@
         self._flashrom_handler.restore_firmware(section)
 
 
+    @allow_multiple_section_input
+    def corrupt_kernel(self, section):
+        """Corrupt the requested kernel section.
+
+        Args:
+            section: A kernel section, either 'a' or 'b'.
+        """
+        self._chromeos_interface.log('Corrupting kernel %s' % section)
+        self._kernel_handler.corrupt_kernel(section)
+
+
+    @allow_multiple_section_input
+    def restore_kernel(self, section):
+        """Restore the requested kernel section (previously corrupted).
+
+        Args:
+            section: A kernel section, either 'a' or 'b'.
+        """
+        self._chromeos_interface.log('restoring kernel %s' % section)
+        self._kernel_handler.restore_kernel(section)
+
+
     def cleanup(self):
         """Cleanup for the RPC server. Currently nothing."""
         pass
diff --git a/server/cros/faftsequence.py b/server/cros/faftsequence.py
index b0d8e37..35b1acb 100644
--- a/server/cros/faftsequence.py
+++ b/server/cros/faftsequence.py
@@ -63,6 +63,14 @@
     """
     version = 1
 
+
+    # Mapping of partition number of kernel and rootfs.
+    KERNEL_MAP = {'a':'2', 'b':'4', '2':'2', '4':'4', '3':'2', '5':'4'}
+    ROOTFS_MAP = {'a':'3', 'b':'5', '2':'3', '4':'5', '3':'3', '5':'5'}
+    OTHER_KERNEL_MAP = {'a':'4', 'b':'2', '2':'4', '4':'2', '3':'4', '5':'2'}
+    OTHER_ROOTFS_MAP = {'a':'5', 'b':'3', '2':'5', '4':'3', '3':'5', '5':'3'}
+
+
     _faft_template = None
     _faft_sequence = ()
 
@@ -191,6 +199,88 @@
         return True
 
 
+    def root_part_checker(self, expected_part):
+        """Check the partition number of the root device matched.
+
+        Args:
+          expected_part: A string containing the number of the expected root
+                         partition.
+
+        Returns:
+          True if the currect root  partition number matched; otherwise, False.
+        """
+        part = self.faft_client.get_root_part()
+        return self.ROOTFS_MAP[expected_part] == part[-1]
+
+
+    def copy_kernel_and_rootfs(self, from_part, to_part):
+        """Copy kernel and rootfs from from_part to to_part.
+
+        Args:
+          from_part: A string of partition number to be copied from.
+          to_part: A string of partition number to be copied to
+        """
+        root_dev = self.faft_client.get_root_dev()
+        self.faft_client.run_shell_command('dd if=%s of=%s bs=4M' %
+                (root_dev + self.KERNEL_MAP[from_part],
+                 root_dev + self.KERNEL_MAP[to_part]))
+        self.faft_client.run_shell_command('dd if=%s of=%s bs=4M' %
+                (root_dev + self.ROOTFS_MAP[from_part],
+                 root_dev + self.ROOTFS_MAP[to_part]))
+
+
+    def ensure_kernel_boot(self, part):
+        """Ensure the request kernel boot.
+
+        If not, it duplicates the current kernel to the requested kernel
+        and sets the requested higher priority to ensure it boot.
+
+        Args:
+            part: A string of kernel partition number or 'a'/'b'.
+        """
+        if not self.root_part_checker(part):
+            self.copy_kernel_and_rootfs(from_part=self.OTHER_KERNEL_MAP[part],
+                                        to_part=part)
+            self.reset_and_prioritize_kernel(part)
+            self.sync_and_hw_reboot()
+            self.wait_for_client_offline()
+            self.wait_for_client()
+
+
+    def setup_kernel(self, part):
+        """Setup for kernel test.
+
+        It makes sure both kernel A and B bootable and the current boot is
+        the requested kernel part.
+
+        Args:
+            part: A string of kernel partition number or 'a'/'b'.
+        """
+        self.ensure_kernel_boot(part)
+        self.copy_kernel_and_rootfs(from_part=part,
+                                    to_part=self.OTHER_KERNEL_MAP[part])
+        self.reset_and_prioritize_kernel(part)
+
+
+    def reset_and_prioritize_kernel(self, part):
+        """Make the requested partition highest priority.
+
+        This function also reset kerenl A and B to bootable.
+
+        Args:
+            part: A string of partition number to be prioritized.
+        """
+        root_dev = self.faft_client.get_root_dev()
+        # Reset kernel A and B to bootable.
+        self.faft_client.run_shell_command('cgpt add -i%s -P1 -S1 -T0 %s' %
+                (self.KERNEL_MAP['a'], root_dev))
+        self.faft_client.run_shell_command('cgpt add -i%s -P1 -S1 -T0 %s' %
+                (self.KERNEL_MAP['b'], root_dev))
+        # Set kernel part highest priority.
+        self.faft_client.run_shell_command('cgpt prioritize -i%s %s' %
+                (self.KERNEL_MAP[part], root_dev))
+
+
     def sync_and_hw_reboot(self):
         """Request the client sync and do a warm reboot.
 
diff --git a/server/site_tests/firmware_CorruptKernelA/control b/server/site_tests/firmware_CorruptKernelA/control
new file mode 100644
index 0000000..5b36be3
--- /dev/null
+++ b/server/site_tests/firmware_CorruptKernelA/control
@@ -0,0 +1,24 @@
+# Copyright (c) 2011 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.
+
+AUTHOR = "Chrome OS Team"
+NAME = "firmware_CorruptKernelA"
+PURPOSE = "Servo based kernel A corruption test"
+CRITERIA = "This test will fail if kernel verification mis-behaved"
+TIME = "LONG"
+TEST_CATEGORY = "Functional"
+TEST_CLASS = "firmware"
+TEST_TYPE = "server"
+
+DOC = """
+This test corrupts kernel A and checks for kernel B on the next boot.
+It will fail if kernel verification mis-behaved.
+"""
+
+def run_corruptkernela(machine):
+    host = hosts.create_host(machine)
+    job.run_test("firmware_CorruptKernelA", host=host, cmdline_args=args,
+                 use_faft=True, disable_sysinfo=True)
+
+parallel_simple(run_corruptkernela, machines)
diff --git a/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py b/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py
new file mode 100644
index 0000000..55d5c44
--- /dev/null
+++ b/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py
@@ -0,0 +1,43 @@
+# Copyright (c) 2011 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 autotest_lib.client.common_lib import error
+from autotest_lib.server.cros.faftsequence import FAFTSequence
+
+
+class firmware_CorruptKernelA(FAFTSequence):
+    """
+    Servo based kernel A corruption test.
+
+    This test corrupts kernel A and checks for kernel B on the next boot.
+    It will fail if kernel verification mis-behaved.
+    """
+    version = 1
+
+
+    def setup(self):
+        super(firmware_CorruptKernelA, self).setup()
+        self.setup_kernel('a')
+
+
+    def cleanup(self):
+        self.ensure_kernel_boot('a')
+        super(firmware_CorruptKernelA, self).cleanup()
+
+
+    def run_once(self, host=None):
+        self.register_faft_sequence((
+            {   # Step 1, corrupt kernel A
+                'state_checker': (self.root_part_checker, 'a'),
+                'userspace_action': (self.faft_client.corrupt_kernel, 'a'),
+            },
+            {   # Step 2, expected kernel B boot and restore kernel A
+                'state_checker': (self.root_part_checker, 'b'),
+                'userspace_action': (self.faft_client.restore_kernel, 'a'),
+            },
+            {   # Step 3, expected kernel A boot
+                'state_checker': (self.root_part_checker, 'a'),
+            },
+        ))
+        self.run_faft_sequence()