restore file contexts for dev and test image.

Dev and test image has more files being installed, or modified later
than base_image_util restores the file contexts.

We need to restore SELinux file contexts every time we modify an image
to make sure all files have their contexts correct.

BUG=b:112616150
TEST=build_image

Change-Id: Ie8b7a2684ec48dc7a10ac3f63f546b895faa20d2
Reviewed-on: https://chromium-review.googlesource.com/1174587
Commit-Ready: Qijiang Fan <fqj@google.com>
Tested-by: Qijiang Fan <fqj@google.com>
Reviewed-by: Kenny Root <kroot@google.com>
diff --git a/build_image b/build_image
index f630430..ff1a4b0 100755
--- a/build_image
+++ b/build_image
@@ -120,6 +120,7 @@
 . "${BUILD_LIBRARY_DIR}/base_image_util.sh" || exit 1
 . "${BUILD_LIBRARY_DIR}/dev_image_util.sh" || exit 1
 . "${BUILD_LIBRARY_DIR}/test_image_util.sh" || exit 1
+. "${BUILD_LIBRARY_DIR}/selinux_util.sh" || exit 1
 
 parse_build_image_args
 
diff --git a/build_library/base_image_util.sh b/build_library/base_image_util.sh
index 9766391..2a957a0 100755
--- a/build_library/base_image_util.sh
+++ b/build_library/base_image_util.sh
@@ -401,19 +401,7 @@
     create_dev_install_lists "${root_fs_dir}"
   fi
 
-  # Restore the extended attributes of necessary files.
-  local selinux_config="${BOARD_ROOT}/etc/selinux/config"
-  if [[ -e "${selinux_config}" ]]; then
-    local selinux_type="$(source "${selinux_config}" && echo "${SELINUXTYPE}")"
-    local file_contexts="${BOARD_ROOT}/etc/selinux/${selinux_type}/contexts/files/file_contexts"
-    # If the selinux_config file exists, file_contexts must also.
-    if ! [[ -e "${file_contexts}" ]]; then
-      local err_msg="The SELinux config file exists at ${selinux_config}, "
-      err_msg+="but an SELinux context file not found at ${file_contexts}."
-      die_notrace "${err_msg}"
-    fi
-    sudo /sbin/setfiles -r "${root_fs_dir}" "${file_contexts}" "${root_fs_dir}"
-  fi
+  restore_fs_contexts "${BOARD_ROOT}" "${root_fs_dir}"
 
   # Zero rootfs free space to make it more compressible so auto-update
   # payloads become smaller.
diff --git a/build_library/dev_image_util.sh b/build_library/dev_image_util.sh
index f9eee75..90535db 100755
--- a/build_library/dev_image_util.sh
+++ b/build_library/dev_image_util.sh
@@ -129,6 +129,8 @@
 
   setup_etc_shadow "${root_fs_dir}"
 
+  restore_fs_contexts "${BOARD_ROOT}" "${root_fs_dir}"
+
   info "Developer image built and stored at ${image_name}"
 
   unmount_image
diff --git a/build_library/selinux_util.sh b/build_library/selinux_util.sh
new file mode 100644
index 0000000..eefb73b
--- /dev/null
+++ b/build_library/selinux_util.sh
@@ -0,0 +1,23 @@
+# Copyright 2018 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.
+
+restore_fs_contexts() {
+  local board_root="$1"
+  local rootfs="$2"
+
+  # Restore the extended attributes of necessary files.
+  local selinux_config="${board_root}/etc/selinux/config"
+  if [[ -e "${selinux_config}" ]]; then
+    info "Restoring SELinux file context."
+    local selinux_type="$(source "${selinux_config}" && echo "${SELINUXTYPE}")"
+    local file_contexts="${board_root}/etc/selinux/${selinux_type}/contexts/files/file_contexts"
+    # If the selinux_config file exists, file_contexts must also.
+    if ! [[ -e "${file_contexts}" ]]; then
+      local err_msg="The SELinux config file exists at ${selinux_config}, "
+      err_msg+="but an SELinux context file not found at ${file_contexts}."
+      die_notrace "${err_msg}"
+    fi
+    sudo /sbin/setfiles -r "${rootfs}" "${file_contexts}" "${rootfs}"
+  fi
+}
diff --git a/build_library/test_image_util.sh b/build_library/test_image_util.sh
index 4eebb78..babfef3 100755
--- a/build_library/test_image_util.sh
+++ b/build_library/test_image_util.sh
@@ -40,6 +40,8 @@
   # Re-run ldconfig to fix /etc/ld.so.cache.
   run_ldconfig "${root_fs_dir}"
 
+  restore_fs_contexts "${BOARD_ROOT}" "${root_fs_dir}"
+
   unmount_image
   trap - EXIT