3 files changed, 65 insertions, 5 deletions

diff --git a/twinstall.cpp b/twinstall.cpp
index 041c78d63..8af709915 100644
--- a/twinstall.cpp
+++ b/twinstall.cpp
@@ -55,13 +55,44 @@ static int Run_Update_Binary(const char *path, ZipArchive *Zip, int* wipe_cache)
 
 	ret_val = mzExtractZipEntryToFile(Zip, binary_location, binary_fd);
 	close(binary_fd);
-	mzCloseZipArchive(Zip);
 
 	if (!ret_val) {
+		mzCloseZipArchive(Zip);
 		LOGERR("Could not extract '%s'\n", ASSUMED_UPDATE_BINARY_NAME);
 		return INSTALL_ERROR;
 	}
 
+	// If exists, extract file_contexts from the zip file
+	const ZipEntry* selinx_contexts = mzFindZipEntry(Zip, "file_contexts");
+	if (selinx_contexts == NULL) {
+		mzCloseZipArchive(Zip);
+		LOGINFO("Zip does not contain SELinux file_contexts file in its root.\n");
+	} else {
+		string output_filename = "/file_contexts";
+		LOGINFO("Zip contains SELinux file_contexts file in its root. Extracting to %s\n", output_filename.c_str());
+		// Delete any file_contexts
+		if (TWFunc::Path_Exists(output_filename) && unlink(output_filename.c_str()) != 0) {
+			LOGINFO("Unable to unlink '%s'\n", output_filename.c_str());
+		}
+
+		int file_contexts_fd = creat(output_filename.c_str(), 0644);
+		if (file_contexts_fd < 0) {
+			mzCloseZipArchive(Zip);
+			LOGERR("Could not extract file_contexts to '%s'\n", output_filename.c_str());
+			return INSTALL_ERROR;
+		}
+
+		ret_val = mzExtractZipEntryToFile(Zip, selinx_contexts, file_contexts_fd);
+		close(file_contexts_fd);
+
+		if (!ret_val) {
+			mzCloseZipArchive(Zip);
+			LOGERR("Could not extract '%s'\n", ASSUMED_UPDATE_BINARY_NAME);
+			return INSTALL_ERROR;
+		}
+	}
+	mzCloseZipArchive(Zip);
+
 	pipe(pipe_fd);
 
 	args[0] = Temp_Binary.c_str();
diff --git a/updater/install.c b/updater/install.c
index 305703cd3..d8c6dc249 100644
--- a/updater/install.c
+++ b/updater/install.c
@@ -37,6 +37,7 @@
 #include "mtdutils/mtdutils.h"
 #include "updater.h"
 #include "applypatch/applypatch.h"
+#include "flashutils/flashutils.h"
 
 #ifdef USE_EXT4
 #include "make_ext4fs.h"
diff --git a/updater/updater.c b/updater/updater.c
index 58ac27f9e..15d3d412f 100644
--- a/updater/updater.c
+++ b/updater/updater.c
@@ -17,6 +17,7 @@
 #include <stdio.h>
 #include <unistd.h>
 #include <stdlib.h>
+#include <fcntl.h>
 
 #include "edify/expr.h"
 #include "updater.h"
@@ -31,6 +32,8 @@
 // Where in the package we expect to find the edify script to execute.
 // (Note it's "updateR-script", not the older "update-script".)
 #define SCRIPT_NAME "META-INF/com/google/android/updater-script"
+#define SELINUX_CONTEXTS_ZIP "file_contexts"
+#define SELINUX_CONTEXTS_TMP "/tmp/file_contexts"
 
 struct selabel_handle *sehandle;
 
@@ -87,6 +90,23 @@ int main(int argc, char** argv) {
     }
     script[script_entry->uncompLen] = '\0';
 
+    const ZipEntry* file_contexts_entry = mzFindZipEntry(&za, SELINUX_CONTEXTS_ZIP);
+    if (script_entry != NULL) {
+        int file_contexts_fd = creat(SELINUX_CONTEXTS_TMP, 0644);
+		if (file_contexts_fd < 0) {
+			fprintf(stderr, "Could not extract %s to '%s'\n", SELINUX_CONTEXTS_ZIP, SELINUX_CONTEXTS_TMP);
+			return 3;
+		}
+
+		int ret_val = mzExtractZipEntryToFile(&za, file_contexts_entry, file_contexts_fd);
+		close(file_contexts_fd);
+
+		if (!ret_val) {
+			fprintf(stderr, "Could not extract '%s'\n", SELINUX_CONTEXTS_ZIP);
+			return 3;
+		}
+    }
+
     // Configure edify's functions.
 
     RegisterBuiltins();
@@ -105,11 +125,19 @@ int main(int argc, char** argv) {
         return 6;
     }
 
-    struct selinux_opt seopts[] = {
-      { SELABEL_OPT_PATH, "/file_contexts" }
-    };
+    if (access(SELINUX_CONTEXTS_TMP, R_OK) == 0) {
+        struct selinux_opt seopts[] = {
+          { SELABEL_OPT_PATH, SELINUX_CONTEXTS_TMP }
+        };
 
-    sehandle = selabel_open(SELABEL_CTX_FILE, seopts, 1);
+        sehandle = selabel_open(SELABEL_CTX_FILE, seopts, 1);
+    } else {
+        struct selinux_opt seopts[] = {
+          { SELABEL_OPT_PATH, "/file_contexts" }
+        };
+
+        sehandle = selabel_open(SELABEL_CTX_FILE, seopts, 1);
+    }
 
     if (!sehandle) {
         fprintf(stderr, "Warning:  No file_contexts\n");