compact.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright 2026 Johnson Ogundeji
#pragma once
 
 
#include <cstdint>
#include <cstring>
#include <limits>
#include <optional>
#include <stack>
#include <stdexcept>
#include <string>
#include <vector>
 
#include "signet/error.hpp"
 
namespace signet::forge::thrift {
 
namespace compact_type {
    inline constexpr uint8_t STOP       = 0;   
    inline constexpr uint8_t BOOL_TRUE  = 1;   
    inline constexpr uint8_t BOOL_FALSE = 2;   
    inline constexpr uint8_t I8         = 3;   
    inline constexpr uint8_t I16        = 4;   
    inline constexpr uint8_t I32        = 5;   
    inline constexpr uint8_t I64        = 6;   
    inline constexpr uint8_t DOUBLE     = 7;   
    inline constexpr uint8_t BINARY     = 8;   
    inline constexpr uint8_t LIST       = 9;   
    inline constexpr uint8_t SET        = 10;  
    inline constexpr uint8_t MAP        = 11;  
    inline constexpr uint8_t STRUCT     = 12;  
} // namespace compact_type
 
struct FieldHeader {
    int16_t  field_id;     
    uint8_t  thrift_type;  
 
    [[nodiscard]] bool is_stop() const { return field_id == 0 && thrift_type == 0; }
};
 
struct ListHeader {
    uint8_t  elem_type;  
    int32_t  size;       
};
 
class CompactEncoder {
public:
    static constexpr size_t MAX_STRING_BYTES = 64u * 1024u * 1024u;  
 
    CompactEncoder() { last_field_ids_.push(0); }
 
    // -- field / struct helpers ------------------------------------------------
 
    void write_field(int16_t field_id, uint8_t thrift_type) {
        int16_t delta = field_id - last_field_ids_.top();
        if (delta > 0 && delta <= 15) {
            buf_.push_back(static_cast<uint8_t>((delta << 4) | thrift_type));
        } else {
            buf_.push_back(thrift_type);
            write_zigzag_i16(field_id);
        }
        last_field_ids_.top() = field_id;
    }
 
    void write_stop() { buf_.push_back(0x00); }
 
    void begin_struct() { last_field_ids_.push(0); }
 
    void end_struct() { last_field_ids_.pop(); }
 
    // -- primitive writers -----------------------------------------------------
 
    void write_bool(bool val) {
        buf_.push_back(val ? 0x01 : 0x00);
    }
 
    void write_i8(int8_t val) {
        buf_.push_back(static_cast<uint8_t>(val));
    }
 
    void write_field_bool(int16_t field_id, bool val) {
        uint8_t thrift_type = val ? compact_type::BOOL_TRUE
                                  : compact_type::BOOL_FALSE;
        int16_t delta = field_id - last_field_ids_.top();
        if (delta > 0 && delta <= 15) {
            buf_.push_back(static_cast<uint8_t>((delta << 4) | thrift_type));
        } else {
            buf_.push_back(thrift_type);
            write_zigzag_i16(field_id);
        }
        last_field_ids_.top() = field_id;
    }
 
    void write_i32(int32_t val) {
        write_varint32(zigzag_encode_i32(val));
    }
 
    void write_i64(int64_t val) {
        write_varint64(zigzag_encode_i64(val));
    }
 
    void write_double(double val) {
        uint64_t bits;
        std::memcpy(&bits, &val, 8);
        for (int i = 0; i < 8; ++i)
            buf_.push_back(static_cast<uint8_t>((bits >> (i * 8)) & 0xFF));
    }
 
    void write_float(float val) {
        uint32_t bits;
        std::memcpy(&bits, &val, 4);
        for (int i = 0; i < 4; ++i)
            buf_.push_back(static_cast<uint8_t>((bits >> (i * 8)) & 0xFF));
    }
 
    void write_string(const std::string& val) {
        if (val.size() > MAX_STRING_BYTES) {
            throw std::overflow_error("CompactEncoder::write_string: length "
                + std::to_string(val.size()) + " exceeds MAX_STRING_BYTES");
        }
        write_varint32(static_cast<uint32_t>(val.size()));
        buf_.insert(buf_.end(), val.begin(), val.end());
    }
 
    void write_binary(const uint8_t* data, size_t len) {
        if (len > MAX_STRING_BYTES) {
            throw std::overflow_error("CompactEncoder::write_binary: length "
                + std::to_string(len) + " exceeds MAX_STRING_BYTES");
        }
        write_varint32(static_cast<uint32_t>(len));
        buf_.insert(buf_.end(), data, data + len);
    }
 
    void write_list_header(uint8_t elem_type, int32_t size) {
        if (size < 0) {
            throw std::invalid_argument("write_list_header: negative list size");
        }
        if (size <= 14) {
            buf_.push_back(static_cast<uint8_t>((size << 4) | elem_type));
        } else {
            buf_.push_back(static_cast<uint8_t>(0xF0 | elem_type));
            write_varint32(static_cast<uint32_t>(size));
        }
    }
 
    // -- access ---------------------------------------------------------------
 
    [[nodiscard]] const std::vector<uint8_t>& data() const { return buf_; }
 
    [[nodiscard]] size_t size() const { return buf_.size(); }
 
    void clear() {
        buf_.clear();
        while (!last_field_ids_.empty()) last_field_ids_.pop();
        last_field_ids_.push(0);
    }
 
private:
    std::vector<uint8_t> buf_;
    std::stack<int16_t>  last_field_ids_;
 
    // -- varint encoding ------------------------------------------------------
 
    void write_varint32(uint32_t val) {
        while (val > 0x7F) {
            buf_.push_back(static_cast<uint8_t>((val & 0x7F) | 0x80));
            val >>= 7;
        }
        buf_.push_back(static_cast<uint8_t>(val));
    }
 
    void write_varint64(uint64_t val) {
        while (val > 0x7F) {
            buf_.push_back(static_cast<uint8_t>((val & 0x7F) | 0x80));
            val >>= 7;
        }
        buf_.push_back(static_cast<uint8_t>(val));
    }
 
    void write_zigzag_i16(int16_t val) {
        uint32_t zz = zigzag_encode_i32(static_cast<int32_t>(val));
        write_varint32(zz);
    }
 
    // -- zigzag encoding ------------------------------------------------------
 
    // CWE-190, C++ [expr.shift] p7.6.7 — left shift on unsigned to avoid UB on signed overflow
    static uint32_t zigzag_encode_i32(int32_t val) {
        return (static_cast<uint32_t>(val) << 1) ^ static_cast<uint32_t>(val >> 31);
    }
 
    // CWE-190, C++ [expr.shift] p7.6.7 — left shift on unsigned to avoid UB
    static uint64_t zigzag_encode_i64(int64_t val) {
        return (static_cast<uint64_t>(val) << 1) ^ static_cast<uint64_t>(val >> 63);
    }
};
 
class CompactDecoder {
public:
    CompactDecoder(const uint8_t* data, size_t size)
        : data_(data), size_(size), pos_(0), error_(false),
          pending_bool_{}, pending_bool_valid_(false) {
        last_field_ids_.push(0);
    }
 
    // -- field / struct helpers ------------------------------------------------
 
    [[nodiscard]] FieldHeader read_field_header() {
        if (!ensure(1)) return {0, 0};
 
        uint8_t byte = data_[pos_++];
 
        // STOP marker
        if (byte == 0x00) return {0, 0};
 
        uint8_t type = byte & 0x0F;
        int16_t delta = static_cast<int16_t>((byte >> 4) & 0x0F);
 
        int16_t field_id;
        if (delta != 0) {
            // Delta-encoded field ID
            field_id = last_field_ids_.top() + delta;
        } else {
            // Full field ID follows as zigzag varint
            int32_t id32 = read_zigzag_i32();
            field_id = static_cast<int16_t>(id32);
        }
        last_field_ids_.top() = field_id;
 
        // For bool fields, the value is embedded in the type nibble.
        // Cache it for the next read_bool() call.
        if (type == compact_type::BOOL_TRUE) {
            pending_bool_ = true;
            pending_bool_valid_ = true;
        } else if (type == compact_type::BOOL_FALSE) {
            pending_bool_ = false;
            pending_bool_valid_ = true;
        }
 
        if (++field_count_ > MAX_FIELD_COUNT) {
            error_ = true;
            return {0, 0};
        }
        // CWE-400: Uncontrolled Resource Consumption (DoS prevention)
        if (++total_fields_read_ > MAX_TOTAL_FIELDS) {
            error_ = true;
            return {0, 0};
        }
 
        return {field_id, type};
    }
 
    [[nodiscard]] bool read_bool() {
        if (pending_bool_valid_) {
            pending_bool_valid_ = false;
            return pending_bool_;
        }
        if (!ensure(1)) return false;
        return data_[pos_++] != 0;
    }
 
    [[nodiscard]] int8_t read_i8() {
        if (!ensure(1)) return 0;
        return static_cast<int8_t>(data_[pos_++]);
    }
 
    [[nodiscard]] int32_t read_i32() {
        return read_zigzag_i32();
    }
 
    [[nodiscard]] int64_t read_i64() {
        return read_zigzag_i64();
    }
 
    [[nodiscard]] double read_double() {
        if (!ensure(8)) return 0.0;
        uint64_t bits = 0;
        for (int i = 0; i < 8; ++i)
            bits |= static_cast<uint64_t>(data_[pos_++]) << (i * 8);
        double val;
        std::memcpy(&val, &bits, 8);
        return val;
    }
 
    [[nodiscard]] float read_float() {
        if (!ensure(4)) return 0.0f;
        uint32_t bits = 0;
        for (int i = 0; i < 4; ++i)
            bits |= static_cast<uint32_t>(data_[pos_++]) << (i * 8);
        float val;
        std::memcpy(&val, &bits, 4);
        return val;
    }
 
    [[nodiscard]] std::string read_string() {
        uint32_t len = read_varint32();
        if (len > MAX_STRING_BYTES) { error_ = true; return {}; }
        if (!ensure(len)) return {};
        std::string result(reinterpret_cast<const char*>(data_ + pos_), len);
        pos_ += len;
        return result;
    }
 
    [[nodiscard]] std::vector<uint8_t> read_binary() {
        uint32_t len = read_varint32();
        if (len > MAX_STRING_BYTES) { error_ = true; return {}; }
        if (!ensure(len)) return {};
        std::vector<uint8_t> result(data_ + pos_, data_ + pos_ + len);
        pos_ += len;
        return result;
    }
 
    [[nodiscard]] ListHeader read_list_header() {
        if (!ensure(1)) return {0, 0};
        uint8_t byte = data_[pos_++];
        uint8_t elem_type = byte & 0x0F;
        int32_t size = (byte >> 4) & 0x0F;
        if (size == 15) {
            // Large list: size follows as varint
            uint32_t raw = read_varint32();
            if (raw > static_cast<uint32_t>((std::numeric_limits<int32_t>::max)())) {
                error_ = true;
                return {0, 0};
            }
            size = static_cast<int32_t>(raw);
        }
        if (size < 0) {
            // CWE-20: Improper Input Validation — negative list size (likely corrupt data)
            error_ = true;
            return {0, 0};
        }
        if (static_cast<uint32_t>(size) > MAX_COLLECTION_SIZE) {
            error_ = true; return {0, 0};
        }
        return {elem_type, size};
    }
 
    void skip_field(uint8_t thrift_type) {
        switch (thrift_type) {
        case compact_type::BOOL_TRUE:
        case compact_type::BOOL_FALSE:
            // Value already consumed in field header; nothing to skip.
            break;
 
        case compact_type::I8:
            // Single byte
            if (ensure(1)) pos_ += 1;
            break;
 
        case compact_type::I16:
        case compact_type::I32:
            // Zigzag varint — just consume it
            (void)read_varint32();
            break;
 
        case compact_type::I64:
            // Zigzag varint64 — just consume it
            (void)read_varint64();
            break;
 
        case compact_type::DOUBLE:
            if (ensure(8)) pos_ += 8;
            break;
 
        case compact_type::BINARY: {
            // Length-prefixed bytes
            uint32_t len = read_varint32();
            if (ensure(len)) pos_ += len;
            break;
        }
 
        case compact_type::LIST:
        case compact_type::SET: {
            auto hdr = read_list_header();
            if (hdr.size < 0 || static_cast<uint32_t>(hdr.size) > MAX_COLLECTION_SIZE) {
                error_ = true; break;
            }
            for (int32_t i = 0; i < hdr.size && good(); ++i) {
                skip_field(hdr.elem_type);
            }
            break;
        }
 
        case compact_type::MAP: {
            uint32_t map_size = read_varint32();
            if (map_size == 0) break;
            if (map_size > MAX_COLLECTION_SIZE) { error_ = true; break; }
            if (!ensure(1)) break;
            uint8_t kv_types = data_[pos_++];
            uint8_t key_type = (kv_types >> 4) & 0x0F;
            uint8_t val_type = kv_types & 0x0F;
            for (uint32_t i = 0; i < map_size && good(); ++i) {
                skip_field(key_type);
                skip_field(val_type);
            }
            break;
        }
 
        case compact_type::STRUCT: {
            // Read fields until STOP
            begin_struct();
            while (good()) {
                auto fh = read_field_header();
                if (fh.is_stop()) break;
                skip_field(fh.thrift_type);
            }
            end_struct();
            break;
        }
 
        default:
            // Unknown type — mark as error
            error_ = true;
            break;
        }
    }
 
    void begin_struct() {
        if (last_field_ids_.size() >= MAX_NESTING_DEPTH) { error_ = true; return; }
        last_field_ids_.push(0);
        field_count_ = 0;
    }
 
    void end_struct() {
        if (last_field_ids_.empty()) { error_ = true; return; } // CWE-124: Buffer Underwrite
        last_field_ids_.pop();
    }
 
    // -- state queries --------------------------------------------------------
 
    [[nodiscard]] size_t remaining() const {
        return (pos_ <= size_) ? (size_ - pos_) : 0;
    }
 
    [[nodiscard]] size_t position() const { return pos_; }
 
    [[nodiscard]] bool good() const { return !error_; }
 
private:
    const uint8_t* data_;
    size_t         size_;
    size_t         pos_;
    bool           error_;
 
    // Bool values embedded in field headers are cached here.
    bool           pending_bool_;
    bool           pending_bool_valid_;
 
    static constexpr size_t   MAX_NESTING_DEPTH   = 64;               
    static constexpr size_t   MAX_FIELD_COUNT     = 65536;            
    static constexpr size_t   MAX_TOTAL_FIELDS    = 1'000'000;        
    static constexpr uint32_t MAX_STRING_BYTES    = 64u * 1024u * 1024u;  
    static constexpr uint32_t MAX_COLLECTION_SIZE = 1'000'000u;       
    std::stack<int16_t> last_field_ids_;
    size_t field_count_ = 0;
    size_t total_fields_read_ = 0;                                    
 
    // -- bounds checking ------------------------------------------------------
 
    [[nodiscard]] bool ensure(size_t n) {
        if (error_ || n > size_ || pos_ > size_ - n) {
            error_ = true;
            return false;
        }
        return true;
    }
 
    // -- varint decoding ------------------------------------------------------
 
    [[nodiscard]] uint32_t read_varint32() {
        uint32_t result = 0;
        int shift = 0;
        while (shift < 35) {
            if (!ensure(1)) return 0;
            uint8_t byte = data_[pos_++];
            result |= static_cast<uint32_t>(byte & 0x7F) << shift;
            if ((byte & 0x80) == 0) return result;
            shift += 7;
        }
        // Varint too long — malformed
        error_ = true;
        return 0;
    }
 
    [[nodiscard]] uint64_t read_varint64() {
        uint64_t result = 0;
        int shift = 0;
        while (shift < 70) {
            if (!ensure(1)) return 0;
            uint8_t byte = data_[pos_++];
            result |= static_cast<uint64_t>(byte & 0x7F) << shift;
            if ((byte & 0x80) == 0) return result;
            shift += 7;
        }
        // Varint too long — malformed
        error_ = true;
        return 0;
    }
 
    // -- zigzag decoding ------------------------------------------------------
 
    [[nodiscard]] int32_t read_zigzag_i32() {
        uint32_t raw = read_varint32();
        return static_cast<int32_t>((raw >> 1) ^ -(static_cast<int32_t>(raw & 1)));
    }
 
    [[nodiscard]] int64_t read_zigzag_i64() {
        uint64_t raw = read_varint64();
        return static_cast<int64_t>((raw >> 1) ^ -(static_cast<int64_t>(raw & 1)));
    }
};
 
} // namespace signet::forge::thrift