tools: add support for ubi tarball from host

Add a parameter that specifies the layout type, static or ubitar.

Signed-off-by: Patrick Venture <venture@google.com>
Change-Id: I86c09d8ffeb55ba789279d2568b579a44b3923d7
diff --git a/tools/main.cpp b/tools/main.cpp
index c944899..c9e76c0 100644
--- a/tools/main.cpp
+++ b/tools/main.cpp
@@ -41,24 +41,38 @@
 #define IPMILPC "ipmilpc"
 #define IPMIPCI "ipmipci"
 #define IPMIBT "ipmibt"
+#define STATIC "static"
+#define UBITAR "ubitar"
 
 namespace
 {
 const std::vector<std::string> interfaceList = {IPMIBT, IPMILPC, IPMIPCI};
-}
+const std::vector<std::string> typeList = {STATIC, UBITAR};
+} // namespace
 
 void usage(const char* program)
 {
     std::fprintf(
         stderr,
         "Usage: %s --command <command> --interface <interface> --image "
-        "<image file> --sig <signature file>\n",
+        "<image file> --sig <signature file> --type <layout>\n",
         program);
 
     std::fprintf(stderr, "interfaces: ");
     std::copy(interfaceList.begin(), interfaceList.end(),
               std::ostream_iterator<std::string>(std::cerr, ", "));
     std::fprintf(stderr, "\n");
+
+    std::fprintf(stderr, "layouts: ");
+    std::copy(typeList.begin(), typeList.end(),
+              std::ostream_iterator<std::string>(std::cerr, ", "));
+    std::fprintf(stderr, "\n");
+}
+
+bool checkType(const std::string& type)
+{
+    auto tFound = std::find(typeList.begin(), typeList.end(), type);
+    return (tFound != typeList.end());
 }
 
 bool checkCommand(const std::string& command)
@@ -75,7 +89,7 @@
 
 int main(int argc, char* argv[])
 {
-    std::string command, interface, imagePath, signaturePath;
+    std::string command, interface, imagePath, signaturePath, type;
     char* valueEnd = nullptr;
     long address = 0;
     long length = 0;
@@ -92,13 +106,14 @@
             {"sig", required_argument, 0, 's'},
             {"address", required_argument, 0, 'a'},
             {"length", required_argument, 0, 'l'},
+            {"type", required_argument, 0, 't'},
             {0, 0, 0, 0}
         };
         // clang-format on
 
         int option_index = 0;
-        int c =
-            getopt_long(argc, argv, "c:i:m:s:a:l", long_options, &option_index);
+        int c = getopt_long(argc, argv, "c:i:m:s:a:l:t:", long_options,
+                            &option_index);
         if (c == -1)
         {
             break;
@@ -159,6 +174,14 @@
                 }
                 hostLength = static_cast<std::uint32_t>(length);
                 break;
+            case 't':
+                type = std::string{optarg};
+                if (!checkType(type))
+                {
+                    usage(argv[0]);
+                    exit(EXIT_FAILURE);
+                }
+                break;
             default:
                 usage(argv[0]);
                 exit(EXIT_FAILURE);
@@ -220,7 +243,7 @@
         try
         {
             host_tool::UpdateHandler updater(&blob, handler.get());
-            host_tool::updaterMain(&updater, imagePath, signaturePath);
+            host_tool::updaterMain(&updater, imagePath, signaturePath, type);
         }
         catch (const host_tool::ToolException& e)
         {
diff --git a/tools/test/tools_updater_unittest.cpp b/tools/test/tools_updater_unittest.cpp
index a2ba2df..85f86bc 100644
--- a/tools/test/tools_updater_unittest.cpp
+++ b/tools/test/tools_updater_unittest.cpp
@@ -110,7 +110,7 @@
     EXPECT_CALL(handler, verifyFile(ipmi_flash::updateBlobId))
         .WillOnce(Return(true));
 
-    updaterMain(&handler, image, signature);
+    updaterMain(&handler, image, signature, "static");
 }
 
 TEST_F(UpdaterTest, UpdateMainCleansUpOnFailure)
@@ -126,7 +126,8 @@
         .WillOnce(Return(false));
     EXPECT_CALL(handler, cleanArtifacts()).WillOnce(Return());
 
-    EXPECT_THROW(updaterMain(&handler, image, signature), ToolException);
+    EXPECT_THROW(updaterMain(&handler, image, signature, "static"),
+                 ToolException);
 }
 
 } // namespace host_tool
diff --git a/tools/updater.cpp b/tools/updater.cpp
index a364a4b..438fc4a 100644
--- a/tools/updater.cpp
+++ b/tools/updater.cpp
@@ -35,13 +35,14 @@
 {
 
 void updaterMain(UpdateHandlerInterface* updater, const std::string& imagePath,
-                 const std::string& signaturePath)
+                 const std::string& signaturePath,
+                 const std::string& layoutType)
 {
-    /* TODO(venture): Add optional parameter to specify the flash type, default
-     * to legacy for now.
-     */
-    bool goalSupported =
-        updater->checkAvailable(ipmi_flash::staticLayoutBlobId);
+    const auto& layout = (layoutType == "static")
+                             ? ipmi_flash::staticLayoutBlobId
+                             : ipmi_flash::ubiTarballBlobId;
+
+    bool goalSupported = updater->checkAvailable(layout);
     if (!goalSupported)
     {
         throw ToolException("Goal firmware or interface not supported");
@@ -50,10 +51,9 @@
     /* Yay, our data handler is supported. */
     try
     {
-
         /* Send over the firmware image. */
         std::fprintf(stderr, "Sending over the firmware image.\n");
-        updater->sendFile(ipmi_flash::staticLayoutBlobId, imagePath);
+        updater->sendFile(layout, imagePath);
 
         /* Send over the hash contents. */
         std::fprintf(stderr, "Sending over the hash file.\n");
diff --git a/tools/updater.hpp b/tools/updater.hpp
index b3ab947..331b551 100644
--- a/tools/updater.hpp
+++ b/tools/updater.hpp
@@ -14,9 +14,11 @@
  * @param[in] updater - update handler object.
  * @param[in] imagePath - the path to the image file.
  * @param[in] signaturePath - the path to the signature file.
+ * @param[in] layoutType - the image update layout type (static/ubi/other)
  * @throws ToolException on failures.
  */
 void updaterMain(UpdateHandlerInterface* updater, const std::string& imagePath,
-                 const std::string& signaturePath);
+                 const std::string& signaturePath,
+                 const std::string& layoutType);
 
 } // namespace host_tool