LiveUpdate now needs registered callbacks for serialization, with keys to uniquely identify future partitions

Author: fwsGonzo

lib/LiveUpdate/liveupdate.hpp

@@ -50,17 +50,20 @@ struct LiveUpdate
   typedef delegate<void(Storage&, const buffer_t*)> storage_func;
   typedef delegate<void(Restore&)> resume_func;
 
+  // Register a function to be called when serialization phase begins
+  static void register_serialization_callback(std::string key, storage_func);
+
   // Start a live update process, storing all user-defined data
   // If no storage function is provided no state will be saved
-  static void begin(buffer_t blob, storage_func = nullptr);
+  static void begin(buffer_t blob);
 
   // In the event that LiveUpdate::begin() fails,
   // call this function in the C++ exception handler:
   static void restore_environment();
 
   // Only store user data, as if there was a live update process
   // Throws exception if process or sanity checks fail
-  static buffer_t store(storage_func);
+  static buffer_t store();
 
   // Returns true if there is stored data from before.
   // It performs an extensive validation process to make sure the data is
@@ -74,7 +77,7 @@ struct LiveUpdate
   // Attempt to restore existing stored entries.
   // Returns false if there was nothing there. or if the process failed
   // to be sure that only failure can return false, use is_resumable first
-  static bool resume(resume_func default_handler);
+  static bool resume(std::string key, resume_func default_handler);
 
   // When explicitly resuming from heap, heap overrun checks are disabled
   static bool resume_from_heap(void* location, resume_func default_handler);

lib/LiveUpdate/resume.cpp

@@ -55,7 +55,7 @@ static bool resume_helper(void* location, LiveUpdate::resume_func func)
   extern bool resume_begin(storage_header&, LiveUpdate::resume_func);
   return resume_begin(*(storage_header*) location, func);
 }
-bool LiveUpdate::resume(resume_func func)
+bool LiveUpdate::resume(std::string key, resume_func func)
 {
   void* location = OS::liveupdate_storage_area();
   /// memory sanity check

lib/LiveUpdate/update.cpp

@@ -55,7 +55,24 @@ bool LIVEUPDATE_PERFORM_SANITY_CHECKS = true;
 
 using namespace liu;
 
-static size_t update_store_data(void* location, LiveUpdate::storage_func, const buffer_t*);
+static size_t update_store_data(void* location, const buffer_t*);
+
+// serialization callbacks
+static std::unordered_map<std::string, LiveUpdate::storage_func> storage_callbacks;
+
+void LiveUpdate::register_serialization_callback(std::string key, storage_func callback)
+{
+  auto it = storage_callbacks.find(key);
+  if (it == storage_callbacks.end())
+  {
+    storage_callbacks.emplace(std::piecewise_construct,
+              std::forward_as_tuple(std::move(key)),
+              std::forward_as_tuple(std::move(callback)));
+  }
+  else {
+    throw std::runtime_error("Storage key '" + key + "' already used");
+  }
+}
 
 template <typename Class>
 inline bool validate_header(const Class* hdr)
@@ -66,8 +83,7 @@ inline bool validate_header(const Class* hdr)
            hdr->e_ident[3] == 'F';
 }
 
-void LiveUpdate::begin(buffer_t     blob,
-                       storage_func storage_callback)
+void LiveUpdate::begin(buffer_t  blob)
 {
   void* location = OS::liveupdate_storage_area();
   LPRINT("LiveUpdate::begin(%p, %p:%d, ...)\n", location, blob.data(), (int) blob.size());
@@ -172,7 +188,7 @@ void LiveUpdate::begin(buffer_t     blob,
   LPRINT("* _start is located at %#x\n", start_offset);
 
   // save ourselves if function passed
-  update_store_data(storage_area, storage_callback, &blob);
+  update_store_data(storage_area, &blob);
 
   // 2. flush all devices with flush() interface
   hw::Devices::flush_all();
@@ -225,10 +241,10 @@ void LiveUpdate::restore_environment()
   // enable interrupts again
   asm volatile("sti");
 }
-buffer_t LiveUpdate::store(storage_func func)
+buffer_t LiveUpdate::store()
 {
   char* location = (char*) OS::liveupdate_storage_area();
-  size_t size = update_store_data(location, func, nullptr);
+  size_t size = update_store_data(location, nullptr);
   return buffer_t(location, location + size);
 }
 
@@ -251,17 +267,17 @@ size_t LiveUpdate::stored_data_length(void* location)
   return storage->total_bytes();
 }
 
-size_t update_store_data(void* location, LiveUpdate::storage_func func, const buffer_t* blob)
+size_t update_store_data(void* location, const buffer_t* blob)
 {
   // create storage header in the fixed location
   new (location) storage_header();
   auto* storage = (storage_header*) location;
 
   /// callback for storing stuff, if provided
-  if (func != nullptr)
+  for (const auto& pair : storage_callbacks)
   {
     Storage wrapper {*storage};
-    func(wrapper, blob);
+    pair.second(wrapper, blob);
   }
 
   /// finalize

lib/uplink/ws_uplink.cpp

@@ -47,6 +47,8 @@ namespace uplink {
   {
     OS::add_stdout({this, &WS_uplink::send_log});
 
+    liu::LiveUpdate::register_serialization_callback("uplink", {this, &WS_uplink::store});
+
     read_config();
     CHECK(config_.reboot, "Reboot on panic");
 
@@ -72,7 +74,7 @@ namespace uplink {
     if(liu::LiveUpdate::is_resumable())
     {
       MYINFO("Found resumable state, try restoring...");
-      auto success = liu::LiveUpdate::resume({this, &WS_uplink::restore});
+      auto success = liu::LiveUpdate::resume("uplink", {this, &WS_uplink::restore});
       CHECK(success, "Success");
     }
 
@@ -254,7 +256,7 @@ namespace uplink {
     ws_->close();
     // do the update
     Timers::oneshot(std::chrono::milliseconds(10), [this, buffer] (auto) {
-      liu::LiveUpdate::begin(buffer, {this, &WS_uplink::store});
+      liu::LiveUpdate::begin(buffer);
     });
   }