ota.hpp

Go to the documentation of this file.

#ifndef _PAINLESS_MESH_PLUGIN_OTA_HPP_
#define _PAINLESS_MESH_PLUGIN_OTA_HPP_
 
#include "painlessmesh/configuration.hpp"
 
#include "painlessmesh/base64.hpp"
#include "painlessmesh/logger.hpp"
#include "painlessmesh/plugin.hpp"
 
#if !defined(USE_FS_SPIFFS) && !defined(USE_FS_LITTLEFS)
#if defined(ESP32) || defined(ESP8266)
#ifdef ESP32
#define USE_FS_SPIFFS
#else
#define USE_FS_LITTLEFS
#endif
#endif
#endif
 
#ifdef ESP32
#include <Update.h>
#endif
 
#if defined(USE_FS_SPIFFS)
#include <SPIFFS.h>
#elif defined(USE_FS_LITTLEFS)
#include <LittleFS.h>
#endif
 
namespace painlessmesh {
namespace plugin {
 
namespace ota {
 
enum class OTA_OP_CODES {
  ANNOUNCE = 10,      // Announce a new update
  DATA_REQUEST = 11,  // Request data from host
  DATA = 12,          // Inbound data to nodes
};
 
class Announce : public BroadcastPackage {
 public:
  TSTRING md5;
  TSTRING hardware;
 
  TSTRING role;
 
  bool forced = false;
 
  bool broadcasted = false;
 
  bool compressed = false;
 
  size_t noPart = 0;
 
  Announce() : BroadcastPackage(10) {}
 
  Announce(JsonObject jsonObj) : BroadcastPackage(jsonObj) {
    md5 = jsonObj["md5"].as<TSTRING>();
    hardware = jsonObj["hardware"].as<TSTRING>();
    role = jsonObj["role"].as<TSTRING>();
#if ARDUINOJSON_VERSION_MAJOR < 7
    if (jsonObj.containsKey("forced")) forced = jsonObj["forced"];
    if (jsonObj.containsKey("broadcasted"))
      broadcasted = jsonObj["broadcasted"];
    if (jsonObj.containsKey("compressed"))
      compressed = jsonObj["compressed"];
#else
    if (jsonObj["forced"].is<bool>()) forced = jsonObj["forced"];
    if (jsonObj["broadcasted"].is<bool>()) broadcasted = jsonObj["broadcasted"];
    if (jsonObj["compressed"].is<bool>()) compressed = jsonObj["compressed"];
#endif
    noPart = jsonObj["noPart"];
  }
 
  JsonObject addTo(JsonObject&& jsonObj) const {
    jsonObj = BroadcastPackage::addTo(std::move(jsonObj));
    jsonObj["md5"] = md5;
    jsonObj["hardware"] = hardware;
    jsonObj["role"] = role;
    if (forced) jsonObj["forced"] = forced;
    jsonObj["noPart"] = noPart;
    jsonObj["broadcasted"] = broadcasted;
    if (compressed) jsonObj["compressed"] = compressed;
    return jsonObj;
  }
 
#if ARDUINOJSON_VERSION_MAJOR < 7
  size_t jsonObjectSize() const {
    return JSON_OBJECT_SIZE(noJsonFields + 5) +
           round(1.1 * (md5.length() + hardware.length() + role.length()));
  }
#endif
 
 protected:
  Announce(int type, router::Type routing) : BroadcastPackage(type) {
    this->routing = routing;
  }
};
 
class Data;
 
class DataRequest : public Announce {
 public:
  size_t partNo = 0;
  uint32_t dest = 0;
 
  DataRequest() : Announce(11, router::SINGLE) {}
 
  DataRequest(JsonObject jsonObj) : Announce(jsonObj) {
    dest = jsonObj["dest"];
    partNo = jsonObj["partNo"];
  }
 
  JsonObject addTo(JsonObject&& jsonObj) const {
    jsonObj = Announce::addTo(std::move(jsonObj));
    jsonObj["dest"] = dest;
    jsonObj["partNo"] = partNo;
    return jsonObj;
  }
 
  static DataRequest replyTo(const Announce& ann, uint32_t from,
                             size_t partNo) {
    DataRequest req;
    req.dest = ann.from;
    req.md5 = ann.md5;
    req.hardware = ann.hardware;
    req.role = ann.role;
    req.forced = ann.forced;
    req.noPart = ann.noPart;
    req.partNo = partNo;
    req.from = from;
    req.broadcasted = ann.broadcasted;
    req.compressed = ann.compressed;
    return req;
  }
 
  static DataRequest replyTo(const Data& d, size_t partNo);
 
#if ARDUINOJSON_VERSION_MAJOR < 7
  size_t jsonObjectSize() const {
    return JSON_OBJECT_SIZE(noJsonFields + 5 + 2) +
           round(1.1 * (md5.length() + hardware.length() + role.length()));
  }
#endif
 
 protected:
  DataRequest(int type) : Announce(type, router::SINGLE) {}
};
 
class Data : public DataRequest {
 public:
  TSTRING data;
 
  Data() : DataRequest(12) {}
 
  Data(JsonObject jsonObj) : DataRequest(jsonObj) {
    data = jsonObj["data"].as<TSTRING>();
  }
 
  JsonObject addTo(JsonObject&& jsonObj) const {
    jsonObj = DataRequest::addTo(std::move(jsonObj));
    jsonObj["data"] = data;
    return jsonObj;
  }
 
  static Data replyTo(const DataRequest& req, TSTRING data, size_t partNo) {
    Data d;
    d.from = req.dest;
    d.dest = req.from;
    d.md5 = req.md5;
    d.hardware = req.hardware;
    d.role = req.role;
    d.forced = req.forced;
    d.noPart = req.noPart;
    d.partNo = partNo;
    d.data = data;
    d.broadcasted = req.broadcasted;
    d.compressed = req.compressed;
    // Phase 2: Set routing to BROADCAST for true broadcast mode
    if (req.broadcasted) {
      d.routing = router::BROADCAST;
    }
    return d;
  }
 
#if ARDUINOJSON_VERSION_MAJOR < 7
  size_t jsonObjectSize() const {
    return JSON_OBJECT_SIZE(noJsonFields + 5 + 2 + 1) +
           round(1.1 * (md5.length() + hardware.length() + role.length() +
                        data.length()));
  }
#endif
};
 
inline DataRequest DataRequest::replyTo(const Data& d, size_t partNo) {
  DataRequest req;
  req.from = d.dest;
  req.dest = d.from;
  req.md5 = d.md5;
  req.hardware = d.hardware;
  req.role = d.role;
  req.forced = d.forced;
  req.noPart = d.noPart;
  req.partNo = partNo;
  req.broadcasted = d.broadcasted;
  req.compressed = d.compressed;
  return req;
}
 
class State : public protocol::PackageInterface {
 public:
  TSTRING md5;
#ifdef ESP32
  TSTRING hardware = "ESP32";
#else
  TSTRING hardware = "ESP8266";
#endif
  TSTRING role;
  size_t noPart = 0;
  size_t partNo = 0;
  bool broadcasted = false;
  bool compressed = false;
  TSTRING ota_fn = "/ota_fw.json";
 
  State() {}
 
  State(JsonObject jsonObj) {
    md5 = jsonObj["md5"].as<TSTRING>();
    hardware = jsonObj["hardware"].as<TSTRING>();
    role = jsonObj["role"].as<TSTRING>();
    broadcasted = jsonObj["broadcasted"].as<bool>();
    compressed = jsonObj["compressed"].as<bool>();
  }
 
  State(const Announce& ann) {
    md5 = ann.md5;
    hardware = ann.hardware;
    role = ann.role;
    noPart = ann.noPart;
    broadcasted = ann.broadcasted;
    compressed = ann.compressed;
  }
 
  JsonObject addTo(JsonObject&& jsonObj) const {
    jsonObj["role"] = role;
    jsonObj["md5"] = md5;
    jsonObj["hardware"] = hardware;
    jsonObj["broadcasted"] = broadcasted;
    jsonObj["compressed"] = compressed;
    return jsonObj;
  }
 
#if ARDUINOJSON_VERSION_MAJOR < 7
  size_t jsonObjectSize() const {
    return JSON_OBJECT_SIZE(3) +
           round(1.1 * (md5.length() + hardware.length() + role.length()));
  }
#endif
 
  std::shared_ptr<Task> task;
};
 
typedef std::function<size_t(painlessmesh::plugin::ota::DataRequest,
                             char* buffer)>
    otaDataPacketCallbackType_t;
 
template <class T>
void addSendPackageCallback(Scheduler& scheduler,
                            plugin::PackageHandler<T>& mesh,
                            otaDataPacketCallbackType_t callback,
                            size_t otaPartSize) {
  using namespace logger;
#if defined(ESP32) || defined(ESP8266)
  mesh.onPackage(
      (int)OTA_OP_CODES::DATA_REQUEST,
      [&mesh, callback, otaPartSize](painlessmesh::protocol::Variant& variant) {
        auto pkg = variant.to<painlessmesh::plugin::ota::DataRequest>();
        char buffer[otaPartSize + 1];
        memset(buffer, 0, otaPartSize + 1);
        auto size = callback(pkg, buffer);
        // Handle zero size
        if (!size) {
          // No data is available by the user app.
 
          // todo - doubtful, shall we return true or false. What is the purpose
          // of this return value.
          return true;
        }
        // Encode data as base64 so there are no null characters and can be
        // shown in plaintext
        auto b64Data =
            painlessmesh::base64::encode((unsigned char*)buffer, size);
        auto reply =
            painlessmesh::plugin::ota::Data::replyTo(pkg, b64Data, pkg.partNo);
        
        // Phase 2: Routing is now handled in replyTo() based on broadcasted flag
        mesh.sendPackage(&reply);
        if (pkg.broadcasted) {
          Log(DEBUG, "OTA: Broadcasting chunk %d/%d\n", pkg.partNo, pkg.noPart);
        }
        return true;
      });
 
#endif
}
 
template <class T>
void addReceivePackageCallback(
    Scheduler& scheduler, plugin::PackageHandler<T>& mesh, TSTRING role = "",
    std::function<void(int, int)> progress_cb = NULL) {
  using namespace logger;
#if defined(ESP32) || defined(ESP8266)
  auto currentFW = std::make_shared<State>();
  currentFW->role = role;
  auto updateFW = std::make_shared<State>();
  updateFW->role = role;
#ifdef USE_FS_SPIFFS
  SPIFFS.begin(true);  // Start the SPI Flash Files System
  if (SPIFFS.exists(currentFW->ota_fn)) {
    auto file = SPIFFS.open(currentFW->ota_fn, "r");
#else
  LittleFS.begin();  // Start the SPI Flash Files System
  if (LittleFS.exists(currentFW->ota_fn)) {
    auto file = LittleFS.open(currentFW->ota_fn, "r");
#endif
    TSTRING msg = "";
    while (file.available()) {
      msg += (char)file.read();
    }
    protocol::Variant var(msg);
    auto fw = var.to<State>();
    if (fw.role == role && fw.hardware == currentFW->hardware) {
      Log(DEBUG, "MD5 found %s\n", fw.md5.c_str());
      currentFW->md5 = fw.md5;
    }
  }
 
  mesh.onPackage((int)OTA_OP_CODES::ANNOUNCE, [currentFW, updateFW, &mesh,
                                               &scheduler](
                                                  protocol::Variant& variant) {
    // convert variant to Announce
    auto pkg = variant.to<Announce>();
    // Check if we want the update
    if (currentFW->role == pkg.role && currentFW->hardware == pkg.hardware) {
      if ((currentFW->md5 == pkg.md5 && !pkg.forced) ||
          updateFW->md5 == pkg.md5) {
        // Either already have it, or already updating to it
        return false;
      } else {
        updateFW->md5 = pkg.md5;
        updateFW->broadcasted = pkg.broadcasted;
        updateFW->compressed = pkg.compressed;
        // If we are not in broadcasted mode, or we are the root node, begin
        // requesting data
        if (!pkg.broadcasted || mesh.isRoot()) {
          auto request =
              DataRequest::replyTo(pkg, mesh.getNodeId(), updateFW->partNo);
          // enable the request task
          updateFW->task =
              mesh.addTask(scheduler, 30 * TASK_SECOND, 10,
                           [request, &mesh]() { mesh.sendPackage(&request); });
          updateFW->task->setOnDisable([updateFW]() {
            Log(ERROR, "OTA: Did not receive the requested data.\n");
            updateFW->md5 = "";
          });
        }
      }
    }
    return false;
  });
 
  // mesh.onPackage(11, [currentFW](protocol::Variant &variant) {
  //   Log(ERROR, "Data request should not be send to this node\n");
  //   return false;
  // });
 
  mesh.onPackage((int)OTA_OP_CODES::DATA, [currentFW, updateFW, progress_cb,
                                           &mesh, &scheduler](
                                              protocol::Variant& variant) {
    auto pkg = variant.to<Data>();
    // Check whether it is a new part, of correct md5 role etc etc
    if (updateFW->md5 == pkg.md5 && updateFW->role == pkg.role &&
        updateFW->hardware == pkg.hardware) {
      // Have we have received the next chunk of firmware in sequence?
      if (updateFW->partNo == pkg.partNo) {
        if (progress_cb != NULL) {
          progress_cb(pkg.partNo, pkg.noPart);
        }
        if (pkg.partNo == 0) {
#ifdef ESP32
          uint32_t maxSketchSpace = UPDATE_SIZE_UNKNOWN;
#else
          uint32_t maxSketchSpace =
                  (ESP.getFreeSketchSpace() - 0x1000) & 0xFFFFF000;
#endif
          Log(DEBUG, "Sketch size %d\n", maxSketchSpace);
          if (Update.isRunning()) {
            Update.end(false);
          }
          if (!Update.begin(maxSketchSpace)) {  // start with max available size
            Log(DEBUG, "handleOTA(): OTA start failed!");
            Update.printError(Serial);
            Update.end();
          } else {
            Update.setMD5(pkg.md5.c_str());
          }
        }
 
        //    write data
        auto b64Data = base64::decode(pkg.data);
        if (Update.write((uint8_t*)b64Data.c_str(), b64Data.length()) !=
            b64Data.length()) {
          Log(ERROR, "handleOTA(): OTA write failed!");
          Update.printError(Serial);
          Update.end();
          updateFW->md5 = "";
          updateFW->partNo = 0;
          return false;
        }
 
        // If last part then write ota_fn and reboot
        if (pkg.partNo == pkg.noPart - 1) {
          // check md5, reboot
          if (Update.end(true)) {  // true to set the size to the
                                   // current progress
#ifdef USE_FS_SPIFFS
            auto file = SPIFFS.open(updateFW->ota_fn, "w");
            if (!file) {
              Log(ERROR,
                  "handleOTA(): Unable to write md5 of new update to the "
                  "SPIFFS "
                  "file. This will result in endless update loops for OTA\n");
            }
#else
            auto file = LittleFS.open(updateFW->ota_fn, "w");
            if (!file) {
              Log(ERROR, "handleOTA(): Unable to write md5 of new binary to the LittleFS file. This will result in endless update loops for OTA\n");
            }
#endif
 
            String msg;
            protocol::Variant var(updateFW.get());
            var.printTo(msg);
            file.print(msg);
            file.close();
 
            Log(DEBUG, "handleOTA(): OTA Success! %s, %s\n", msg.c_str(),
                updateFW->role.c_str());
            // Delay restart by 2 seconds to allow mesh activity to finish
            mesh.addTask(scheduler, 2 * TASK_SECOND, TASK_ONCE,
                         []() { ESP.restart(); })
                ->enableDelayed();
          } else {
            Log(DEBUG, "handleOTA(): OTA failed!\n");
            Update.printError(Serial);
            updateFW->md5 = "";
            updateFW->partNo = 0;
          }
          if (updateFW->task != NULL) {
            updateFW->task->setOnDisable(NULL);
            updateFW->task->disable();
          }
        } else {
          // else request more
          ++updateFW->partNo;
          if (updateFW->task != NULL) {
            auto request = DataRequest::replyTo(pkg, updateFW->partNo);
            updateFW->task->setCallback(
                [request, &mesh]() { mesh.sendPackage(&request); });
            // updateFW->task->disable();
            updateFW->task->restart();
          }
        }
      } else {
        // We are out-of-sequence, resume with non-broadcasted update
        if (updateFW->broadcasted && pkg.broadcasted) {
          Log(DEBUG, "Out of sequence packet! We may have missed a packet?");
          // Drop of out broadcasted mode
          updateFW->broadcasted = false;
          auto request = DataRequest::replyTo(pkg, updateFW->partNo);
          request.broadcasted = false;
          updateFW->task =
              mesh.addTask(scheduler, 30 * TASK_SECOND, 10,
                           [request, &mesh]() { mesh.sendPackage(&request); });
          updateFW->task->setOnDisable([updateFW]() {
            Log(ERROR, "OTA: Did not receive the requested data.\n");
            updateFW->md5 = "";
          });
        }
      }
    }
    return false;
  });
 
#endif
}
 
}  // namespace ota
}  // namespace plugin
}  // namespace painlessmesh
 
#endif