Merge "libfiemap: Handle EAGAIN in fallocate()."
diff --git a/bootstat/bootstat.cpp b/bootstat/bootstat.cpp
index 844357c..3b8866e 100644
--- a/bootstat/bootstat.cpp
+++ b/bootstat/bootstat.cpp
@@ -459,6 +459,8 @@
     {"reboot,sys_ldo_ok,pmic,main", 227},
     {"reboot,sys_ldo_ok,pmic,sub", 228},
     {"reboot,smpl_timeout,pmic,main", 229},
+    {"reboot,ota,.*", 230},
+    {"reboot,periodic,.*", 231},
 };
 
 // Converts a string value representing the reason the system booted to an
diff --git a/debuggerd/client/debuggerd_client_test.cpp b/debuggerd/client/debuggerd_client_test.cpp
index ebb8d86..33ff05f 100644
--- a/debuggerd/client/debuggerd_client_test.cpp
+++ b/debuggerd/client/debuggerd_client_test.cpp
@@ -18,6 +18,7 @@
 
 #include <fcntl.h>
 #include <stdio.h>
+#include <sys/eventfd.h>
 #include <unistd.h>
 
 #include <chrono>
@@ -51,23 +52,35 @@
 
 TEST(debuggerd_client, race) {
   static int THREAD_COUNT = getThreadCount();
+
+  // Semaphore incremented once per thread started.
+  unique_fd barrier(eventfd(0, EFD_SEMAPHORE));
+  ASSERT_NE(-1, barrier.get());
+
   pid_t forkpid = fork();
-
   ASSERT_NE(-1, forkpid);
-
   if (forkpid == 0) {
     // Spawn a bunch of threads, to make crash_dump take longer.
     std::vector<std::thread> threads;
+    threads.reserve(THREAD_COUNT);
     for (int i = 0; i < THREAD_COUNT; ++i) {
-      threads.emplace_back([]() {
-        while (true) {
-          std::this_thread::sleep_for(60s);
+      threads.emplace_back([&barrier]() {
+        uint64_t count = 1;
+        ASSERT_NE(-1, write(barrier.get(), &count, sizeof(count)));
+        for (;;) {
+          pause();
         }
       });
     }
+    for (;;) {
+      pause();
+    }
+  }
 
-    std::this_thread::sleep_for(60s);
-    exit(0);
+  // Wait for the child to spawn all of its threads.
+  for (int i = 0; i < THREAD_COUNT; ++i) {
+    uint64_t count;
+    ASSERT_NE(-1, read(barrier.get(), &count, sizeof(count)));
   }
 
   unique_fd pipe_read, pipe_write;
@@ -77,9 +90,6 @@
   constexpr int PIPE_SIZE = 16 * 1024 * 1024;
   ASSERT_EQ(PIPE_SIZE, fcntl(pipe_read.get(), F_SETPIPE_SZ, PIPE_SIZE));
 
-  // Wait for a bit to let the child spawn all of its threads.
-  std::this_thread::sleep_for(1s);
-
   ASSERT_TRUE(
       debuggerd_trigger_dump(forkpid, kDebuggerdNativeBacktrace, 60000, std::move(pipe_write)));
   // Immediately kill the forked child, to make sure that the dump didn't return early.
diff --git a/fastboot/Android.bp b/fastboot/Android.bp
index 7794c4b..56cac88 100644
--- a/fastboot/Android.bp
+++ b/fastboot/Android.bp
@@ -196,6 +196,7 @@
         "libfastbootshim",
         "libsnapshot_cow",
         "liblz4",
+        "libzstd",
         "libsnapshot_nobinder",
         "update_metadata-protos",
         "liburing",
diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index 3dd1f1a..d3bd904 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -163,6 +163,7 @@
         "libbrotli",
         "libz",
         "liblz4",
+        "libzstd",
     ],
     export_include_dirs: ["include"],
 }
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h b/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h
index 6baf4be..c3ca00a 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/cow_format.h
@@ -157,7 +157,8 @@
     kCowCompressNone = 0,
     kCowCompressGz = 1,
     kCowCompressBrotli = 2,
-    kCowCompressLz4 = 3
+    kCowCompressLz4 = 3,
+    kCowCompressZstd = 4,
 };
 
 static constexpr uint8_t kCowReadAheadNotStarted = 0;
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/cow_compress.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/cow_compress.cpp
index d06c904..a4a0ad6 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/cow_compress.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/cow_compress.cpp
@@ -29,6 +29,7 @@
 #include <libsnapshot/cow_writer.h>
 #include <lz4.h>
 #include <zlib.h>
+#include <zstd.h>
 
 namespace android {
 namespace snapshot {
@@ -40,6 +41,8 @@
         return {kCowCompressBrotli};
     } else if (name == "lz4") {
         return {kCowCompressLz4};
+    } else if (name == "zstd") {
+        return {kCowCompressZstd};
     } else if (name == "none" || name.empty()) {
         return {kCowCompressNone};
     } else {
@@ -112,6 +115,23 @@
             }
             return buffer;
         }
+        case kCowCompressZstd: {
+            std::basic_string<uint8_t> buffer(ZSTD_compressBound(length), '\0');
+            const auto compressed_size =
+                    ZSTD_compress(buffer.data(), buffer.size(), data, length, 0);
+            if (compressed_size <= 0) {
+                LOG(ERROR) << "ZSTD compression failed " << compressed_size;
+                return {};
+            }
+            // Don't run compression if the compressed output is larger
+            if (compressed_size >= length) {
+                buffer.resize(length);
+                memcpy(buffer.data(), data, length);
+            } else {
+                buffer.resize(compressed_size);
+            }
+            return buffer;
+        }
         default:
             LOG(ERROR) << "unhandled compression type: " << compression;
             break;
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.cpp
index 3d34413..da90cc0 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.cpp
@@ -17,12 +17,15 @@
 #include "cow_decompress.h"
 
 #include <array>
+#include <cstring>
 #include <utility>
+#include <vector>
 
 #include <android-base/logging.h>
 #include <brotli/decode.h>
 #include <lz4.h>
 #include <zlib.h>
+#include <zstd.h>
 
 namespace android {
 namespace snapshot {
@@ -336,9 +339,56 @@
     }
 };
 
+class ZstdDecompressor final : public IDecompressor {
+  public:
+    ssize_t Decompress(void* buffer, size_t buffer_size, size_t decompressed_size,
+                       size_t ignore_bytes = 0) override {
+        if (buffer_size < decompressed_size - ignore_bytes) {
+            LOG(INFO) << "buffer size " << buffer_size
+                      << " is not large enough to hold decompressed data. Decompressed size "
+                      << decompressed_size << ", ignore_bytes " << ignore_bytes;
+            return -1;
+        }
+        if (ignore_bytes == 0) {
+            if (!Decompress(buffer, decompressed_size)) {
+                return -1;
+            }
+            return decompressed_size;
+        }
+        std::vector<unsigned char> ignore_buf(decompressed_size);
+        if (!Decompress(buffer, decompressed_size)) {
+            return -1;
+        }
+        memcpy(buffer, ignore_buf.data() + ignore_bytes, buffer_size);
+        return decompressed_size;
+    }
+    bool Decompress(void* output_buffer, const size_t output_size) {
+        std::string input_buffer;
+        input_buffer.resize(stream_->Size());
+        size_t bytes_read = stream_->Read(input_buffer.data(), input_buffer.size());
+        if (bytes_read != input_buffer.size()) {
+            LOG(ERROR) << "Failed to read all input at once. Expected: " << input_buffer.size()
+                       << " actual: " << bytes_read;
+            return false;
+        }
+        const auto bytes_decompressed = ZSTD_decompress(output_buffer, output_size,
+                                                        input_buffer.data(), input_buffer.size());
+        if (bytes_decompressed != output_size) {
+            LOG(ERROR) << "Failed to decompress ZSTD block, expected output size: " << output_size
+                       << ", actual: " << bytes_decompressed;
+            return false;
+        }
+        return true;
+    }
+};
+
 std::unique_ptr<IDecompressor> IDecompressor::Lz4() {
     return std::make_unique<Lz4Decompressor>();
 }
 
+std::unique_ptr<IDecompressor> IDecompressor::Zstd() {
+    return std::make_unique<ZstdDecompressor>();
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.h b/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.h
index 9e83f3c..52b3d76 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.h
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/cow_decompress.h
@@ -47,6 +47,7 @@
     static std::unique_ptr<IDecompressor> Gz();
     static std::unique_ptr<IDecompressor> Brotli();
     static std::unique_ptr<IDecompressor> Lz4();
+    static std::unique_ptr<IDecompressor> Zstd();
 
     static std::unique_ptr<IDecompressor> FromString(std::string_view compressor);
 
diff --git a/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp b/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp
index 4c5a8bf..6749cf1 100644
--- a/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp
+++ b/fs_mgr/libsnapshot/libsnapshot_cow/cow_reader.cpp
@@ -30,6 +30,7 @@
 #include <zlib.h>
 
 #include "cow_decompress.h"
+#include "libsnapshot/cow_format.h"
 
 namespace android {
 namespace snapshot {
@@ -777,6 +778,11 @@
         case kCowCompressBrotli:
             decompressor = IDecompressor::Brotli();
             break;
+        case kCowCompressZstd:
+            if (header_.block_size != op->data_length) {
+                decompressor = IDecompressor::Zstd();
+            }
+            break;
         case kCowCompressLz4:
             if (header_.block_size != op->data_length) {
                 decompressor = IDecompressor::Lz4();
diff --git a/fs_mgr/libsnapshot/snapuserd/Android.bp b/fs_mgr/libsnapshot/snapuserd/Android.bp
index 1e03683..9fe567a 100644
--- a/fs_mgr/libsnapshot/snapuserd/Android.bp
+++ b/fs_mgr/libsnapshot/snapuserd/Android.bp
@@ -112,6 +112,7 @@
         "liblz4",
         "libext4_utils",
         "liburing",
+        "libzstd",
     ],
 
     header_libs: [
diff --git a/gatekeeperd/gatekeeperd.cpp b/gatekeeperd/gatekeeperd.cpp
index e5241b5..7987167 100644
--- a/gatekeeperd/gatekeeperd.cpp
+++ b/gatekeeperd/gatekeeperd.cpp
@@ -144,14 +144,22 @@
     }
 }
 
-uint32_t GateKeeperProxy::adjust_userId(uint32_t userId) {
+Status GateKeeperProxy::adjust_userId(uint32_t userId, uint32_t* hw_userId) {
     static constexpr uint32_t kGsiOffset = 1000000;
-    CHECK(userId < kGsiOffset);
-    CHECK((aidl_hw_device != nullptr) || (hw_device != nullptr));
-    if (is_running_gsi) {
-        return userId + kGsiOffset;
+    if (userId >= kGsiOffset) {
+        return Status::fromExceptionCode(Status::EX_ILLEGAL_ARGUMENT);
     }
-    return userId;
+
+    if ((aidl_hw_device == nullptr) && (hw_device == nullptr)) {
+        return Status::fromExceptionCode(Status::EX_ILLEGAL_STATE);
+    }
+
+    if (is_running_gsi) {
+        *hw_userId = userId + kGsiOffset;
+        return Status::ok();
+    }
+    *hw_userId = userId;
+    return Status::ok();
 }
 
 #define GK_ERROR *gkResponse = GKResponse::error(), Status::ok()
@@ -201,7 +209,12 @@
     android::hardware::hidl_vec<uint8_t> newPwd;
     newPwd.setToExternal(const_cast<uint8_t*>(desiredPassword.data()), desiredPassword.size());
 
-    uint32_t hw_userId = adjust_userId(userId);
+    uint32_t hw_userId = 0;
+    Status result = adjust_userId(userId, &hw_userId);
+    if (!result.isOk()) {
+        return result;
+    }
+
     uint64_t secureUserId = 0;
     if (aidl_hw_device) {
         // AIDL gatekeeper service
@@ -300,7 +313,12 @@
         }
     }
 
-    uint32_t hw_userId = adjust_userId(userId);
+    uint32_t hw_userId = 0;
+    Status result = adjust_userId(userId, &hw_userId);
+    if (!result.isOk()) {
+        return result;
+    }
+
     android::hardware::hidl_vec<uint8_t> curPwdHandle;
     curPwdHandle.setToExternal(const_cast<uint8_t*>(enrolledPasswordHandle.data()),
                                enrolledPasswordHandle.size());
@@ -410,7 +428,12 @@
     }
     clear_sid(userId);
 
-    uint32_t hw_userId = adjust_userId(userId);
+    uint32_t hw_userId = 0;
+    Status result = adjust_userId(userId, &hw_userId);
+    if (!result.isOk()) {
+        return result;
+    }
+
     if (aidl_hw_device) {
         aidl_hw_device->deleteUser(hw_userId);
     } else if (hw_device) {
diff --git a/gatekeeperd/gatekeeperd.h b/gatekeeperd/gatekeeperd.h
index 29873da..b1f08c6 100644
--- a/gatekeeperd/gatekeeperd.h
+++ b/gatekeeperd/gatekeeperd.h
@@ -47,7 +47,7 @@
 
     // This should only be called on userIds being passed to the GateKeeper HAL. It ensures that
     // secure storage shared across a GSI image and a host image will not overlap.
-    uint32_t adjust_userId(uint32_t userId);
+    Status adjust_userId(uint32_t userId, uint32_t* hw_userId);
 
 #define GK_ERROR *gkResponse = GKResponse::error(), Status::ok()
 
diff --git a/init/Android.bp b/init/Android.bp
index 7b52903..41c7a95 100644
--- a/init/Android.bp
+++ b/init/Android.bp
@@ -169,6 +169,7 @@
         "libfsverity_init",
         "liblmkd_utils",
         "liblz4",
+        "libzstd",
         "libmini_keyctl_static",
         "libmodprobe",
         "libprocinfo",
@@ -370,6 +371,7 @@
         "libprotobuf-cpp-lite",
         "libsnapshot_cow",
         "liblz4",
+        "libzstd",
         "libsnapshot_init",
         "update_metadata-protos",
         "libprocinfo",
diff --git a/init/init.cpp b/init/init.cpp
index be1ebee..da63fdc 100644
--- a/init/init.cpp
+++ b/init/init.cpp
@@ -1043,6 +1043,12 @@
     SetProperty(gsi::kGsiBootedProp, is_running);
     auto is_installed = android::gsi::IsGsiInstalled() ? "1" : "0";
     SetProperty(gsi::kGsiInstalledProp, is_installed);
+    if (android::gsi::IsGsiRunning()) {
+        std::string dsu_slot;
+        if (android::gsi::GetActiveDsu(&dsu_slot)) {
+            SetProperty(gsi::kDsuSlotProp, dsu_slot);
+        }
+    }
 
     am.QueueBuiltinAction(SetupCgroupsAction, "SetupCgroups");
     am.QueueBuiltinAction(SetKptrRestrictAction, "SetKptrRestrict");
diff --git a/init/test_kill_services/init_kill_services_test.cpp b/init/test_kill_services/init_kill_services_test.cpp
index 5355703..227380d 100644
--- a/init/test_kill_services/init_kill_services_test.cpp
+++ b/init/test_kill_services/init_kill_services_test.cpp
@@ -16,6 +16,7 @@
 
 #include <gtest/gtest.h>
 
+#include <android-base/logging.h>
 #include <android-base/properties.h>
 
 #include <iostream>
@@ -24,6 +25,7 @@
 using ::android::base::SetProperty;
 
 void ExpectKillingServiceRecovers(const std::string& service_name) {
+    LOG(INFO) << "hello " << service_name << "!";
     const std::string status_prop = "init.svc." + service_name;
     const std::string pid_prop = "init.svc_debug_pid." + service_name;
 
@@ -32,6 +34,7 @@
     ASSERT_EQ("running", GetProperty(status_prop, "")) << status_prop;
     ASSERT_NE("", initial_pid) << pid_prop;
 
+    LOG(INFO) << "okay, now goodbye " << service_name;
     EXPECT_EQ(0, system(("kill -9 " + initial_pid).c_str()));
 
     constexpr size_t kMaxWaitMilliseconds = 10000;
@@ -42,11 +45,16 @@
     for (size_t retry = 0; retry < kRetryTimes; retry++) {
         const std::string& pid = GetProperty(pid_prop, "");
         if (pid != initial_pid && pid != "") break;
+        LOG(INFO) << "I said goodbye " << service_name << "!";
         usleep(kRetryWaitMilliseconds * 1000);
     }
 
+    LOG(INFO) << "are you still there " << service_name << "?";
+
     // svc_debug_pid is set after svc property
     EXPECT_EQ("running", GetProperty(status_prop, ""));
+
+    LOG(INFO) << "I'm done with " << service_name;
 }
 
 class InitKillServicesTest : public ::testing::TestWithParam<std::string> {};
diff --git a/libsysutils/Android.bp b/libsysutils/Android.bp
index 5f472b2..1b41a6b 100644
--- a/libsysutils/Android.bp
+++ b/libsysutils/Android.bp
@@ -29,6 +29,10 @@
         "liblog",
     ],
 
+    header_libs: [
+        "bpf_headers",
+    ],
+
     export_include_dirs: ["include"],
 
     tidy: true,
diff --git a/libsysutils/src/NetlinkEvent.cpp b/libsysutils/src/NetlinkEvent.cpp
index 515cc10..cd9db54 100644
--- a/libsysutils/src/NetlinkEvent.cpp
+++ b/libsysutils/src/NetlinkEvent.cpp
@@ -37,10 +37,12 @@
 #include <sys/utsname.h>
 
 #include <android-base/parseint.h>
+#include <bpf/KernelUtils.h>
 #include <log/log.h>
 #include <sysutils/NetlinkEvent.h>
 
 using android::base::ParseInt;
+using android::bpf::isKernel64Bit;
 
 /* From kernel's net/netfilter/xt_quota2.c */
 const int LOCAL_QLOG_NL_EVENT = 112;
@@ -138,60 +140,6 @@
 static_assert(sizeof(ulog_packet_msg32_t) == 168);
 static_assert(sizeof(ulog_packet_msg64_t) == 192);
 
-// Figure out the bitness of userspace.
-// Trivial and known at compile time.
-static bool isUserspace64bit(void) {
-    return sizeof(long) == 8;
-}
-
-// Figure out the bitness of the kernel.
-static bool isKernel64Bit(void) {
-    // a 64-bit userspace requires a 64-bit kernel
-    if (isUserspace64bit()) return true;
-
-    static bool init = false;
-    static bool cache = false;
-    if (init) return cache;
-
-    // Retrieve current personality - on Linux this system call *cannot* fail.
-    int p = personality(0xffffffff);
-    // But if it does just assume kernel and userspace (which is 32-bit) match...
-    if (p == -1) return false;
-
-    // This will effectively mask out the bottom 8 bits, and switch to 'native'
-    // personality, and then return the previous personality of this thread
-    // (likely PER_LINUX or PER_LINUX32) with any extra options unmodified.
-    int q = personality((p & ~PER_MASK) | PER_LINUX);
-    // Per man page this theoretically could error out with EINVAL,
-    // but kernel code analysis suggests setting PER_LINUX cannot fail.
-    // Either way, assume kernel and userspace (which is 32-bit) match...
-    if (q != p) return false;
-
-    struct utsname u;
-    (void)uname(&u);  // only possible failure is EFAULT, but u is on stack.
-
-    // Switch back to previous personality.
-    // Theoretically could fail with EINVAL on arm64 with no 32-bit support,
-    // but then we wouldn't have fetched 'p' from the kernel in the first place.
-    // Either way there's nothing meaningul we can do in case of error.
-    // Since PER_LINUX32 vs PER_LINUX only affects uname.machine it doesn't
-    // really hurt us either.  We're really just switching back to be 'clean'.
-    (void)personality(p);
-
-    // Possible values of utsname.machine observed on x86_64 desktop (arm via qemu):
-    //   x86_64 i686 aarch64 armv7l
-    // additionally observed on arm device:
-    //   armv8l
-    // presumably also might just be possible:
-    //   i386 i486 i586
-    // and there might be other weird arm32 cases.
-    // We note that the 64 is present in both 64-bit archs,
-    // and in general is likely to be present in only 64-bit archs.
-    cache = !!strstr(u.machine, "64");
-    init = true;
-    return cache;
-}
-
 /******************************************************************************/
 
 NetlinkEvent::NetlinkEvent() {
diff --git a/trusty/confirmationui/fuzz/Android.bp b/trusty/confirmationui/fuzz/Android.bp
index 4780943..96804bd 100644
--- a/trusty/confirmationui/fuzz/Android.bp
+++ b/trusty/confirmationui/fuzz/Android.bp
@@ -26,7 +26,8 @@
         "-DTRUSTY_APP_FILENAME=\"confirmationui.syms.elf\"",
     ],
     fuzz_config: {
-       cc: ["trong@google.com"],
+       cc: ["mikemcternan@google.com"],
+       componentid: 1290237,
     },
 
 }
@@ -40,7 +41,8 @@
         "libdmabufheap",
     ],
     fuzz_config: {
-       cc: ["trong@google.com"],
+       cc: ["mikemcternan@google.com"],
+       componentid: 1290237,
     },
 
     // The initial corpus for this fuzzer was derived by dumping messages from/to
diff --git a/trusty/keymaster/set_attestation_ids/set_attestation_ids.cpp b/trusty/keymaster/set_attestation_ids/set_attestation_ids.cpp
index e944167..6b8f90f 100644
--- a/trusty/keymaster/set_attestation_ids/set_attestation_ids.cpp
+++ b/trusty/keymaster/set_attestation_ids/set_attestation_ids.cpp
@@ -17,8 +17,10 @@
 #include <getopt.h>
 
 #include <string>
+#include <vector>
 
 #include <android-base/properties.h>
+#include <android-base/strings.h>
 #include <trusty_keymaster/ipc/trusty_keymaster_ipc.h>
 
 namespace {
@@ -34,14 +36,66 @@
         {"model", required_argument, nullptr, 'm'},
         {"imei", required_argument, nullptr, 'i'},
         {"meid", required_argument, nullptr, 'c'},
+        {"imei2", required_argument, nullptr, '2'},
         {0, 0, 0, 0},
 };
 
+std::string TELEPHONY_CMD_GET_IMEI = "cmd phone get-imei ";
+
+// Run a shell command and collect the output of it. If any error, set an empty string as the
+// output.
+std::string exec_command(const std::string& command) {
+    char buffer[128];
+    std::string result = "";
+
+    FILE* pipe = popen(command.c_str(), "r");
+    if (!pipe) {
+        fprintf(stderr, "popen('%s') failed\n", command.c_str());
+        return result;
+    }
+
+    while (!feof(pipe)) {
+        if (fgets(buffer, 128, pipe) != NULL) {
+            result += buffer;
+        }
+    }
+
+    pclose(pipe);
+    return result;
+}
+
+// Get IMEI using Telephony service shell command. If any error while executing the command
+// then empty string will be returned as output.
+std::string get_imei(int slot) {
+    std::string cmd = TELEPHONY_CMD_GET_IMEI + std::to_string(slot);
+    std::string output = exec_command(cmd);
+
+    if (output.empty()) {
+        fprintf(stderr, "Retrieve IMEI command ('%s') failed\n", cmd.c_str());
+        return "";
+    }
+
+    std::vector<std::string> out =
+            ::android::base::Tokenize(::android::base::Trim(output), "Device IMEI:");
+
+    if (out.size() != 1) {
+        fprintf(stderr, "Error parsing command ('%s') output '%s'\n", cmd.c_str(), output.c_str());
+        return "";
+    }
+
+    std::string imei = ::android::base::Trim(out[0]);
+    if (imei.compare("null") == 0) {
+        fprintf(stderr, "IMEI value from command ('%s') is null, skipping", cmd.c_str());
+        return "";
+    }
+    return imei;
+}
+
 std::string buf2string(const keymaster::Buffer& buf) {
     return std::string(reinterpret_cast<const char*>(buf.peek_read()), buf.available_read());
 }
 
-void print_usage(const char* prog, const keymaster::SetAttestationIdsRequest& req) {
+void print_usage(const char* prog, const keymaster::SetAttestationIdsKM3Request& req) {
     fprintf(stderr,
             "Usage: %s [options]\n"
             "\n"
@@ -55,34 +109,53 @@
             "  -m, --model <val>          set model (default '%s')\n"
             "  -i, --imei <val>           set IMEI (default '%s')\n"
             "  -c, --meid <val>           set MEID (default '%s')\n"
+            "  -2, --imei2 <val>          set second IMEI (default '%s')\n"
             "\n",
-            prog, buf2string(req.brand).c_str(), buf2string(req.device).c_str(),
-            buf2string(req.product).c_str(), buf2string(req.serial).c_str(),
-            buf2string(req.manufacturer).c_str(), buf2string(req.model).c_str(),
-            buf2string(req.imei).c_str(), buf2string(req.meid).c_str());
+            prog, buf2string(req.base.brand).c_str(), buf2string(req.base.device).c_str(),
+            buf2string(req.base.product).c_str(), buf2string(req.base.serial).c_str(),
+            buf2string(req.base.manufacturer).c_str(), buf2string(req.base.model).c_str(),
+            buf2string(req.base.imei).c_str(), buf2string(req.base.meid).c_str(),
+            buf2string(req.second_imei).c_str());
+}
+
+void set_to(keymaster::Buffer* buf, const std::string& value) {
+    if (!value.empty()) {
+        buf->Reinitialize(value.data(), value.size());
+    }
 }
 
 void set_from_prop(keymaster::Buffer* buf, const std::string& prop) {
     std::string prop_value = ::android::base::GetProperty(prop, /* default_value = */ "");
-    if (!prop_value.empty()) {
-        buf->Reinitialize(prop_value.data(), prop_value.size());
-    }
+    set_to(buf, prop_value);
 }
 
-void populate_ids(keymaster::SetAttestationIdsRequest* req) {
+void populate_base_ids(keymaster::SetAttestationIdsRequest* req) {
     set_from_prop(&req->brand, "ro.product.brand");
     set_from_prop(&req->device, "ro.product.device");
     set_from_prop(&req->product, "ro.product.name");
     set_from_prop(&req->serial, "ro.serialno");
     set_from_prop(&req->manufacturer, "ro.product.manufacturer");
     set_from_prop(&req->model, "ro.product.model");
+    std::string imei = get_imei(0);
+    set_to(&req->imei, imei);
+}
+
+void populate_ids(keymaster::SetAttestationIdsKM3Request* req) {
+    populate_base_ids(&req->base);
+
+    // - "What about IMEI?"
+    // - "You've already had it."
+    // - "We've had one, yes. What about second IMEI?"
+    // - "I don't think he knows about second IMEI, Pip."
+    std::string imei2 = get_imei(1);
+    set_to(&req->second_imei, imei2);
 }
 
 }  // namespace
 
 int main(int argc, char** argv) {
     // By default, set attestation IDs to the values in userspace properties.
-    keymaster::SetAttestationIdsRequest req(/* ver = */ 4);
+    keymaster::SetAttestationIdsKM3Request req(/* ver = */ 4);
     populate_ids(&req);
 
     while (true) {
@@ -94,28 +167,31 @@
 
         switch (c) {
             case 'b':
-                req.brand.Reinitialize(optarg, strlen(optarg));
+                req.base.brand.Reinitialize(optarg, strlen(optarg));
                 break;
             case 'd':
-                req.device.Reinitialize(optarg, strlen(optarg));
+                req.base.device.Reinitialize(optarg, strlen(optarg));
                 break;
             case 'p':
-                req.product.Reinitialize(optarg, strlen(optarg));
+                req.base.product.Reinitialize(optarg, strlen(optarg));
                 break;
             case 's':
-                req.serial.Reinitialize(optarg, strlen(optarg));
+                req.base.serial.Reinitialize(optarg, strlen(optarg));
                 break;
             case 'M':
-                req.manufacturer.Reinitialize(optarg, strlen(optarg));
+                req.base.manufacturer.Reinitialize(optarg, strlen(optarg));
                 break;
             case 'm':
-                req.model.Reinitialize(optarg, strlen(optarg));
+                req.base.model.Reinitialize(optarg, strlen(optarg));
                 break;
             case 'i':
-                req.imei.Reinitialize(optarg, strlen(optarg));
+                req.base.imei.Reinitialize(optarg, strlen(optarg));
                 break;
             case 'c':
-                req.meid.Reinitialize(optarg, strlen(optarg));
+                req.base.meid.Reinitialize(optarg, strlen(optarg));
+                break;
+            case '2':
+                req.second_imei.Reinitialize(optarg, strlen(optarg));
                 break;
             case 'h':
                 print_usage(argv[0], req);
@@ -144,19 +220,33 @@
            "  manufacturer: %s\n"
            "  model:        %s\n"
            "  IMEI:         %s\n"
-           "  MEID:         %s\n",
-           buf2string(req.brand).c_str(), buf2string(req.device).c_str(),
-           buf2string(req.product).c_str(), buf2string(req.serial).c_str(),
-           buf2string(req.manufacturer).c_str(), buf2string(req.model).c_str(),
-           buf2string(req.imei).c_str(), buf2string(req.meid).c_str());
+           "  MEID:         %s\n"
+           "  SECOND_IMEI:  %s\n\n",
+           buf2string(req.base.brand).c_str(), buf2string(req.base.device).c_str(),
+           buf2string(req.base.product).c_str(), buf2string(req.base.serial).c_str(),
+           buf2string(req.base.manufacturer).c_str(), buf2string(req.base.model).c_str(),
+           buf2string(req.base.imei).c_str(), buf2string(req.base.meid).c_str(),
+           buf2string(req.second_imei).c_str());
+    fflush(stdout);
 
     keymaster::EmptyKeymasterResponse rsp(/* ver = */ 4);
-    ret = trusty_keymaster_send(KM_SET_ATTESTATION_IDS, req, &rsp);
-    if (ret) {
-        fprintf(stderr, "SET_ATTESTATION_IDS failed: %d\n", ret);
-        trusty_keymaster_disconnect();
-        return EXIT_FAILURE;
+    const char* msg;
+    if (req.second_imei.available_read() == 0) {
+        // No SECOND_IMEI set, use base command.
+        ret = trusty_keymaster_send(KM_SET_ATTESTATION_IDS, req.base, &rsp);
+        msg = "SET_ATTESTATION_IDS";
+    } else {
+        // SECOND_IMEI is set, use updated command.
+        ret = trusty_keymaster_send(KM_SET_ATTESTATION_IDS_KM3, req, &rsp);
+        msg = "SET_ATTESTATION_IDS_KM3";
     }
+    trusty_keymaster_disconnect();
 
-    return EXIT_SUCCESS;
+    if (ret) {
+        fprintf(stderr, "%s failed: %d\n", msg, ret);
+        return EXIT_FAILURE;
+    } else {
+        printf("done\n");
+        return EXIT_SUCCESS;
+    }
 }