esp_alox/main/main.c

#include "client_handler.h"
#include "driver/gpio.h"
#include "driver/uart.h"
#include "esp_app_desc.h"
#include "esp_app_format.h"
#include "esp_err.h"
#include "esp_flash_partitions.h"
#include "esp_image_format.h"
#include "esp_log.h"
#include "esp_log_buffer.h"
#include "esp_ota_ops.h"
#include "esp_partition.h"
#include "esp_phy_init.h"
#include "esp_rom_gpio.h"
#include "esp_timer.h"
#include "esp_wifi.h"
#include "freertos/idf_additions.h"
#include "hal/uart_types.h"
#include "message_handler.h"
#include "message_parser.h"
#include "nvs.h"
#include "nvs_flash.h"

#include "communication_handler.h"
#include "main.h"
#include "ota_update.h"
#include "uart_handler.h"
#include <math.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>

#include "message_builder.h"
#include "uart_msg_ids.h"

static const char *TAG = "ALOX - MAIN";
static const uint16_t version = 0x0001;
static uint8_t send_message_buffer[1024];
static uint8_t send_message_payload_buffer[512];
uint32_t g_uart_firmware_total_size = 0;

static MessageBrokerTaskParams_t broker_task_params;
static MasterOTA_TaskParams_t master_ota_task_params;
static ESP_MessageBrokerTaskParams_t esp_broker_task_params;

ClientList clientList = {.Clients = {{0}}, .ClientCount = 0};

size_t build_ClientInfoPart(uint8_t clientid, float_t lagex, float_t lagey,
                            int32_t bitmask, uint8_t *outputArray,
                            size_t outputArrayOffset, size_t outputArraySize) {

  size_t offset = outputArrayOffset;
  memcpy(&outputArray[offset], &clientid, sizeof(clientid));
  offset += sizeof(clientid);

  // lagex (typischerweise 4 Bytes)
  memcpy(&outputArray[offset], &lagex, sizeof(lagex));
  offset += sizeof(lagex);

  // lagey (typischerweise 4 Bytes)
  memcpy(&outputArray[offset], &lagey, sizeof(lagey));
  offset += sizeof(lagey);

  // bitmask (4 Bytes)
  memcpy(&outputArray[offset], &bitmask, sizeof(bitmask));
  offset += sizeof(bitmask);

  return offset - outputArrayOffset;
}

void fakeDataCallback(uint8_t msgid, const uint8_t *payload, size_t payload_len,
                      uint8_t *send_payload_buffer,
                      size_t send_payload_buffer_size, uint8_t *send_buffer,
                      size_t send_buffer_size) {

  uint8_t seed = payload[1];
  ESP_LOGI(TAG, "Sending Fake Client Infos with seed %d", seed);
  srand(seed);

  size_t offset = 1;
  send_payload_buffer[0] = 3; // Client Count
  offset +=
      build_ClientInfoPart(1, rand() * 1.0, rand() * 1.0, rand() * 1,
                           send_payload_buffer, 1, send_payload_buffer_size);
  offset += build_ClientInfoPart(2, rand() * 2.0, rand() * 2.0, rand() * 2,
                                 send_payload_buffer, offset,
                                 send_payload_buffer_size);
  offset += build_ClientInfoPart(3, rand() * 3.0, rand() * 3.0, rand() * 3,
                                 send_payload_buffer, offset,
                                 send_payload_buffer_size);

  int len = build_message(UART_CLIENT_INPUT, send_payload_buffer, offset,
                          send_buffer, send_buffer_size);

  if (len < 0) {
    ESP_LOGE(TAG,
             "Error Building UART Message: payload_len, %d, sendbuffer_size: "
             "%d, mes_len(error): %d",
             payload_len, send_buffer_size, len);
    return;
  }

  uart_write_bytes(MASTER_UART, send_buffer, len);
}

void echoCallback(uint8_t msgid, const uint8_t *payload, size_t payload_len,
                  uint8_t *send_payload_buffer, size_t send_payload_buffer_size,
                  uint8_t *send_buffer, size_t send_buffer_size) {
  ESP_LOGI(TAG, "Echo command 0x01...");
  int len = build_message(UART_ECHO, payload, payload_len, send_buffer,
                          send_buffer_size);
  if (len < 0) {
    ESP_LOGE(TAG,
             "Error Building UART Message: payload_len, %d, sendbuffer_size: "
             "%d, mes_len(error): %d",
             payload_len, send_buffer_size, len);
    return;
  }

  uart_write_bytes(MASTER_UART, send_buffer, len);
}

void versionCallback(uint8_t msgid, const uint8_t *payload, size_t payload_len,
                     uint8_t *send_payload_buffer,
                     size_t send_payload_buffer_size, uint8_t *send_buffer,
                     size_t send_buffer_size) {
  ESP_LOGI(TAG, "Version command 0x02...");
  size_t git_build_hash_len = strlen(BUILD_GIT_HASH);
  size_t needed_buffer_size = 2 + git_build_hash_len;

  if (send_payload_buffer_size < needed_buffer_size) {
    ESP_LOGE(TAG, "send_payload_buffer to small size %d need %d",
             send_payload_buffer_size, needed_buffer_size);
    return;
  }

  send_payload_buffer[0] = (uint8_t)(version & 0xFF);
  send_payload_buffer[1] = (uint8_t)((version >> 8) & 0xFF);
  memcpy(&send_payload_buffer[2], &BUILD_GIT_HASH, git_build_hash_len);

  int len = build_message(UART_VERSION, send_payload_buffer, needed_buffer_size,
                          send_buffer, send_buffer_size);
  if (len < 0) {
    ESP_LOGE(TAG,
             "Error Building UART Message: payload_len, %d, sendbuffer_size: "
             "%d, mes_len(error): %d",
             payload_len, send_buffer_size, len);
    return;
  }
  uart_write_bytes(MASTER_UART, send_buffer, len);
}

void clientInfoCallback(uint8_t msgid, const uint8_t *payload,
                        size_t payload_len, uint8_t *send_payload_buffer,
                        size_t send_payload_buffer_size, uint8_t *send_buffer,
                        size_t send_buffer_size) {
  ESP_LOGI(TAG, "Client Info Command 0x03...");
  static uint8_t entryLength = 19;
  uint8_t needed_buffer_size = 1 + entryLength * clientList.ClientCount;

  if (send_payload_buffer_size < needed_buffer_size) {
    ESP_LOGE(TAG, "send_payload_buffer to small size %d need %d",
             send_payload_buffer_size, needed_buffer_size);
    return;
  }

  send_payload_buffer[0] = clientList.ClientCount;
  uint8_t offsetMult = 0;

  for (int i = 0; i < MAX_CLIENTS; i++) {
    if (clientList.Clients[i].slotIsUsed) {
      size_t offset = 1 + (entryLength * offsetMult++);
      send_payload_buffer[offset] = i;
      send_payload_buffer[offset + 1] = clientList.Clients[i].isAvailable;
      send_payload_buffer[offset + 2] = clientList.Clients[i].slotIsUsed;
      memcpy(&send_payload_buffer[offset + 3], clientList.Clients[i].macAddr,
             MAC_LENGTH);
      memcpy(&send_payload_buffer[offset + 9], &clientList.Clients[i].lastPing,
             4);
      memcpy(&send_payload_buffer[offset + 13],
             &clientList.Clients[i].lastSuccessfullPing, 4);
      memcpy(&send_payload_buffer[offset + 17],
             &clientList.Clients[i].clientVersion, 2);
    }
  }

  int len = build_message(UART_CLIENT_INFO, send_payload_buffer,
                          needed_buffer_size, send_buffer, send_buffer_size);

  if (len < 0) {
    ESP_LOGE(TAG,
             "Error Building UART Message: payload_len, %d, sendbuffer_size: "
             "%d, mes_len(error): %d",
             needed_buffer_size, send_buffer_size, len);
    return;
  }
  uart_write_bytes(MASTER_UART, send_buffer, len);
}

bool g_ota_in_progress = false;

void ota_monitor_task(void *param) {
  const TickType_t timeout = pdMS_TO_TICKS(10000); // 10 seconds

  while (g_ota_in_progress) {
    bool all_done = true;
    for (int i = 0; i < MAX_CLIENTS; i++) {
      if (clientList.Clients[i].slotIsUsed) {
        if (clientList.Clients[i].ota_status != OTA_SUCCESS &&
            clientList.Clients[i].ota_status != OTA_FAILED) {
          all_done = false;

          if ((xTaskGetTickCount() - clientList.Clients[i].last_seen) >
              timeout) {
            ESP_LOGE(TAG, "Client %d timed out", i);
            clientList.Clients[i].ota_status = OTA_FAILED;
          }
        }
      }
    }

    if (all_done) {
      g_ota_in_progress = false;
    }

    vTaskDelay(pdMS_TO_TICKS(1000));
  }

  ESP_LOGI(TAG, "OTA Operation Finished.");
  for (int i = 0; i < MAX_CLIENTS; i++) {
    if (clientList.Clients[i].slotIsUsed) {
      ESP_LOGI(TAG, "Client %d [MAC: " MACSTR "]: %s, Resent Chunks: %d", i,
               MAC2STR(clientList.Clients[i].macAddr),
               clientList.Clients[i].ota_status == OTA_SUCCESS ? "SUCCESS"
                                                               : "FAILED",
               clientList.Clients[i].resent_chunks_counter);
    }
  }

  vTaskDelete(NULL);
}

void send_client_ota_start_message(uint8_t clientID, uint32_t app_size) {
  BaseMessage message = {};
  OTA_PREPARE_FOR_UPDATE_Payload ota_payload = {
      .total_size = app_size,
  };
  message = MessageBuilder(OTA_PREPARE_FOR_UPDATE,
                           *(PayloadUnion *)&ota_payload, sizeof(ota_payload));

  esp_err_t err = esp_now_send(clientList.Clients[clientID].macAddr,
                               (uint8_t *)&message, sizeof(BaseMessage));
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Could not send OTA PREPARE FOR UPDATE to " MACSTR ", %s",
             MAC2STR(clientList.Clients[clientID].macAddr),
             esp_err_to_name(err));
  } else {
    ESP_LOGI(TAG, "Sent OTA PREPARE FOR UPDATE to " MACSTR,
             MAC2STR(clientList.Clients[clientID].macAddr));
  }
}

void app_main(void) {
  ESP_LOGI(TAG, "Starting Alox Powerpod Version %d Build: %s", version,
           BUILD_GIT_HASH);

  esp_err_t ret = nvs_flash_init();
  if (ret == ESP_ERR_NVS_NO_FREE_PAGES ||
      ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
    ESP_ERROR_CHECK(nvs_flash_erase());
    ret = nvs_flash_init();
  }
  ESP_ERROR_CHECK(ret);

  gpio_reset_pin(MASTER_MODE_PIN);
  gpio_set_direction(MASTER_MODE_PIN, GPIO_MODE_INPUT);
  bool isMaster = (gpio_get_level(MASTER_MODE_PIN) == 0);

  ESP_ERROR_CHECK(esp_netif_init());
  ESP_ERROR_CHECK(esp_event_loop_create_default());
  wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
  ESP_ERROR_CHECK(esp_wifi_init(&cfg));
  wifi_config_t wifi_config = {
      .sta =
          {
              .channel = 1,
          },
  };
  ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
  ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
  ESP_ERROR_CHECK(esp_wifi_start());
  ESP_ERROR_CHECK(esp_now_init());

  if (isMaster) {
    ESP_ERROR_CHECK(esp_now_register_recv_cb(master_receive_callback));
  } else {
    ESP_ERROR_CHECK(esp_now_register_recv_cb(client_receive_callback));
  }

  ret = init_com(&clientList, 1);
  if (ret < 0) {
    ESP_LOGE(TAG, "Could not Init ESP NOW Communication!");
  }

  const esp_partition_t *running = esp_ota_get_running_partition();
  ESP_LOGI(TAG, "OTA: Running Partition: %s", running->label);

  esp_ota_img_states_t ota_state;
  if (esp_ota_get_state_partition(running, &ota_state) == ESP_OK) {
    if (ota_state == ESP_OTA_IMG_PENDING_VERIFY) {
      bool diagnostic_is_ok = true; // TODO build in valid diagnostics
      if (diagnostic_is_ok) {
        esp_ota_mark_app_valid_cancel_rollback();
      } else {
        // esp_ota_mark_app_invalid_rollback(); Put this function at the start
        // so when the esp crashes it can rollback
        esp_ota_mark_app_invalid_rollback_and_reboot();
      }
    }
  }

  const esp_partition_t *next_ota_partition =
      esp_ota_get_next_update_partition(NULL);
  int app_size = get_app_size(next_ota_partition);
  ESP_LOGE(TAG, "App Size in Other Partition %d", app_size);

  QueueHandle_t espnow_message_queue =
      xQueueCreate(10, sizeof(ESPNOW_MessageInfo));
  ESP_InitMessageBroker(espnow_message_queue);

  esp_broker_task_params.message_queue = espnow_message_queue;
  xTaskCreate(ESP_MessageBrokerTask, "espnow_message_broker_task", 4096,
              (void *)&esp_broker_task_params, 4, NULL);

  init_ota();

  if (isMaster) {
    ESP_LOGI(TAG, "Started in Mastermode");
    ESPNOW_RegisterMasterCallbacks();
    ESPNOW_RegisterOTAMaster();

    add_peer(broadcast_address);
    xTaskCreate(master_broadcast_task, "MasterBroadcast", 4096, NULL, 1, NULL);
    xTaskCreate(master_broadcast_ping, "MasterBroadcastPing", 4096, NULL, 1,
                NULL);

    QueueHandle_t parsed_message_queue =
        xQueueCreate(10, sizeof(ParsedMessage_t));
    init_uart(parsed_message_queue);
    InitMessageBroker();

    broker_task_params.message_queue = parsed_message_queue;
    broker_task_params.send_buffer = send_message_buffer;
    broker_task_params.send_buffer_size = sizeof(send_message_buffer);
    broker_task_params.payload_buffer = send_message_payload_buffer;
    broker_task_params.payload_buffer_size =
        sizeof(send_message_payload_buffer);

    xTaskCreate(MessageBrokerTask, "MessageHandler", 4096,
                (void *)&broker_task_params, 5, NULL);

    master_ota_task_params.client_list = &clientList;
    xTaskCreate(MasterOTATask, "MasterOTATask", 4096,
                (void *)&master_ota_task_params, 4, NULL);

    RegisterCallback(UART_ECHO, echoCallback);
    RegisterCallback(UART_VERSION, versionCallback);
    RegisterCallback(UART_CLIENT_INFO, clientInfoCallback);
    RegisterCallback(UART_CLIENT_INPUT, fakeDataCallback);
    RegisterCallback(UART_OTA_START_ESPNOW, start_ota_update_espnow);

  } else {
    ESP_LOGI(TAG, "Started in Slavemode");
    ESPNOW_RegisterSlaveCallbacks();
    xTaskCreate(slave_ota_task, "SlaveOTATask", 4096, NULL, 4, NULL);
    ESPNOW_RegisterOTASlave();
  }
}