tools: add cleanup blob on failure

If the update process fails, call the cleanup blob and try to cleanup
artifacts automatically.

Signed-off-by: Patrick Venture <venture@google.com>
Change-Id: Ib57e5cc7f13b73e2d7371116058469a4d4bfd0c4
diff --git a/tools/test/tools_updater_unittest.cpp b/tools/test/tools_updater_unittest.cpp
index e54e613..5b3cb5b 100644
--- a/tools/test/tools_updater_unittest.cpp
+++ b/tools/test/tools_updater_unittest.cpp
@@ -1,5 +1,6 @@
 #include "data_interface_mock.hpp"
 #include "status.hpp"
+#include "tool_errors.hpp"
 #include "updater.hpp"
 #include "updater_mock.hpp"
 #include "util.hpp"
@@ -138,4 +139,20 @@
     updaterMain(&handler, image, signature);
 }
 
+TEST_F(UpdaterTest, UpdateMainCleansUpOnFailure)
+{
+    const std::string image = "image.bin";
+    const std::string signature = "signature.bin";
+    UpdateHandlerMock handler;
+
+    EXPECT_CALL(handler, checkAvailable(_)).WillOnce(Return(true));
+    EXPECT_CALL(handler, sendFile(_, image)).WillOnce(Return());
+    EXPECT_CALL(handler, sendFile(_, signature)).WillOnce(Return());
+    EXPECT_CALL(handler, verifyFile(ipmi_flash::verifyBlobId))
+        .WillOnce(Return(false));
+    EXPECT_CALL(handler, cleanArtifacts()).WillOnce(Return());
+
+    EXPECT_THROW(updaterMain(&handler, image, signature), ToolException);
+}
+
 } // namespace host_tool
diff --git a/tools/test/updater_mock.hpp b/tools/test/updater_mock.hpp
index ca4c4b9..d67991b 100644
--- a/tools/test/updater_mock.hpp
+++ b/tools/test/updater_mock.hpp
@@ -15,6 +15,7 @@
     MOCK_METHOD1(checkAvailable, bool(const std::string&));
     MOCK_METHOD2(sendFile, void(const std::string&, const std::string&));
     MOCK_METHOD1(verifyFile, bool(const std::string&));
+    MOCK_METHOD0(cleanArtifacts, void());
 };
 
 } // namespace host_tool
diff --git a/tools/updater.cpp b/tools/updater.cpp
index a5007b2..0c4b740 100644
--- a/tools/updater.cpp
+++ b/tools/updater.cpp
@@ -231,6 +231,28 @@
     return (success == true);
 }
 
+void UpdateHandler::cleanArtifacts()
+{
+    /* open(), commit(), close() */
+    std::uint16_t session;
+
+    /* Errors aren't important for this call. */
+    try
+    {
+        std::fprintf(stderr, "Opening the cleanup blob\n");
+        session =
+            blob->openBlob(ipmi_flash::cleanupBlobId,
+                           static_cast<std::uint16_t>(blobs::OpenFlags::write));
+        std::fprintf(stderr, "Committing to the cleanup blob\n");
+        blob->commit(session, {});
+        std::fprintf(stderr, "Closing cleanup blob\n");
+        blob->closeBlob(session);
+    }
+    catch (...)
+    {
+    }
+}
+
 void updaterMain(UpdateHandlerInterface* updater, const std::string& imagePath,
                  const std::string& signaturePath)
 {
@@ -245,42 +267,51 @@
     }
 
     /* Yay, our data handler is supported. */
-
-    /* Send over the firmware image. */
-    std::fprintf(stderr, "Sending over the firmware image.\n");
-    updater->sendFile(ipmi_flash::staticLayoutBlobId, imagePath);
-
-    /* Send over the hash contents. */
-    std::fprintf(stderr, "Sending over the hash file.\n");
-    updater->sendFile(ipmi_flash::hashBlobId, signaturePath);
-
-    /* Trigger the verification by opening and committing the verify file. */
-    std::fprintf(stderr, "Opening the verification file\n");
-    if (updater->verifyFile(ipmi_flash::verifyBlobId))
+    try
     {
-        std::fprintf(stderr, "succeeded\n");
-    }
-    else
-    {
-        std::fprintf(stderr, "failed\n");
-        throw ToolException("Verification failed");
-    }
 
-    /* Trigger the update by opening and committing the update file. */
-    std::fprintf(stderr, "Opening the update file\n");
-    if (updater->verifyFile(ipmi_flash::updateBlobId))
-    {
-        std::fprintf(stderr, "succeeded\n");
-    }
-    else
-    {
-        /* Depending on the update mechanism used, this may be uninteresting.
-         * For instance, for the static layout, we use the reboot update
-         * mechanism.  Which doesn't always lead to a successful return before
-         * the BMC starts shutting down services.
+        /* Send over the firmware image. */
+        std::fprintf(stderr, "Sending over the firmware image.\n");
+        updater->sendFile(ipmi_flash::staticLayoutBlobId, imagePath);
+
+        /* Send over the hash contents. */
+        std::fprintf(stderr, "Sending over the hash file.\n");
+        updater->sendFile(ipmi_flash::hashBlobId, signaturePath);
+
+        /* Trigger the verification by opening and committing the verify file.
          */
-        std::fprintf(stderr, "failed\n");
-        throw ToolException("Update failed");
+        std::fprintf(stderr, "Opening the verification file\n");
+        if (updater->verifyFile(ipmi_flash::verifyBlobId))
+        {
+            std::fprintf(stderr, "succeeded\n");
+        }
+        else
+        {
+            std::fprintf(stderr, "failed\n");
+            throw ToolException("Verification failed");
+        }
+
+        /* Trigger the update by opening and committing the update file. */
+        std::fprintf(stderr, "Opening the update file\n");
+        if (updater->verifyFile(ipmi_flash::updateBlobId))
+        {
+            std::fprintf(stderr, "succeeded\n");
+        }
+        else
+        {
+            /* Depending on the update mechanism used, this may be
+             * uninteresting. For instance, for the static layout, we use the
+             * reboot update mechanism.  Which doesn't always lead to a
+             * successful return before the BMC starts shutting down services.
+             */
+            std::fprintf(stderr, "failed\n");
+            throw ToolException("Update failed");
+        }
+    }
+    catch (...)
+    {
+        updater->cleanArtifacts();
+        throw;
     }
 }
 
diff --git a/tools/updater.hpp b/tools/updater.hpp
index 670d981..4352c82 100644
--- a/tools/updater.hpp
+++ b/tools/updater.hpp
@@ -39,6 +39,11 @@
      * @return true if verified, false if verification errors.
      */
     virtual bool verifyFile(const std::string& target) = 0;
+
+    /**
+     * Cleanup the artifacts by triggering this action.
+     */
+    virtual void cleanArtifacts() = 0;
 };
 
 /** Object that actually handles the update itself. */
@@ -64,6 +69,8 @@
      */
     bool verifyFile(const std::string& target) override;
 
+    void cleanArtifacts() override;
+
   private:
     ipmiblob::BlobInterface* blob;
     DataInterface* handler;