roro9stack/src/services/debug_console.cpp

#ifdef RORO_DEBUG

#include "services/debug_console.h"

#include <WiFi.h>

#include <algorithm>
#include <esp_core_dump.h>
#include <esp_flash.h>

#include <SD.h>
#include <unistd.h>

#include <atomic>
#include <memory>

#include "file_receiver.h"
#include "platform/console.h"
#include "version.h"

namespace roro {

namespace {

constexpr uint32_t kAuthTimeoutMs = 10000;
constexpr size_t kMaxLine = 240;
constexpr size_t kMaxQueued = 8;

// Compares without stopping at the first difference, so timing says nothing about the token.
bool sameToken(const std::string& a, const char* b) {
    size_t n = strlen(b);
    uint8_t diff = a.size() != n;
    for (size_t i = 0; i < n; i++) diff |= (i < a.size() ? a[i] : 0) ^ b[i];
    return diff == 0;
}

// Reads one line (without its \r\n) within `timeoutMs`; false on a timeout, a drop or an overlong line.
bool readLine(NetworkClient& c, std::string& line, uint32_t timeoutMs) {
    line.clear();
    uint32_t since = millis();
    while (c.connected() && millis() - since < timeoutMs) {
        int ch = c.read();
        if (ch < 0) {
            delay(10);
            continue;
        }
        if (ch == '\n') {
            if (!line.empty() && line.back() == '\r') line.pop_back();
            return true;
        }
        if (line.size() >= kMaxLine) return false;
        line += static_cast<char>(ch);
    }
    return false;
}

// The raw core dump partition contents, as esp-coredump reads them ("-t raw").
void sendCoreDump(NetworkClient& client) {
    size_t addr = 0, size = 0;
    if (esp_core_dump_image_check() != ESP_OK || esp_core_dump_image_get(&addr, &size) != ESP_OK) {
        client.print("coredump: none\n");
        return;
    }
    client.printf("coredump: data %u\n", (unsigned)size);
    uint8_t buf[1024];
    for (size_t off = 0; off < size;) {
        size_t n = std::min(sizeof buf, size - off);
        if (esp_flash_read(nullptr, buf, addr + off, n) != ESP_OK) break;  // the host sees a short file
        if (client.write(buf, n) != n) return;
        off += n;
    }
    client.print("coredump: end\n");
}

}  // namespace

DebugConsole::DebugConsole(WifiService& wifi, StorageService& storage)
    : wifi_(wifi), storage_(storage), lock_(xSemaphoreCreateMutex()) {}

void DebugConsole::start() {
    console.captureEspLogs();
    if (!task_) xTaskCreate(taskEntry, "debug", 6144, this, 1, &task_);
}

bool DebugConsole::takeCommand(std::string& line) {
    xSemaphoreTake(lock_, portMAX_DELAY);
    bool any = !commands_.empty();
    if (any) {
        line = std::move(commands_.front());
        commands_.pop_front();
    }
    xSemaphoreGive(lock_);
    return any;
}

void DebugConsole::taskEntry(void* self) { static_cast<DebugConsole*>(self)->listen(); }

void DebugConsole::listen() {
    NetworkServer server(kPort);
    bool listening = false;
    for (;;) {
        bool up = wifi_.state() == WifiController::State::Connected;
        if (up && !listening) {
            server.begin();
            listening = true;
        } else if (!up && listening) {
            server.end();
            listening = false;
        }
        if (listening) {
            NetworkClient client = server.accept();
            if (client) {
                client.setNoDelay(true);
                if (authenticate(client)) serve(client);
                client.stop();
            }
        }
        vTaskDelay(pdMS_TO_TICKS(200));
    }
}

bool DebugConsole::authenticate(NetworkClient& client) {
    std::string token;
    if (readLine(client, token, kAuthTimeoutMs) && sameToken(token, RORO_DEBUG_TOKEN)) return true;
    console.printf("debug: refused a client from %s\n", client.remoteIP().toString().c_str());
    delay(1000);  // no quick retries
    client.print("denied\n");
    return false;
}

void DebugConsole::serve(NetworkClient& client) {
    console.printf("debug: client %s connected\n", client.remoteIP().toString().c_str());
    client.printf("%s %s debug console. 'help' lists the commands. Backlog follows.\n", kProductName, versionString());
    connected_ = true;
    uint32_t pos = console.oldest();
    uint8_t buf[512];
    std::string line;
    while (client.connected()) {
        // Console output since last time, including the replies to this client's commands.
        uint32_t skipped = 0;
        size_t n;
        while ((n = console.readSince(pos, buf, sizeof buf, skipped)) > 0) {
            if (skipped) client.printf("\n[... %u bytes lost: the console ran faster than the network]\n", (unsigned)skipped);
            if (client.write(buf, n) != n) break;
        }
        // Commands, one line at a time.
        while (client.available()) {
            int ch = client.read();
            if (ch == '\r') continue;
            if (ch != '\n') {
                if (line.size() < kMaxLine) line += static_cast<char>(ch);
                continue;
            }
            if (line == "quit" || line == "exit") {
                client.stop();
                break;
            }
            if (binaryCommand(client, line)) {
                line.clear();
                continue;
            }
            xSemaphoreTake(lock_, portMAX_DELAY);
            bool full = commands_.size() >= kMaxQueued;
            if (!full && !line.empty()) commands_.push_back(line);
            xSemaphoreGive(lock_);
            if (full) client.print("debug: busy, command dropped\n");
            line.clear();
        }
        vTaskDelay(pdMS_TO_TICKS(20));
    }
    connected_ = false;
    console.println("debug: client disconnected");
}

bool DebugConsole::binaryCommand(NetworkClient& client, const std::string& line) {
    if (line == "reset") {  // works even when the main loop is stuck
        client.print("debug: restarting now\n");
        client.flush();
        delay(200);
        esp_restart();
    }
    if (line == "coredump get") sendCoreDump(client);
    else if (line.rfind("get ", 0) == 0) get(client, line.substr(4));
    else if (line.rfind("put ", 0) == 0) put(client, line.substr(4));
    else if (line == "screenshot") screenshot(client);
    else return false;
    return true;
}

bool DebugConsole::onStorage(std::function<void()> job) {
    // Shared with the job, which outlives this wait if the card is missing and it never starts.
    // Exactly one side wins the state change: the job (Queued -> Running) or the wait giving up
    // (Queued -> Abandoned), so an abandoned job can never touch this stack frame.
    enum : int { Queued, Running, Abandoned };
    struct Run {
        std::atomic<int> state{Queued};
        SemaphoreHandle_t done = xSemaphoreCreateBinary();
        ~Run() { vSemaphoreDelete(done); }
    };
    auto run = std::make_shared<Run>();
    storage_.runJob([run, job]() {
        int expected = Queued;
        if (!run->state.compare_exchange_strong(expected, Running)) return;
        job();
        xSemaphoreGive(run->done);
    });
    for (int waited = 0; xSemaphoreTake(run->done, pdMS_TO_TICKS(500)) != pdTRUE; waited += 500) {
        int expected = Queued;
        if (waited >= 5000 && run->state.compare_exchange_strong(expected, Abandoned)) return false;
    }
    return true;
}

void DebugConsole::get(NetworkClient& client, const std::string& path) {
    bool ran = onStorage([&]() {
        File f = SD.open(path.c_str());
        if (!f || f.isDirectory()) return (void)client.printf("get: error cannot open %s\n", path.c_str());
        size_t size = f.size();
        client.printf("get: data %u\n", (unsigned)size);
        uint8_t buf[1024];
        for (size_t sent = 0; sent < size;) {
            int n = f.read(buf, std::min(sizeof buf, size - sent));
            if (n <= 0 || client.write(buf, n) != static_cast<size_t>(n)) break;  // the host sees it short
            sent += n;
        }
        f.close();
        client.print("get: end\n");
    });
    if (!ran) client.print("get: error no SD card\n");
}

// The same checks as `sd put` over serial (FileReceiver), but TCP does the flow control, so the
// file is written in one job with the file kept open.
void DebugConsole::put(NetworkClient& client, const std::string& args) {
    FileReceiver r;
    std::string error = r.begin(args, millis());
    StorageState card = storage_.state();
    if (error.empty() && !card.present) error = "no SD card";
    if (error.empty() && card.totalBytes - card.usedBytes < r.size() + 64 * 1024) error = "not enough space";
    if (!error.empty()) return (void)client.printf("put: error %s\n", error.c_str());

    bool ran = onStorage([&]() {
        std::string part = r.partPath();
        for (size_t slash = part.find('/', 1); slash != std::string::npos; slash = part.find('/', slash + 1)) {
            std::string dir = part.substr(0, slash);
            if (!SD.exists(dir.c_str())) SD.mkdir(dir.c_str());
        }
        File f = SD.open(part.c_str(), FILE_WRITE);
        if (!f) return (void)client.print("put: error card not writable\n");
        client.printf("put: ready %u\n", (unsigned)r.size());
        uint8_t buf[1024];
        using S = FileReceiver::State;
        while (r.state() == S::Receiving || r.state() == S::Writing) {
            if (r.state() == S::Writing) {
                const auto& c = r.chunk();
                bool ok = f.write(c.data(), c.size()) == c.size();
                // A card can fail one write and take the next. FATFS keeps a failed file in error,
                // so: close, cut back to the last good byte, reopen, try again.
                for (int retry = 1; !ok && retry <= 3; retry++) {
                    console.printf("put: write failed at %u, retry %d\n", (unsigned)r.received(), retry);
                    f.close();
                    delay(50 * retry);
                    truncate(("/sd" + part).c_str(), r.received());
                    f = SD.open(part.c_str(), FILE_APPEND);
                    ok = f && f.size() == r.received() && f.write(c.data(), c.size()) == c.size();
                }
                r.chunkWritten(ok, millis());
                continue;
            }
            int n = client.read(buf, std::min(sizeof buf, r.wanted()));
            if (n > 0) r.feed(buf, n, millis());
            else if (!client.connected()) break;
            else {
                delay(2);
                r.tick(millis());
            }
        }
        f.close();
        if (r.state() == S::Finishing) {
            if (SD.exists(r.path().c_str())) SD.remove(r.path().c_str());
            r.finished(SD.rename(part.c_str(), r.path().c_str()));
        }
        if (r.state() == S::Done) client.printf("put: done %s %u B\n", r.path().c_str(), (unsigned)r.size());
        else {
            SD.remove(part.c_str());
            client.printf("put: error %s\n", r.error().empty() ? "connection lost" : r.error().c_str());
            // The rest of the file is still on its way: never read it as commands.
            client.flush();
            client.stop();
        }
    });
    if (!ran) client.print("put: error no SD card\n");
}

// The frame as composed off-screen (RGB332, one byte a pixel). Read while the UI may be drawing,
// so it can tear; it's for looking, not for pixel-exact tests.
void DebugConsole::screenshot(NetworkClient& client) {
    const uint8_t* pixels = frame_ ? static_cast<const uint8_t*>(frame_->getBuffer()) : nullptr;
    if (!pixels) return (void)client.print("screenshot: error no frame\n");
    int w = frame_->width(), h = frame_->height();
    client.printf("screenshot: rgb332 %d %d\n", w, h);
    for (int y = 0; y < h; y += 16) client.write(pixels + y * w, w * std::min(16, h - y));
    client.print("screenshot: end\n");
}

}  // namespace roro

#endif