picomap/firmware/include/wire.h

#pragma once
#include <array>
#include <cstdint>
#include <span>
#include <string>
#include <tuple>
#include <vector>
#include "msgpack.h"
#include "halfsiphash.h"
#include "static_vector.h"

struct Envelope {
    static constexpr int8_t ext_id = 0;
    uint32_t message_id;
    uint32_t checksum;
    std::span<const uint8_t> payload;
    auto as_tuple() const { return std::tie(message_id, checksum, payload); }
    auto as_tuple() { return std::tie(message_id, checksum, payload); }
};

struct DeviceError {
    static constexpr int8_t ext_id = 1;
    uint32_t code;
    std::string message;
    auto as_tuple() const { return std::tie(code, message); }
    auto as_tuple() { return std::tie(code, message); }
};

struct RequestPICOBOOT {
    static constexpr int8_t ext_id = 2;
    auto as_tuple() const { return std::tie(); }
    auto as_tuple() { return std::tie(); }
};

struct ResponsePICOBOOT {
    static constexpr int8_t ext_id = 3;
    auto as_tuple() const { return std::tie(); }
    auto as_tuple() { return std::tie(); }
};

struct RequestInfo {
    static constexpr int8_t ext_id = 4;
    auto as_tuple() const { return std::tie(); }
    auto as_tuple() { return std::tie(); }
};

struct ResponseInfo {
    static constexpr int8_t ext_id = 5;
    std::array<uint8_t, 8> board_id;
    std::array<uint8_t, 6> mac;
    std::array<uint8_t, 4> ip;
    std::string firmware_name;
    auto as_tuple() const { return std::tie(board_id, mac, ip, firmware_name); }
    auto as_tuple() { return std::tie(board_id, mac, ip, firmware_name); }
};

struct RequestLog {
    static constexpr int8_t ext_id = 6;
    auto as_tuple() const { return std::tie(); }
    auto as_tuple() { return std::tie(); }
};

struct LogEntry {
    uint32_t timestamp_us;
    std::string message;
    auto as_tuple() const { return std::tie(timestamp_us, message); }
    auto as_tuple() { return std::tie(timestamp_us, message); }
};

struct ResponseLog {
    static constexpr int8_t ext_id = 7;
    std::vector<LogEntry> entries;
    auto as_tuple() const { return std::tie(entries); }
    auto as_tuple() { return std::tie(entries); }
};

struct RequestTest {
    static constexpr int8_t ext_id = 127;
    std::string name;
    auto as_tuple() const { return std::tie(name); }
    auto as_tuple() { return std::tie(name); }
};

struct ResponseTest {
    static constexpr int8_t ext_id = 126;
    bool pass;
    std::vector<std::string> messages;
    auto as_tuple() const { return std::tie(pass, messages); }
    auto as_tuple() { return std::tie(pass, messages); }
};

static constexpr uint8_t hash_key[8] = {};

struct DecodedMessage {
    uint32_t message_id;
    int8_t type_id;
    std::span<const uint8_t> payload;
};

static constexpr size_t ext16_header_len = 4;
static constexpr size_t array3_header_len = 1;
static constexpr size_t uint32_fixed_len = 5;
static constexpr size_t bin16_header_len = 3;

static constexpr size_t envelope_hdr_len =
    ext16_header_len + array3_header_len +
    uint32_fixed_len + uint32_fixed_len + bin16_header_len;

static constexpr size_t response_prefix_len = envelope_hdr_len + ext16_header_len;

template <typename T>
inline msgpack::result<size_t> encode_response_into(span_writer &out, uint32_t message_id, const T &msg) {
    auto body = out.subspan(response_prefix_len);
    msgpack::packer body_p(body);
    body_p.pack(msg.as_tuple());

    auto inner_ext = out.subspan(envelope_hdr_len, ext16_header_len);
    msgpack::packer(inner_ext).pack_ext16_header(T::ext_id, static_cast<uint16_t>(body.size()));

    size_t bin_len = inner_ext.size() + body.size();
    uint32_t checksum = halfsiphash::hash32({inner_ext.data(), bin_len}, hash_key);

    auto env_hdr = out.subspan(0, envelope_hdr_len);
    size_t env_body_len = array3_header_len + uint32_fixed_len + uint32_fixed_len + bin16_header_len + bin_len;
    msgpack::packer hdr(env_hdr);
    hdr.pack_ext16_header(Envelope::ext_id, static_cast<uint16_t>(env_body_len));
    hdr.pack_array(3);
    hdr.pack_uint32_fixed(message_id);
    hdr.pack_uint32_fixed(checksum);
    hdr.pack_bin16_header(static_cast<uint16_t>(bin_len));

    if (body.overflow() || inner_ext.overflow() || env_hdr.overflow())
        return std::unexpected(msgpack::error_code::overflow);
    return response_prefix_len + body.size();
}

inline msgpack::result<DecodedMessage> try_decode(const uint8_t *data, size_t len) {
    msgpack::parser p(data, static_cast<int>(len));

    Envelope env;
    auto r = msgpack::unpack(p, env);
    if (!r) return std::unexpected(r.error());

    uint32_t expected = halfsiphash::hash32(env.payload, hash_key);
    if (env.checksum != expected) return std::unexpected(msgpack::error_code::invalid);

    msgpack::parser inner(env.payload.data(), static_cast<int>(env.payload.size()));
    if (!inner.is_ext()) return std::unexpected(msgpack::error_code::type_error);
    auto ext = inner.get_ext();
    if (!ext) return std::unexpected(ext.error());

    auto& [type_id, ext_data] = *ext;
    return DecodedMessage{env.message_id, type_id,
        std::span<const uint8_t>(reinterpret_cast<const uint8_t*>(ext_data.data()), ext_data.size())};
}

template <size_t N>
inline msgpack::result<DecodedMessage> try_decode(const static_vector<uint8_t, N> &buf) {
    return try_decode(buf.data(), buf.size());
}

template <typename T>
inline msgpack::result<T> decode_response(const uint8_t *data, size_t len) {
    msgpack::parser p(data, static_cast<int>(len));

    Envelope env;
    auto r = msgpack::unpack(p, env);
    if (!r) return std::unexpected(r.error());

    uint32_t expected = halfsiphash::hash32(env.payload, hash_key);
    if (env.checksum != expected) return std::unexpected(msgpack::error_code::invalid);

    msgpack::parser inner(env.payload.data(), static_cast<int>(env.payload.size()));
    T out;
    auto r2 = msgpack::unpack(inner, out);
    if (!r2) return std::unexpected(r2.error());
    return out;
}