image service: generate recovery images

Add the ability to generate recovery images via the same Create
endpoint. This adds a bit of management around ensuring that
dependent image types (base) are also generated.

BUG=b:181231500
TEST=./run_tests.py

Change-Id: I0e5c9203987bc3e983db3e820aafc40fa6b50c72
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2743542
Reviewed-by: Alex Klein <saklein@chromium.org>
Commit-Queue: George Engelbrecht <engeg@google.com>
Tested-by: George Engelbrecht <engeg@google.com>
diff --git a/api/controller/image.py b/api/controller/image.py
index 7c3f8b3..a28bb0c 100644
--- a/api/controller/image.py
+++ b/api/controller/image.py
@@ -54,6 +54,7 @@
     constants.IMAGE_TYPE_FIRMWARE: _FIRMWARE_ID,
 }
 
+# Dict to describe the prerequisite built images for each VM image type.
 _VM_IMAGE_MAPPING = {
     _BASE_VM_ID: _IMAGE_MAPPING[_BASE_ID],
     _TEST_VM_ID: _IMAGE_MAPPING[_TEST_ID],
@@ -61,6 +62,11 @@
     _TEST_GUEST_VM_ID: _IMAGE_MAPPING[_TEST_ID],
 }
 
+# Dict to describe the prerequisite built images for each mod image type.
+_MOD_IMAGE_MAPPING = {
+    _RECOVERY_ID: _IMAGE_MAPPING[_BASE_ID],
+}
+
 # Supported image types for PushImage.
 SUPPORTED_IMAGE_TYPES = {
     common_pb2.IMAGE_TYPE_RECOVERY: constants.IMAGE_TYPE_RECOVERY,
@@ -69,10 +75,18 @@
     common_pb2.IMAGE_TYPE_ACCESSORY_USBPD: constants.IMAGE_TYPE_ACCESSORY_USBPD,
     common_pb2.IMAGE_TYPE_ACCESSORY_RWSIG: constants.IMAGE_TYPE_ACCESSORY_RWSIG,
     common_pb2.IMAGE_TYPE_BASE: constants.IMAGE_TYPE_BASE,
-    common_pb2.IMAGE_TYPE_GSC_FIRMWARE: constants.IMAGE_TYPE_GSC_FIRMWARE
+    common_pb2.IMAGE_TYPE_GSC_FIRMWARE: constants.IMAGE_TYPE_GSC_FIRMWARE,
 }
 
 
+def _add_image_to_proto(output_proto, path, image_type, board):
+  """Quick helper function to add a new image to the output proto."""
+  new_image = output_proto.images.add()
+  new_image.path = path
+  new_image.type = image_type
+  new_image.build_target.name = board
+
+
 def _CreateResponse(_input_proto, output_proto, _config):
   """Set output_proto success field on a successful Create response."""
   output_proto.success = True
@@ -84,7 +98,7 @@
 @validate.validation_complete
 @metrics.collect_metrics
 def Create(input_proto, output_proto, _config):
-  """Build an image.
+  """Build images.
 
   Args:
     input_proto (image_pb2.CreateImageRequest): The input message.
@@ -96,7 +110,7 @@
   # Build the base image if no images provided.
   to_build = input_proto.image_types or [_BASE_ID]
 
-  image_types, vm_types = _ParseImagesToCreate(to_build)
+  image_types, vm_types, mod_image_types = _ParseImagesToCreate(to_build)
   build_config = _ParseCreateBuildConfig(input_proto)
 
   # Sorted isn't really necessary here, but it's much easier to test.
@@ -109,22 +123,30 @@
     # Success -- we need to list out the images we built in the output.
     _PopulateBuiltImages(board, image_types, output_proto)
 
-    if vm_types:
-      for vm_type in vm_types:
-        is_test = vm_type in [_TEST_VM_ID, _TEST_GUEST_VM_ID]
-        try:
-          if vm_type in [_BASE_GUEST_VM_ID, _TEST_GUEST_VM_ID]:
-            vm_path = image.CreateGuestVm(board, is_test=is_test)
-          else:
-            vm_path = image.CreateVm(
-                board, disk_layout=build_config.disk_layout, is_test=is_test)
-        except image.ImageToVmError as e:
-          cros_build_lib.Die(e)
+    for vm_type in vm_types:
+      is_test = vm_type in [_TEST_VM_ID, _TEST_GUEST_VM_ID]
+      try:
+        if vm_type in [_BASE_GUEST_VM_ID, _TEST_GUEST_VM_ID]:
+          vm_path = image.CreateGuestVm(board, is_test=is_test)
+        else:
+          vm_path = image.CreateVm(
+              board, disk_layout=build_config.disk_layout, is_test=is_test)
+      except image.ImageToVmError as e:
+        cros_build_lib.Die(e)
 
-        new_image = output_proto.images.add()
-        new_image.path = vm_path
-        new_image.type = vm_type
-        new_image.build_target.name = board
+      _add_image_to_proto(output_proto, vm_path, vm_type, board)
+
+    for mod_type in mod_image_types:
+      if mod_type == _RECOVERY_ID:
+        base_image_path = _GetBaseImagePath(output_proto)
+        result = image.BuildRecoveryImage(board=board,
+                                          image_path=base_image_path)
+        if result.success:
+          _PopulateBuiltImages(board, [_IMAGE_MAPPING[mod_type]], output_proto)
+        else:
+          cros_build_lib.Die('Failed to create recovery image.')
+      else:
+        cros_build_lib.Die('_RECOVERY_ID is the only mod_image_type.')
 
     # Read metric events log and pipe them into output_proto.events.
     deserialize_metrics_log(output_proto.events, prefix=board)
@@ -141,27 +163,41 @@
 
     return controller.RETURN_CODE_UNSUCCESSFUL_RESPONSE_AVAILABLE
 
+def _GetBaseImagePath(output_proto):
+  """From image_pb2.CreateImageResult, return base image path or None."""
+  ret = None
+  for i in output_proto.images:
+    if i.type == _BASE_ID:
+      ret = i.path
+  return ret
+
 
 def _ParseImagesToCreate(to_build):
   """Helper function to parse the image types to build.
 
-  This function exists just to clean up the Create function.
+  This function expresses the dependencies of each image type and adds
+  the requisite image types if they're not explicitly defined.
 
   Args:
     to_build (list[int]): The image type list.
 
   Returns:
-    (set, set): The image and vm types, respectively, that need to be built.
+    (set, set, set): The image, vm, and mod_image types that need to be built.
   """
   image_types = set()
   vm_types = set()
+  mod_image_types = set()
   for current in to_build:
-    if current in _IMAGE_MAPPING:
-      image_types.add(_IMAGE_MAPPING[current])
-    elif current in _VM_IMAGE_MAPPING:
+    # Find out if it's a special case (vm, img mod), or just any old image.
+    if current in _VM_IMAGE_MAPPING:
       vm_types.add(current)
       # Make sure we build the image required to build the VM.
       image_types.add(_VM_IMAGE_MAPPING[current])
+    elif current in _MOD_IMAGE_MAPPING:
+      mod_image_types.add(current)
+      image_types.add(_MOD_IMAGE_MAPPING[current])
+    elif current in _IMAGE_MAPPING:
+      image_types.add(_IMAGE_MAPPING[current])
     else:
       # Not expected, but at least it will be obvious if this comes up.
       cros_build_lib.Die(
@@ -173,7 +209,7 @@
   if vm_types.issuperset({_BASE_VM_ID, _TEST_VM_ID}):
     cros_build_lib.Die('Cannot create more than one VM.')
 
-  return image_types, vm_types
+  return image_types, vm_types, mod_image_types
 
 
 def _ParseCreateBuildConfig(input_proto):
@@ -202,11 +238,7 @@
   for current in image_types:
     type_id = _IMAGE_MAPPING[current]
     path = os.path.join(base_path, constants.IMAGE_TYPE_TO_NAME[current])
-
-    new_image = output_proto.images.add()
-    new_image.path = path
-    new_image.type = type_id
-    new_image.build_target.name = board
+    _add_image_to_proto(output_proto, path, type_id, board)
 
 
 def _SignerTestResponse(_input_proto, output_proto, _config):