Patternflow

// ═══════════════════════════════════════════════════════════
// PatternFlow - Browser self-update (drop a .bin onto patternflow.local/update)
//
// Wireless upload without TLS (issue #232). The device already serves plain
// HTTP on the LAN (the audio-react UI on port 80); this module adds /update
// to that same server: a GET page with a drop zone, and a POST handler that
// streams the multipart body straight into Update.h (U_FLASH). The artifact
// is the application image the web build service already emits — the browser
// downloads it over HTTPS from wherever it likes, and the device only ever
// sees a same-origin LAN upload, so there is no certificate to get wrong and
// no mixed-content wall to hit.
//
// Access: by default (PF_WEBUPDATE_ALWAYS_ARMED 1) the POST endpoint accepts
// firmware at any time — drop-a-.bin-whenever is the UX the project chose,
// accepting that anyone on the LAN can flash the device (ArduinoOTA's
// no-password default has the same exposure). Builds for shared networks set
// the flag to 0, and then uploads are only accepted while the device is
// ARMED — the UPDATE screen physically opened from an encoder (hold K2 →
// NETWORK, turn K4). The GET page itself is always served (harmless, and it
// tells you how to arm when arming is required).
//
// Rollback, honestly: once the app has been up for a few seconds, handle()
// calls esp_ota_mark_app_valid_cancel_rollback(). With a bootloader built
// with CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE that makes a crash-looping
// image fall back to the previous slot; the stock Arduino core's bootloader
// does not enable it, so there the call is a harmless no-op and the real
// protection is Update.h's own checks (magic byte, complete write) rejecting
// truncated or wrong files before the boot partition ever changes.
//
// Progress display: the entire POST is consumed inside ONE handleClient()
// call, so the main loop cannot redraw the panel while the image streams in.
// The sketch installs progressCallback to draw from inside the upload
// handler instead (same task, so touching the DMA display is safe).
//
// License: MIT
// ═══════════════════════════════════════════════════════════
#pragma once

#include <Arduino.h>
#include "config.h"
#include "core_wifi.h"
#include "core_names.h"

#if PF_WEBUPDATE_ENABLED
#include <WiFi.h>
#include "webserver/WebServer.h"  // vendored: fixes the 5 s final-chunk stall (see src/webserver/VENDORED.md)
#include <Update.h>
#include <ESPmDNS.h>
#include <esp_ota_ops.h>
#include "core_send.h"
#include "web_update_index.h"
#include "core_http.h"
#endif

namespace PatternflowWebUpdate {

#if PF_WEBUPDATE_ENABLED

// One core-owned console server (core_http.h). This used to fork on
// PF_AUDIO_ENABLED and ride the audio module's server.
constexpr uint16_t HTTP_PORT = PatternflowHttp::HTTP_PORT;
inline WebServer& server() { return PatternflowHttp::server(); }

// How long the app must stay up before this image is declared good (see the
// rollback note in the header comment). Reaching the main loop and staying
// there for this long is the "successful boot" signal.
constexpr uint32_t BOOT_VALID_AFTER_MS = 8000;

inline bool initialized = false;
inline bool armed = false;        // set only from the on-device UPDATE screen
inline bool uploading = false;
inline bool rejected = false;     // this POST arrived while not armed
inline bool completedOk = false;
inline bool bootMarkedValid = false;
inline unsigned progressPct = 0;
inline size_t expectedBytes = 0;
inline size_t receivedBytes = 0;
// Counts POSTs the upload handler actually saw. Distinguishes "the device
// never received the request" from "it received it and got nowhere", which
// look identical from the browser.
inline unsigned uploadAttempts = 0;
inline uint32_t rebootAtMs = 0;   // nonzero = flash landed, reboot scheduled
inline String lastError;

// Installed by the sketch: draws upload progress on the panel from inside
// the upload handler (see header comment). pct is 0..100, or -1 when the
// client sent no ?size= and the total is unknown.
inline void (*progressCallback)(int pct) = nullptr;

inline void arm() {
  armed = true;
  lastError = "";
  Serial.println("[UPDATE] armed — /update accepts firmware");
}

inline void disarm() {
  armed = false;
  Serial.println("[UPDATE] disarmed");
}

// Effective gate: the physical UPDATE screen, or the opt-in always-armed
// build flag (see net_config.h for what that trades away).
inline bool isArmed()         { return armed || PF_WEBUPDATE_ALWAYS_ARMED != 0; }
inline bool isUploading()     { return uploading; }
inline bool isRebootPending() { return rebootAtMs != 0; }
inline bool hasError()        { return lastError.length() > 0; }
inline unsigned progressPercent() { return progressPct; }

inline void failUpload(const char* stage) {
  lastError = String(stage) + ": " + Update.errorString();
  Update.abort();
  uploading = false;
  Serial.printf("[UPDATE] failed — %s\n", lastError.c_str());
}

inline void noteProgress(unsigned pct) {
  if (pct > 100) pct = 100;
  if (pct == progressPct) return;
  progressPct = pct;
  if (progressCallback) progressCallback((int)pct);
}

// Upload body handler — multipart FormData POST (what the /update page uses).
inline void handleUpload() {
  HTTPUpload& up = server().upload();
  switch (up.status) {
    case UPLOAD_FILE_START: {
      uploadAttempts++;
      rejected = !isArmed() || uploading || rebootAtMs != 0;
      if (rejected) {
        Serial.println("[UPDATE] upload refused (not armed)");
        break;
      }
      lastError = "";
      completedOk = false;
      receivedBytes = 0;
      progressPct = 0;
      // The page sends the file size as ?size= because HTTPUpload has no
      // total; with it Update erases exactly what it needs, and the panel
      // can show a real percentage.
      expectedBytes = server().hasArg("size")
                      ? (size_t)strtoul(server().arg("size").c_str(), nullptr, 10)
                      : 0;
      uploading = true;
      // Long flash: do not drop the TCP session on a brief Wi-Fi hitch.
      server().client().setTimeout(120);
      if (!Update.begin(expectedBytes ? expectedBytes : UPDATE_SIZE_UNKNOWN,
                        U_FLASH)) {
        failUpload("begin");
        break;
      }
      Serial.printf("[UPDATE] receiving \"%s\" (%u bytes)\n",
                    up.filename.c_str(), (unsigned)expectedBytes);
      if (progressCallback) progressCallback(expectedBytes ? 0 : -1);
      break;
    }
    case UPLOAD_FILE_WRITE: {
      if (!uploading) break;  // refused or already failed — drain silently
      if (Update.write(up.buf, up.currentSize) != up.currentSize) {
        failUpload("write");
        break;
      }
      receivedBytes += up.currentSize;
      if (expectedBytes) {
        noteProgress((unsigned)((uint64_t)receivedBytes * 100 / expectedBytes));
      }
      break;
    }
    case UPLOAD_FILE_END: {
      if (!uploading) break;
      if (Update.end(true)) {
        completedOk = true;
        progressPct = 100;
        uploading = false;
        Serial.printf("[UPDATE] flashed %u bytes OK\n", (unsigned)receivedBytes);
        if (progressCallback) progressCallback(100);
      } else {
        failUpload("end");
      }
      break;
    }
    case UPLOAD_FILE_ABORTED: {
      if (uploading) {
        Update.abort();
        uploading = false;
        lastError = "upload aborted";
        Serial.println("[UPDATE] upload aborted by client");
      }
      break;
    }
  }
}

// Optional raw PUT (curl --upload-file). Not used by the browser page —
// FormData POST is the path that drives xhr.upload.onprogress reliably.
inline void handleRawBody() {
  HTTPRaw& raw = server().raw();
  switch (raw.status) {
    case RAW_START: {
      uploadAttempts++;
      rejected = !isArmed() || uploading || rebootAtMs != 0;
      if (rejected) {
        Serial.println("[UPDATE] raw upload refused (not armed)");
        break;
      }
      lastError = "";
      completedOk = false;
      receivedBytes = 0;
      progressPct = 0;
      expectedBytes = server().hasArg("size")
                      ? (size_t)strtoul(server().arg("size").c_str(), nullptr, 10)
                      : 0;
      uploading = true;
      server().client().setTimeout(120);
      if (!Update.begin(expectedBytes ? expectedBytes : UPDATE_SIZE_UNKNOWN,
                        U_FLASH)) {
        failUpload("begin");
        break;
      }
      Serial.printf("[UPDATE] receiving raw (%u bytes)\n", (unsigned)expectedBytes);
      if (progressCallback) progressCallback(expectedBytes ? 0 : -1);
      break;
    }
    case RAW_WRITE: {
      if (!uploading) break;
      if (Update.write(raw.buf, raw.currentSize) != raw.currentSize) {
        failUpload("write");
        break;
      }
      receivedBytes += raw.currentSize;
      if (expectedBytes) {
        noteProgress((unsigned)((uint64_t)receivedBytes * 100 / expectedBytes));
      }
      break;
    }
    case RAW_END: {
      if (!uploading) break;
      if (Update.end(true)) {
        completedOk = true;
        progressPct = 100;
        uploading = false;
        Serial.printf("[UPDATE] flashed %u bytes OK (raw)\n", (unsigned)receivedBytes);
        if (progressCallback) progressCallback(100);
      } else {
        failUpload("end");
      }
      break;
    }
    case RAW_ABORTED: {
      if (uploading) {
        Update.abort();
        uploading = false;
        lastError = "upload aborted";
        Serial.println("[UPDATE] raw upload aborted");
      }
      break;
    }
  }
}

// Completion handler — runs after the whole POST body is consumed.
inline void handleUploadDone() {
  if (rejected) {
    server().send(403, "application/json",
        "{\"error\":\"locked - on the device: hold K2 for NETWORK, then turn K4 for UPDATE\"}");
    return;
  }
  if (completedOk) {
    server().send(200, "application/json", "{\"ok\":true}");
    rebootAtMs = millis() + 1200;  // let the response flush before restarting
    return;
  }
  server().send(500, "application/json",
                "{\"error\":\"" + lastError + "\"}");
}

// Reports what the LAST upload attempt actually did, not just whether one is
// running. "Connection lost during upload" is what the browser says for every
// network-level failure, and without these fields there is no way to tell a
// device that never saw a byte from one that took 900 KB and then had the
// client vanish from one that failed the flash write. The counters are only
// cleared when the next upload starts, so they can be read after the fact —
// the same reason /api/status carries loadError.
inline void handleStatus() {
  char buf[224];
  String error = lastError;
  error.replace("\"", "'");
  snprintf(buf, sizeof(buf),
           "{\"armed\":%s,\"busy\":%s,\"version\":\"%s\","
           "\"lastError\":\"%s\",\"lastRejected\":%s,\"lastOk\":%s,"
           "\"received\":%u,\"expected\":%u,\"attempts\":%u}",
           isArmed() ? "true" : "false",
           (uploading || rebootAtMs != 0) ? "true" : "false",
           PF_IMPROV_FW_VERSION,
           error.c_str(),
           rejected ? "true" : "false",
           completedOk ? "true" : "false",
           (unsigned)receivedBytes, (unsigned)expectedBytes, uploadAttempts);
  server().send(200, "application/json", buf);
}

#endif  // PF_WEBUPDATE_ENABLED

inline bool isCompiledIn() {
#if PF_WEBUPDATE_ENABLED
  return true;
#else
  return false;
#endif
}

// Register routes (and, standalone, start the server). Wi-Fi is owned by
// PatternflowWifi; this runs on the connect edge. Idempotent.
inline void begin() {
#if PF_WEBUPDATE_ENABLED
  if (!PatternflowWifi::linkUp()) return;
  if (initialized) {
    // Every reconnect, not just the first: the per-panel mDNS alias carries
    // an address and a new lease may have changed it (core_names.h).
    PatternflowNames::announce();
    return;
  }

  PatternflowNames::announce();

  server().on("/update", HTTP_GET, []() {
    PFSend::gz(server(), WEB_UPDATE_HTML_GZ, WEB_UPDATE_HTML_GZ_LEN);
  });
  // Multipart POST — what the /update page uses (FormData + xhr progress).
  server().on("/update", HTTP_POST, handleUploadDone, handleUpload);
  // Optional raw PUT for curl --upload-file / scripted flashes.
  server().on("/update", HTTP_PUT, handleUploadDone, handleRawBody);
  server().on("/update/status", HTTP_GET, handleStatus);

  PatternflowHttp::begin();  // idempotent; whoever is first starts it

  initialized = true;
  Serial.printf("[UPDATE] Ready — http://%s.local/update (IP %s)\n",
                PF_OTA_HOSTNAME, WiFi.localIP().toString().c_str());
#endif
}

// Called every main loop: services the standalone server (shared-server
// traffic is pumped by PatternflowAudio::handle()), marks the running image
// valid after a healthy stretch of uptime, and fires the deferred reboot.
inline void handle() {
#if PF_WEBUPDATE_ENABLED

  if (!bootMarkedValid && millis() > BOOT_VALID_AFTER_MS) {
    bootMarkedValid = true;
    // No-op unless the bootloader has rollback enabled — see header comment.
    esp_ota_mark_app_valid_cancel_rollback();
  }

  if (rebootAtMs != 0 && (int32_t)(millis() - rebootAtMs) >= 0) {
    Serial.println("[UPDATE] rebooting into new firmware");
    Serial.flush();
    ESP.restart();
  }
#endif
}

#if !PF_WEBUPDATE_ENABLED
// Stubs so the sketch compiles unchanged with the feature off.
inline void (*progressCallback)(int) = nullptr;
inline void arm() {}
inline void disarm() {}
inline bool isArmed()         { return false; }
inline bool isUploading()     { return false; }
inline bool isRebootPending() { return false; }
inline bool hasError()        { return false; }
inline unsigned progressPercent() { return 0; }
#endif

} // namespace PatternflowWebUpdate