Flipdot-display/src/esp32/system/firmware_guard.cpp

#include "system/firmware_guard.h"

#include <Arduino.h>
#include <esp_ota_ops.h>
#include <esp_system.h>
#include <esp_task_wdt.h>

#include "config.h"

// Survives crashes and software resets, not a power loss
static const uint32_t RTC_MAGIC = 0xF11D0C0A;
RTC_NOINIT_ATTR static uint32_t rtc_magic;
RTC_NOINIT_ATTR static uint8_t rtc_crash_count;

static uint32_t main_start = 0; // millis() when the main code started
static bool watchdog_running = false;
static bool firmware_pending = false; // new OTA firmware, not yet marked as good
static bool crash_count_cleared = false;
static char status[64];

// The Arduino core marks a new OTA firmware as good right at boot. Returning true here
// leaves that to us, so the bootloader can roll back a firmware that fails early.
extern "C" bool verifyRollbackLater()
{
	return true;
}

static ResetKind resetKind()
{
	switch (esp_reset_reason())
	{
	case ESP_RST_POWERON:
		return ResetKind::PowerOn;
	case ESP_RST_SW:
		return ResetKind::Software;
	case ESP_RST_PANIC:
	case ESP_RST_INT_WDT:
	case ESP_RST_TASK_WDT:
	case ESP_RST_WDT:
		return ResetKind::Crash;
	default:
		return ResetKind::Other;
	}
}

static void setCrashCount(uint8_t count)
{
	rtc_magic = RTC_MAGIC;
	rtc_crash_count = count;
}

static const char *firmwareState()
{
	esp_ota_img_states_t state;
	if (esp_ota_get_state_partition(esp_ota_get_running_partition(), &state) != ESP_OK)
		return "usb";
	switch (state)
	{
	case ESP_OTA_IMG_PENDING_VERIFY:
		return "testing";
	case ESP_OTA_IMG_VALID:
		return "ok";
	case ESP_OTA_IMG_NEW:
		// The bootloader turns NEW into PENDING_VERIFY when it supports rollback
		return "ok, no rollback in bootloader";
	default:
		return "usb";
	}
}

static void updateStatus()
{
	snprintf(status, sizeof(status), "%s, %s, %u crashes", FIRMWARE_VERSION, firmwareState(),
			 rtc_magic == RTC_MAGIC ? rtc_crash_count : 0);
}

BootMode firmwareGuardBegin()
{
	uint8_t count = rtc_magic == RTC_MAGIC ? rtc_crash_count : 0;
	BootDecision decision = decideBoot(resetKind(), count);
	setCrashCount(decision.crash_count);

	esp_ota_img_states_t state;
	firmware_pending = esp_ota_get_state_partition(esp_ota_get_running_partition(), &state) == ESP_OK &&
					   state == ESP_OTA_IMG_PENDING_VERIFY;

	updateStatus();
	Serial.printf("Firmware %s, reset reason %d, crash count %u\n", status, esp_reset_reason(),
				  decision.crash_count);
	return decision.mode;
}

void firmwareGuardSafeModeEnded()
{
	setCrashCount(crashCountAfterSafeMode());
}

void firmwareGuardStartWatchdog()
{
	// Reconfigures the watchdog the ESP-IDF already started (default 5 s)
	esp_task_wdt_init(WATCHDOG_TIMEOUT_S, true);
	esp_task_wdt_add(NULL);
	watchdog_running = true;
	main_start = millis();
}

void firmwareGuardStopWatchdog()
{
	if (watchdog_running)
		esp_task_wdt_delete(NULL);
	watchdog_running = false;
}

void firmwareGuardLoop(bool wifi_connected)
{
	if (watchdog_running)
		esp_task_wdt_reset();

	uint32_t running = millis() - main_start;

	if (firmware_pending && shouldMarkFirmwareValid(running, wifi_connected))
	{
		esp_ota_mark_app_valid_cancel_rollback();
		firmware_pending = false;
		updateStatus();
		Serial.println("New firmware marked as good");
	}

	if (!crash_count_cleared && shouldClearCrashCount(running))
	{
		setCrashCount(0);
		crash_count_cleared = true;
		updateStatus();
	}
}

const char *firmwareGuardStatus()
{
	return status;
}