battleship/src/main.c

#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>

#include "app_config.h"
#include "application.h"
#include "esp_check.h"
#include "esp_event.h"
#include "esp_heap_caps.h"
#include "esp_http_server.h"
#include "esp_littlefs.h"
#include "esp_log.h"
#include "esp_netif.h"
#include "esp_random.h"
#include "esp_timer.h"
#include "esp_wifi.h"
#include "freertos/FreeRTOS.h"
#include "freertos/portmacro.h"
#include "http_api.h"
#include "nvs_flash.h"
#include "sync_service.h"

#if __has_include("wifi_config.h")
#include "wifi_config.h"
#define WIFI_CONFIG_AVAILABLE 1
#else
#define WIFI_CONFIG_AVAILABLE 0
#endif

static const char *const kLogTag = "battleship";
static const char *const kLittlefsBasePath = "/littlefs";
static const char *const kLittlefsPartitionLabel = "littlefs";
static const size_t kFileChunkBytes = 1024U;
static const uint8_t kHttpMaxOpenSockets = 12U;

typedef struct { bool configured; bool connected; bool retry_scheduled; uint8_t reconnect_attempt; } wifi_state_t;
static portMUX_TYPE s_wifi_lock = portMUX_INITIALIZER_UNLOCKED;
static portMUX_TYPE s_diagnostics_lock = portMUX_INITIALIZER_UNLOCKED;
static wifi_state_t s_wifi_state = {0};
static bool s_littlefs_mounted;
static httpd_handle_t s_server;
static application_t s_application;
static http_api_t s_api;
static sync_service_t s_sync;
static uint16_t s_rejected_input;
static esp_timer_handle_t s_sync_timer;
static esp_timer_handle_t s_bot_timer;
#if WIFI_CONFIG_AVAILABLE
static esp_timer_handle_t s_reconnect_timer;
#endif

static uint32_t platform_random(void *unused) { (void)unused; return esp_random(); }

static void record_rejected_input(void) {
  portENTER_CRITICAL(&s_diagnostics_lock);
  if (s_rejected_input < UINT16_MAX) ++s_rejected_input;
  portEXIT_CRITICAL(&s_diagnostics_lock);
}

static uint16_t rejected_input_snapshot(void) {
  uint16_t rejected = 0;
  portENTER_CRITICAL(&s_diagnostics_lock);
  rejected = s_rejected_input;
  portEXIT_CRITICAL(&s_diagnostics_lock);
  return rejected;
}

static wifi_state_t wifi_state_snapshot(void) {
  wifi_state_t snapshot;
  portENTER_CRITICAL(&s_wifi_lock);
  snapshot = s_wifi_state;
  portEXIT_CRITICAL(&s_wifi_lock);
  return snapshot;
}

static const char *wifi_state_name(const wifi_state_t *state) {
  if (!state->configured) return "not_configured";
  return state->connected ? "connected" : "connecting";
}

static void refresh_health(void) {
  const wifi_state_t wifi_state = wifi_state_snapshot();
  size_t clients = kHttpMaxOpenSockets;
  int client_fds[kHttpMaxOpenSockets];
  if (s_server == NULL || httpd_get_client_list(s_server, &clients, client_fds) != ESP_OK) clients = 0U;
  http_api_set_health(&s_api, &(http_api_health_t){
      .uptime_ms = (uint32_t)(esp_timer_get_time() / 1000U),
      .free_heap_bytes = esp_get_free_heap_size(),
      .minimum_free_heap_bytes = esp_get_minimum_free_heap_size(),
      .largest_free_block_bytes = heap_caps_get_largest_free_block(MALLOC_CAP_8BIT),
      .connected_clients = (uint8_t)clients,
      .rejected_input = rejected_input_snapshot(),
      .reset_reason = (int8_t)esp_reset_reason(),
      .wifi_state = wifi_state_name(&wifi_state),
  });
}

static void set_status(httpd_req_t *request, uint16_t status) {
  const char *const values[] = {"200 OK", "400 Bad Request", "401 Unauthorized", "403 Forbidden", "404 Not Found", "409 Conflict", "413 Payload Too Large", "503 Service Unavailable"};
  const uint16_t codes[] = {200U, 400U, 401U, 403U, 404U, 409U, 413U, 503U};
  for (size_t index = 0; index < sizeof(codes) / sizeof(codes[0]); ++index) if (status == codes[index]) { httpd_resp_set_status(request, values[index]); return; }
  httpd_resp_set_status(request, "503 Service Unavailable");
}

static esp_err_t send_api_response(httpd_req_t *request, const http_api_response_t *response) {
  set_status(request, response->status);
  httpd_resp_set_type(request, "application/json");
  httpd_resp_set_hdr(request, "Cache-Control", "no-store");
  return httpd_resp_send(request, response->body, response->body_length);
}

static bool websocket_send(void *unused, int client_id, const char *payload, size_t length) {
  (void)unused;
  httpd_ws_frame_t frame = {.final = true, .type = HTTPD_WS_TYPE_TEXT, .payload = (uint8_t *)payload, .len = length};
  const esp_err_t result = httpd_ws_send_frame_async(s_server, client_id, &frame);
  if (result != ESP_OK) httpd_sess_trigger_close(s_server, client_id);
  return result == ESP_OK;
}

static void sync_broadcast_work(void *unused) {
  (void)unused;
  sync_service_broadcast(&s_sync, websocket_send, NULL);
}

static void queue_state_broadcast(void) {
  if (s_server != NULL && httpd_queue_work(s_server, sync_broadcast_work, NULL) != ESP_OK) record_rejected_input();
}

static bool schedule_bot_turn(void *unused, uint32_t delay_ms) {
  (void)unused;
  return s_bot_timer != NULL && esp_timer_start_once(s_bot_timer, (uint64_t)delay_ms * 1000U) == ESP_OK;
}

static void bot_turn_work(void *unused) {
  (void)unused;
  const uint32_t version_before = s_application.lifecycle.game.state.version;
  if (application_bot_take_turn(&s_application) && s_application.lifecycle.game.state.version != version_before) queue_state_broadcast();
}

static void bot_timer_callback(void *unused) {
  (void)unused;
  if (s_server != NULL && httpd_queue_work(s_server, bot_turn_work, NULL) != ESP_OK) record_rejected_input();
}

static void sync_expire_work(void *unused) {
  (void)unused;
  int closed[kSessionCapacity] = {0};
  size_t closed_count = 0U;
  sync_service_expire(&s_sync, (uint64_t)(esp_timer_get_time() / 1000U), closed, &closed_count);
  for (size_t index = 0; index < closed_count; ++index) httpd_sess_trigger_close(s_server, closed[index]);
}

static void sync_timer_callback(void *unused) {
  (void)unused;
  if (s_server != NULL) httpd_queue_work(s_server, sync_expire_work, NULL);
}

static bool request_has_json_content_type(httpd_req_t *request) {
  const size_t length = httpd_req_get_hdr_value_len(request, "Content-Type");
  if (length == 0U || length >= 64U) return false;
  char value[64];
  return httpd_req_get_hdr_value_str(request, "Content-Type", value, sizeof(value)) == ESP_OK && strncmp(value, "application/json", 16U) == 0;
}

static size_t route_body_limit(http_api_route_t route) {
  switch (route) {
    case HTTP_API_ROUTE_JOIN: return 192U;
    case HTTP_API_ROUTE_RESUME:
    case HTTP_API_ROUTE_CONFIG:
    case HTTP_API_ROUTE_SHOT: return 96U;
    case HTTP_API_ROUTE_START:
    case HTTP_API_ROUTE_REMATCH:
    case HTTP_API_ROUTE_ABORT: return 80U;
    default: return 0U;
  }
}

static esp_err_t api_handler(httpd_req_t *request) {
  const http_api_route_t route = (http_api_route_t)(uintptr_t)request->user_ctx;
  char body[kRequestBodyCapacity + 1U] = {0};
  char token[kSessionTokenBytes * 2U + 1U] = {0};
  const size_t maximum = route_body_limit(route);
  size_t body_length = 0U;
  if (request->method == HTTP_POST) {
    if ((size_t)request->content_len > maximum) { body_length = maximum + 1U; record_rejected_input(); }
    else {
      while (body_length < (size_t)request->content_len) {
        const int received = httpd_req_recv(request, body + body_length, request->content_len - body_length);
        if (received <= 0) { record_rejected_input(); body_length = maximum + 1U; break; }
        body_length += (size_t)received;
      }
      body[body_length <= kRequestBodyCapacity ? body_length : 0U] = '\0';
    }
  }
  const bool target_too_large = request->method == HTTP_GET && httpd_req_get_url_query_len(request) > 128U;
  if (target_too_large) record_rejected_input();
  if (route == HTTP_API_ROUTE_STATE || route == HTTP_API_ROUTE_STATISTICS) {
    const size_t token_length = httpd_req_get_hdr_value_len(request, "X-Session-Token");
    if (token_length > 0U && token_length < sizeof(token) && httpd_req_get_hdr_value_str(request, "X-Session-Token", token, sizeof(token)) != ESP_OK) token[0] = '\0';
    else if (token_length >= sizeof(token)) snprintf(token, sizeof(token), "%s", "invalid");
  }
  refresh_health();
  http_api_response_t response;
  const http_api_request_t api_request = {.method = request->method == HTTP_GET ? HTTP_API_GET : HTTP_API_POST, .route = route,
      .content_type_json = request_has_json_content_type(request), .body = body, .body_length = body_length,
      .session_token = token[0] == '\0' ? NULL : token, .target_too_large = target_too_large};
  const uint32_t version_before = s_application.lifecycle.game.state.version;
  if (!http_api_handle(&s_api, &api_request, &response)) return ESP_FAIL;
  if (s_application.lifecycle.game.state.version != version_before) queue_state_broadcast();
  return send_api_response(request, &response);
}

static esp_err_t websocket_handler(httpd_req_t *request) {
  const int client_id = httpd_req_to_sockfd(request);
  if (request->method == HTTP_GET) {
    if (sync_service_open(&s_sync, client_id, (uint64_t)(esp_timer_get_time() / 1000U))) return ESP_OK;
    httpd_sess_trigger_close(s_server, client_id);
    return ESP_FAIL;
  }
  httpd_ws_frame_t frame = {0};
  if (httpd_ws_recv_frame(request, &frame, 0U) != ESP_OK || frame.type != HTTPD_WS_TYPE_TEXT ||
      frame.len > kWebSocketFrameCapacity) {
    record_rejected_input();
    sync_service_close(&s_sync, client_id);
    httpd_sess_trigger_close(s_server, client_id);
    return ESP_OK;
  }
  char input[kWebSocketFrameCapacity + 1U] = {0};
  frame.payload = (uint8_t *)input;
  if (httpd_ws_recv_frame(request, &frame, sizeof(input) - 1U) != ESP_OK) {
    sync_service_close(&s_sync, client_id);
    httpd_sess_trigger_close(s_server, client_id);
    return ESP_OK;
  }
  char output[kStateMessageCapacity] = {0};
  size_t output_length = 0U;
  bool state_changed = false;
  bool close_client = false;
  if (!sync_service_receive(&s_sync, client_id, input, frame.len, (uint64_t)(esp_timer_get_time() / 1000U),
                            output, &output_length, &state_changed, &close_client)) return ESP_FAIL;
  if (output_length > 0U) {
    httpd_ws_frame_t response = {.final = true, .type = HTTPD_WS_TYPE_TEXT,
        .payload = (uint8_t *)output, .len = output_length};
    if (httpd_ws_send_frame(request, &response) != ESP_OK) close_client = true;
  }
  if (state_changed) queue_state_broadcast();
  if (close_client) {
    sync_service_close(&s_sync, client_id);
    httpd_sess_trigger_close(s_server, client_id);
  }
  return ESP_OK;
}

#if WIFI_CONFIG_AVAILABLE
static void reconnect_timer_callback(void *unused) {
  (void)unused;
  portENTER_CRITICAL(&s_wifi_lock); s_wifi_state.retry_scheduled = false; portEXIT_CRITICAL(&s_wifi_lock);
  if (esp_wifi_connect() != ESP_OK) ESP_LOGW(kLogTag, "wifi connect request failed");
}

static void schedule_reconnect(uint32_t delay_ms) {
  bool schedule = false;
  portENTER_CRITICAL(&s_wifi_lock);
  if (s_wifi_state.configured && !s_wifi_state.connected && !s_wifi_state.retry_scheduled) { s_wifi_state.retry_scheduled = true; schedule = true; }
  portEXIT_CRITICAL(&s_wifi_lock);
  if (schedule && esp_timer_start_once(s_reconnect_timer, (uint64_t)delay_ms * 1000U) != ESP_OK) {
    portENTER_CRITICAL(&s_wifi_lock); s_wifi_state.retry_scheduled = false; portEXIT_CRITICAL(&s_wifi_lock);
  }
}

static void wifi_event_handler(void *argument, esp_event_base_t event_base, int32_t event_id, void *event_data) {
  (void)argument; (void)event_data;
  if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START) schedule_reconnect(0U);
  else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
    uint8_t attempt;
    portENTER_CRITICAL(&s_wifi_lock);
    if (s_wifi_state.reconnect_attempt < 5U) ++s_wifi_state.reconnect_attempt;
    attempt = s_wifi_state.reconnect_attempt; s_wifi_state.connected = false;
    portEXIT_CRITICAL(&s_wifi_lock);
    schedule_reconnect(1000U << attempt);
  } else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
    portENTER_CRITICAL(&s_wifi_lock); s_wifi_state.connected = true; s_wifi_state.reconnect_attempt = 0U; portEXIT_CRITICAL(&s_wifi_lock);
  }
}
#endif

static esp_err_t start_wifi(void) {
#if !WIFI_CONFIG_AVAILABLE
  ESP_LOGE(kLogTag, "wifi is not configured; copy include/wifi_config.h.example to include/wifi_config.h");
  return ESP_ERR_INVALID_STATE;
#else
  wifi_init_config_t init_config = WIFI_INIT_CONFIG_DEFAULT();
  wifi_config_t station_config = {0};
  const esp_timer_create_args_t timer_args = {.callback = reconnect_timer_callback, .name = "wifi_reconnect"};
  portENTER_CRITICAL(&s_wifi_lock); s_wifi_state.configured = true; portEXIT_CRITICAL(&s_wifi_lock);
  esp_netif_create_default_wifi_sta();
  ESP_RETURN_ON_ERROR(esp_wifi_init(&init_config), kLogTag, "wifi initialization failed");
  ESP_RETURN_ON_ERROR(esp_event_handler_instance_register(WIFI_EVENT, ESP_EVENT_ANY_ID, wifi_event_handler, NULL, NULL), kLogTag, "wifi event registration failed");
  ESP_RETURN_ON_ERROR(esp_event_handler_instance_register(IP_EVENT, IP_EVENT_STA_GOT_IP, wifi_event_handler, NULL, NULL), kLogTag, "ip event registration failed");
  ESP_RETURN_ON_ERROR(esp_timer_create(&timer_args, &s_reconnect_timer), kLogTag, "wifi timer creation failed");
  snprintf((char *)station_config.sta.ssid, sizeof(station_config.sta.ssid), "%s", WIFI_CONFIG_SSID);
  snprintf((char *)station_config.sta.password, sizeof(station_config.sta.password), "%s", WIFI_CONFIG_PASSWORD);
  station_config.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
  station_config.sta.pmf_cfg.capable = true;
  ESP_RETURN_ON_ERROR(esp_wifi_set_mode(WIFI_MODE_STA), kLogTag, "wifi mode setup failed");
  ESP_RETURN_ON_ERROR(esp_wifi_set_config(WIFI_IF_STA, &station_config), kLogTag, "wifi configuration failed");
  return esp_wifi_start();
#endif
}

static const char *mime_type_for_path(const char *path) {
  if (strcmp(path, "/index.html") == 0) return "text/html; charset=utf-8";
  if (strcmp(path, "/ship-sprite.svg") == 0) return "image/svg+xml";
  return strcmp(path, "/styles.css") == 0 ? "text/css; charset=utf-8" : "application/javascript; charset=utf-8";
}

static bool request_accepts_gzip(httpd_req_t *request) {
  const size_t length = httpd_req_get_hdr_value_len(request, "Accept-Encoding");
  if (length == 0U || length >= 96U) return false;
  char value[96];
  return httpd_req_get_hdr_value_str(request, "Accept-Encoding", value, sizeof(value)) == ESP_OK &&
      strstr(value, "gzip") != NULL;
}

static esp_err_t send_static_file(httpd_req_t *request, const char *asset_path) {
  char path[80];
  char gzip_path[84];
  snprintf(path, sizeof(path), "%s%s", kLittlefsBasePath, asset_path);
  const char *path_to_open = path;
  bool gzip = false;
  if (request_accepts_gzip(request)) {
    snprintf(gzip_path, sizeof(gzip_path), "%s.gz", path);
    FILE *compressed = fopen(gzip_path, "rb");
    if (compressed != NULL) { fclose(compressed); path_to_open = gzip_path; gzip = true; }
  }
  FILE *file = fopen(path_to_open, "rb");
  if (file == NULL) return httpd_resp_send_err(request, HTTPD_404_NOT_FOUND, "static asset not found");
  httpd_resp_set_type(request, mime_type_for_path(asset_path));
  httpd_resp_set_hdr(request, "Cache-Control", strcmp(asset_path, "/index.html") == 0 ? "no-cache" : "public, max-age=86400");
  if (gzip) { httpd_resp_set_hdr(request, "Content-Encoding", "gzip"); httpd_resp_set_hdr(request, "Vary", "Accept-Encoding"); }
  char chunk[kFileChunkBytes]; size_t read = 0U; esp_err_t result = ESP_OK;
  while ((read = fread(chunk, 1U, sizeof(chunk), file)) > 0U && result == ESP_OK) result = httpd_resp_send_chunk(request, chunk, read);
  fclose(file);
  return result == ESP_OK ? httpd_resp_send_chunk(request, NULL, 0U) : result;
}

static esp_err_t root_handler(httpd_req_t *request) {
  if (s_littlefs_mounted) return send_static_file(request, "/index.html");
  httpd_resp_set_status(request, "503 Service Unavailable");
  return httpd_resp_send(request, "LittleFS is unavailable", HTTPD_RESP_USE_STRLEN);
}

static esp_err_t static_file_handler(httpd_req_t *request) {
  if (!s_littlefs_mounted) {
    httpd_resp_set_status(request, "503 Service Unavailable");
    return httpd_resp_send(request, "LittleFS is unavailable", HTTPD_RESP_USE_STRLEN);
  }
  if (strcmp(request->uri, "/styles.css") != 0 && strcmp(request->uri, "/app.js") != 0 &&
      strcmp(request->uri, "/target_interaction.js") != 0 && strcmp(request->uri, "/ship-sprite.svg") != 0) return httpd_resp_send_err(request, HTTPD_404_NOT_FOUND, "asset not found");
  return send_static_file(request, request->uri);
}

static esp_err_t start_http_server(void) {
  httpd_config_t config = HTTPD_DEFAULT_CONFIG();
  config.max_uri_handlers = 16U; config.max_open_sockets = kHttpMaxOpenSockets; config.uri_match_fn = httpd_uri_match_wildcard; config.lru_purge_enable = true;
  ESP_RETURN_ON_ERROR(httpd_start(&s_server, &config), kLogTag, "http server start failed");
  const httpd_uri_t routes[] = {
      {.uri = "/", .method = HTTP_GET, .handler = root_handler},
      {.uri = "/api/info", .method = HTTP_GET, .handler = api_handler, .user_ctx = (void *)(uintptr_t)HTTP_API_ROUTE_INFO},
      {.uri = "/api/health", .method = HTTP_GET, .handler = api_handler, .user_ctx = (void *)(uintptr_t)HTTP_API_ROUTE_HEALTH},
      {.uri = "/api/session/join", .method = HTTP_POST, .handler = api_handler, .user_ctx = (void *)(uintptr_t)HTTP_API_ROUTE_JOIN},
      {.uri = "/api/session/resume", .method = HTTP_POST, .handler = api_handler, .user_ctx = (void *)(uintptr_t)HTTP_API_ROUTE_RESUME},
      {.uri = "/api/game/config", .method = HTTP_POST, .handler = api_handler, .user_ctx = (void *)(uintptr_t)HTTP_API_ROUTE_CONFIG},
      {.uri = "/api/game/start", .method = HTTP_POST, .handler = api_handler, .user_ctx = (void *)(uintptr_t)HTTP_API_ROUTE_START},
      {.uri = "/api/game/shot", .method = HTTP_POST, .handler = api_handler, .user_ctx = (void *)(uintptr_t)HTTP_API_ROUTE_SHOT},
      {.uri = "/api/game/rematch", .method = HTTP_POST, .handler = api_handler, .user_ctx = (void *)(uintptr_t)HTTP_API_ROUTE_REMATCH},
      {.uri = "/api/game/abort", .method = HTTP_POST, .handler = api_handler, .user_ctx = (void *)(uintptr_t)HTTP_API_ROUTE_ABORT},
      {.uri = "/api/state", .method = HTTP_GET, .handler = api_handler, .user_ctx = (void *)(uintptr_t)HTTP_API_ROUTE_STATE},
      {.uri = "/api/statistics", .method = HTTP_GET, .handler = api_handler, .user_ctx = (void *)(uintptr_t)HTTP_API_ROUTE_STATISTICS},
      {.uri = "/ws", .method = HTTP_GET, .handler = websocket_handler, .is_websocket = true},
      {.uri = "/*", .method = HTTP_GET, .handler = static_file_handler},
  };
  for (size_t index = 0; index < sizeof(routes) / sizeof(routes[0]); ++index) ESP_RETURN_ON_ERROR(httpd_register_uri_handler(s_server, &routes[index]), kLogTag, "route registration failed");
  return ESP_OK;
}

static void mount_littlefs(void) {
  const esp_vfs_littlefs_conf_t config = {.base_path = kLittlefsBasePath, .partition_label = kLittlefsPartitionLabel, .format_if_mount_failed = false, .dont_mount = false, .grow_on_mount = false};
  if (esp_vfs_littlefs_register(&config) != ESP_OK) { ESP_LOGE(kLogTag, "littlefs mount failed"); return; }
  s_littlefs_mounted = true;
}

void app_main(void) {
  esp_err_t result = nvs_flash_init();
  if (result == ESP_ERR_NVS_NO_FREE_PAGES || result == ESP_ERR_NVS_NEW_VERSION_FOUND) { ESP_ERROR_CHECK(nvs_flash_erase()); result = nvs_flash_init(); }
  ESP_ERROR_CHECK(result); ESP_ERROR_CHECK(esp_netif_init()); ESP_ERROR_CHECK(esp_event_loop_create_default());
  application_init(&s_application, (random_source_t){.next_u32 = platform_random, .context = NULL});
  const esp_timer_create_args_t bot_timer_args = {.callback = bot_timer_callback, .name = "bot_turn"};
  ESP_ERROR_CHECK(esp_timer_create(&bot_timer_args, &s_bot_timer));
  application_set_bot_scheduler(&s_application, (scheduler_t){.schedule_after_ms = schedule_bot_turn, .context = NULL});
  http_api_init(&s_api, &s_application);
  sync_service_init(&s_sync, &s_application);
  mount_littlefs();
  ESP_ERROR_CHECK(start_http_server());
  const esp_timer_create_args_t sync_timer_args = {.callback = sync_timer_callback, .name = "sync_expire"};
  ESP_ERROR_CHECK(esp_timer_create(&sync_timer_args, &s_sync_timer));
  ESP_ERROR_CHECK(esp_timer_start_periodic(s_sync_timer, 1000000U));
  result = start_wifi();
  if (result != ESP_OK) ESP_LOGW(kLogTag, "wifi unavailable: %s", esp_err_to_name(result));
}