[MTE] add test for mtectrl

Bug: 245624194
Test: atest --no-bazel-mode mtectrl_test
Change-Id: I9f6920a1154d33b58fa3f62f809f54dce1e4448d

3 files changed, 171 insertions, 5 deletions

diff --git a/mtectrl/Android.bp b/mtectrl/Android.bp
index bc84f8cd..688ca7dd 100644
--- a/mtectrl/Android.bp
+++ b/mtectrl/Android.bp
@@ -25,3 +25,18 @@ cc_binary {
   ],
   init_rc: ["mtectrl.rc"],
 }
+
+cc_test {
+  name: "mtectrl_test",
+  srcs: ["mtectrl_test.cc"],
+  test_suites: ["general-tests"],
+  // shell cannot use /system/bin/mtectrl
+  require_root: true,
+  shared_libs: [
+    "libbootloader_message",
+    "libbase",
+  ],
+  static_libs: [
+    "libgmock",
+  ]
+}
diff --git a/mtectrl/mtectrl.cc b/mtectrl/mtectrl.cc
index dca37ca7..8ec518a4 100644
--- a/mtectrl/mtectrl.cc
+++ b/mtectrl/mtectrl.cc
@@ -15,17 +15,19 @@
  */
 
 #include <getopt.h>
+#include <unistd.h>
 
+#include <android-base/file.h>
 #include <android-base/logging.h>
 #include <android-base/properties.h>
 #include <android-base/strings.h>
 #include <bootloader_message/bootloader_message.h>
 
+#include <functional>
 #include <iostream>
 
 void AddItem(std::string* s, const char* item) {
-  if (!s->empty())
-    *s += ",";
+  if (!s->empty()) *s += ",";
   *s += item;
 }
 
@@ -85,11 +87,31 @@ bool HandleOverride(const std::string& override_value, misc_memtag_message* m) {
 int main(int argc, char** argv) {
   const char* set_prop = nullptr;
   int opt;
-  while ((opt = getopt(argc, argv, "s:")) != -1) {
+  std::function<bool(misc_memtag_message*, std::string*)> read_memtag_message =
+      ReadMiscMemtagMessage;
+  std::function<bool(const misc_memtag_message&, std::string*)> write_memtag_message =
+      WriteMiscMemtagMessage;
+  while ((opt = getopt(argc, argv, "s:t:")) != -1) {
     switch (opt) {
       case 's':
         set_prop = optarg;
         break;
+      case 't': {
+        // Use different fake misc partition for testing.
+        const char* filename = optarg;
+        int fd = open(filename, O_RDWR | O_CLOEXEC);
+        CHECK_NE(fd, -1);
+        CHECK_NE(ftruncate(fd, sizeof(misc_memtag_message)), -1);
+        read_memtag_message = [fd](misc_memtag_message* m, std::string*) {
+          CHECK(android::base::ReadFully(fd, m, sizeof(*m)));
+          return true;
+        };
+        write_memtag_message = [fd](const misc_memtag_message& m, std::string*) {
+          CHECK(android::base::WriteFully(fd, &m, sizeof(m)));
+          return true;
+        };
+        break;
+      }
       default:
         PrintUsage(argv[0]);
         return 1;
@@ -107,7 +129,7 @@ int main(int argc, char** argv) {
   if (!value && set_prop) {
     std::string err;
     misc_memtag_message m = {};
-    if (!ReadMiscMemtagMessage(&m, &err)) {
+    if (!read_memtag_message(&m, &err)) {
       LOG(ERROR) << "Failed to read memtag message: " << err;
       return 1;
     }
@@ -137,7 +159,7 @@ int main(int argc, char** argv) {
     return 1;
   }
   std::string err;
-  if (!WriteMiscMemtagMessage(m, &err)) {
+  if (!write_memtag_message(m, &err)) {
     LOG(ERROR) << "Failed to apply mode: " << value << ", override: " << override_value << err;
     return 1;
   } else {
diff --git a/mtectrl/mtectrl_test.cc b/mtectrl/mtectrl_test.cc
new file mode 100644
index 00000000..d99c6b44
--- /dev/null
+++ b/mtectrl/mtectrl_test.cc
@@ -0,0 +1,129 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <stdio.h>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include <android-base/file.h>
+#include <android-base/logging.h>
+#include <android-base/properties.h>
+#include <bootloader_message/bootloader_message.h>
+
+namespace {
+using ::testing::StartsWith;
+
+int mtectrl(const char* arg) {
+  std::string cmd = "mtectrl -t /data/local/tmp/misc_memtag ";
+  cmd += arg;
+  return system(cmd.c_str());
+}
+
+std::string GetMisc() {
+  std::string data;
+  CHECK(android::base::ReadFileToString("/data/local/tmp/misc_memtag", &data, false));
+  return data;
+}
+
+std::string TestProperty() {
+  return android::base::GetProperty("arm64.memtag.test_bootctl", "");
+}
+}  // namespace
+
+class MteCtrlTest : public ::testing::Test {
+  void SetUp() override {
+    // Empty fake misc partition.
+    int fd = creat("/data/local/tmp/misc_memtag", 0600);
+    CHECK(fd != -1);
+    CHECK(ftruncate(fd, sizeof(misc_memtag_message)) != -1);
+    close(fd);
+    android::base::SetProperty("arm64.memtag.test_bootctl", "INVALID");
+  }
+  void TearDown() override {
+    CHECK(unlink("/data/local/tmp/misc_memtag") == 0);
+  }
+};
+
+TEST_F(MteCtrlTest, invalid) {
+  EXPECT_NE(mtectrl("memtag-invalid"), 0);
+  EXPECT_NE(mtectrl("memtag override-invalid"), 0);
+}
+
+TEST_F(MteCtrlTest, set_once) {
+  ASSERT_EQ(mtectrl("memtag-once"), 0);
+  EXPECT_THAT(GetMisc(), StartsWith("\x01\x5a\xfe\xfe\x5a\x02"));
+}
+
+TEST_F(MteCtrlTest, set_once_kernel) {
+  ASSERT_EQ(mtectrl("memtag-once,memtag-kernel"), 0);
+  EXPECT_THAT(GetMisc(), StartsWith("\x01\x5a\xfe\xfe\x5a\x06"));
+}
+
+TEST_F(MteCtrlTest, set_memtag) {
+  ASSERT_EQ(mtectrl("memtag"), 0);
+  EXPECT_THAT(GetMisc(), StartsWith("\x01\x5a\xfe\xfe\x5a\x01"));
+}
+
+TEST_F(MteCtrlTest, set_memtag_force_off) {
+  ASSERT_EQ(mtectrl("memtag force_off"), 0);
+  EXPECT_THAT(GetMisc(), StartsWith("\x01\x5a\xfe\xfe\x5a\x10"));
+}
+
+TEST_F(MteCtrlTest, read_memtag) {
+  ASSERT_EQ(mtectrl("memtag"), 0);
+  ASSERT_EQ(mtectrl("-s arm64.memtag.test_bootctl"), 0);
+  EXPECT_EQ(TestProperty(), "memtag");
+}
+
+TEST_F(MteCtrlTest, set_read_memtag) {
+  ASSERT_EQ(mtectrl("-s arm64.memtag.test_bootctl memtag"), 0);
+  EXPECT_EQ(TestProperty(), "memtag");
+}
+
+TEST_F(MteCtrlTest, set_read_force_off) {
+  ASSERT_EQ(mtectrl("-s arm64.memtag.test_bootctl memtag,memtag-once force_off"), 0);
+  EXPECT_EQ(TestProperty(), "memtag-once,memtag-off");
+}
+
+TEST_F(MteCtrlTest, override) {
+  ASSERT_EQ(mtectrl("memtag"), 0);
+  ASSERT_EQ(mtectrl("memtag-once"), 0);
+  EXPECT_THAT(GetMisc(), StartsWith("\x01\x5a\xfe\xfe\x5a\x02"));
+}
+
+TEST_F(MteCtrlTest, read_empty) {
+  ASSERT_EQ(mtectrl("-s arm64.memtag.test_bootctl"), 0);
+  EXPECT_EQ(TestProperty(), "");
+}
+
+TEST_F(MteCtrlTest, force_off_invalid_mode) {
+  mtectrl("-s arm64.memtag.test_bootctl memtag-invalid force_off");
+  EXPECT_EQ(TestProperty(), "memtag-off");
+  EXPECT_THAT(GetMisc(), StartsWith("\x01\x5a\xfe\xfe\x5a\x10"));
+}
+
+TEST_F(MteCtrlTest, force_on_invalid_mode) {
+  mtectrl("-s arm64.memtag.test_bootctl memtag-invalid force_on");
+  EXPECT_EQ(TestProperty(), "memtag");
+  EXPECT_THAT(GetMisc(), StartsWith("\x01\x5a\xfe\xfe\x5a\x01"));
+}
+
+TEST_F(MteCtrlTest, mode_invalid_override) {
+  mtectrl("-s arm64.memtag.test_bootctl memtag force_invalid");
+  EXPECT_EQ(TestProperty(), "memtag");
+  EXPECT_THAT(GetMisc(), StartsWith("\x01\x5a\xfe\xfe\x5a\x01"));
+}