sink.h

// Copyright (c) 2020-now by the Zeek Project. See LICENSE for details.
 
#pragma once
 
#include <list>
#include <string>
#include <string_view>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <hilti/rt/exception.h>
#include <hilti/rt/extension-points.h>
#include <hilti/rt/types/bytes.h>
#include <hilti/rt/types/integer.h>
#include <hilti/rt/types/reference.h>
#include <hilti/rt/types/stream.h>
 
#include <spicy/rt/debug.h>
#include <spicy/rt/filter.h>
#include <spicy/rt/mime.h>
#include <spicy/rt/parser-fwd.h>
#include <spicy/rt/typedefs.h>
 
namespace spicy::rt {
 
HILTI_EXCEPTION(SinkError, UsageError)
 
namespace sink {
enum class ReassemblerPolicy { First };
} // namespace sink
 
namespace sink::detail {
 
template<typename T, typename = int>
struct supports_sinks : std::false_type {};
 
template<typename T>
struct supports_sinks<T, decltype((void)T::__sink, 0)> : std::true_type {};
 
struct State {
    hilti::rt::ValueReference<hilti::rt::Stream> data;
 
    hilti::rt::Resumable resumable;
 
    Parser* parser;
 
    bool skip_delivery = false;
};
 
template<typename U>
auto connectUnit(UnitRef<U>& unit) {
    auto parse2 = hilti::rt::any_cast<spicy::rt::Parse2Function<U>>(U::__parser.parse2);
 
    auto self = hilti::rt::ValueReference<U>::self(&*unit);
 
    auto& state = unit->__sink;
    state = new sink::detail::State();                       // NOLINT
    state->resumable = (*parse2)(self, state->data, {}, {}); // Kick-off parsing with empty data.
    state->parser = &U::__parser;
    return state;
}
 
// Name used as template parameter for sink's filter state. */
inline const char sink_name[] = "__sink__";
} // namespace sink::detail
 
class Sink {
public:
    Sink() { _init(); } // NOLINT(hicpp-member-init)
    ~Sink() {
        try {
            _close(true);
        } catch ( ... ) {
            // Ignore errors.
        }
    }
 
    Sink(const Sink&) = delete;
    Sink(Sink&&) = default;
    Sink& operator=(const Sink&) = delete;
    Sink& operator=(Sink&&) = default;
 
    template<typename T>
    void connect(spicy::rt::UnitRef<T> unit) {
        SPICY_RT_DEBUG_VERBOSE(hilti::rt::fmt("connecting parser %s [%p] to sink %p", T::__parser.name, &*unit, this));
        auto state = spicy::rt::sink::detail::connectUnit(unit);
        _units.emplace_back(std::move(unit));
        _states.emplace_back(std::move(state));
    }
 
    template<typename T>
    void connect_filter(spicy::rt::UnitRef<T> filter_unit) {
        if ( _size )
            throw SinkError("cannot connect filter after data has been forwarded already");
 
        SPICY_RT_DEBUG_VERBOSE(
            hilti::rt::fmt("connecting filter unit %s [%p] to sink %p", T::__parser.name, &*filter_unit, this));
        spicy::rt::filter::detail::connect(_filter, filter_unit);
    }
 
    void close() { _close(true); }
 
    void connect_mime_type(const MIMEType& mt, uint64_t scope);
 
    void connect_mime_type(const std::string& mt, uint64_t scope) { connect_mime_type(MIMEType(mt), scope); }
 
    void connect_mime_type(const hilti::rt::Bytes& mt, uint64_t scope) { connect_mime_type(MIMEType(mt.str()), scope); }
 
    void gap(uint64_t seq, uint64_t len);
 
    uint64_t sequence_number() const { return _initial_seq + _cur_rseq; }
 
    void set_auto_trim(bool enable) { _auto_trim = enable; }
 
    void set_initial_sequence_number(uint64_t seq) {
        if ( _haveInput() ) {
            _close(false);
            throw SinkError("sink cannot update initial sequence number after activity has already been seen");
        }
 
        _initial_seq = seq;
    }
 
    void set_policy(sink::ReassemblerPolicy policy) { _policy = policy; }
 
    hilti::rt::integer::safe<uint64_t> size() const { return _size; }
 
    void skip(uint64_t seq);
 
    void trim(uint64_t seq);
 
    void write(hilti::rt::Bytes data, std::optional<uint64_t> seq = {}, std::optional<uint64_t> len = {});
 
    filter::State<sink::detail::sink_name> _filter;
 
private:
    struct Chunk {
        std::optional<hilti::rt::Bytes> data; // Data at +1; unset for gap
        uint64_t rseq;                        // Sequence number of first byte.
        uint64_t rupper;                      // Sequence number of last byte + 1.
 
        Chunk(std::optional<hilti::rt::Bytes> data, uint64_t rseq, uint64_t rupper)
            : data(std::move(data)), rseq(rseq), rupper(rupper) {}
    };
 
    using ChunkList = std::list<Chunk>;
 
    // Returns true if any input has been passed in already (including gaps).
    bool _haveInput() { return _cur_rseq || _chunks.size(); }
 
    // Backend for disconnecting the sink. If orderly, connected units get a
    // chance to parse any remaining input; otherwise we abort directly.
    void _close(bool orderly);
 
    // Turns an absolute sequence number into a relative one.
    uint64_t _rseq(uint64_t seq) const {
        // I believe this does the right thing for wrap-around ...
        return seq - _initial_seq;
    }
 
    // Turns a relative sequence number into an absolute one.
    uint64_t _aseq(uint64_t rseq) const {
        // I believe this does the right thing for wrap-around ...
        return _initial_seq + rseq;
    }
 
    // (Re-)initialize instance.
    void _init();
 
    // Add new data to buffer, beginning search for insert position at given start *c*.
    ChunkList::iterator _addAndCheck(std::optional<hilti::rt::Bytes> data, uint64_t rseq, uint64_t rupper,
                                     ChunkList::iterator c);
 
    // Deliver data to connected parsers. Returns false if the data is empty (i.e., a gap).
    bool _deliver(std::optional<hilti::rt::Bytes> data, uint64_t rseq, uint64_t rupper);
 
    // Entry point for all new data. If not bytes instance is given, that signals a gap.
    void _newData(std::optional<hilti::rt::Bytes> data, uint64_t rseq, uint64_t len);
 
    // Skip up to sequence number.
    void _skip(uint64_t rseq);
 
    // Trim up to sequence number.
    void _trim(uint64_t rseq);
 
    // Deliver as much as possible starting at given buffer position.
    void _tryDeliver(ChunkList::iterator c);
 
    // Trigger various hooks.
    void _reportGap(uint64_t rseq, uint64_t len) const;
    void _reportOverlap(uint64_t rseq, const hilti::rt::Bytes& old, const hilti::rt::Bytes& new_) const;
    void _reportSkipped(uint64_t rseq) const;
    void _reportUndelivered(uint64_t rseq, const hilti::rt::Bytes& data) const;
    void _reportUndeliveredUpTo(uint64_t rupper) const;
 
    // Output reassembler state for debugging.
    void _debugReassembler(std::string_view msg, const std::optional<hilti::rt::Bytes>& data, uint64_t seq,
                           uint64_t len) const;
    void _debugReassemblerBuffer(std::string_view msg) const;
    void _debugDeliver(const hilti::rt::Bytes& data) const;
 
    // States for connected units.
    std::vector<sink::detail::State*> _states;
 
    // Must come after `_state` as it's keeping the states around.
    std::vector<hilti::rt::StrongReferenceGeneric> _units;
 
    // Filter input and output.
    struct FilterData {
        hilti::rt::ValueReference<hilti::rt::Stream> input;
        hilti::rt::StrongReference<hilti::rt::Stream> output;
        hilti::rt::stream::View output_cur;
    };
 
    std::optional<FilterData> _filter_data;
 
    // Reassembly state.
    sink::ReassemblerPolicy _policy; // Current policy
    bool _auto_trim{};               // True if automatic trimming is enabled.
    uint64_t _size{};
    uint64_t _initial_seq{};       // Initial sequence number.
    uint64_t _cur_rseq{};          // Sequence of last delivered byte + 1 (i.e., seq of next)
    uint64_t _last_reassem_rseq{}; // Sequence of last byte reassembled and delivered + 1.
    uint64_t _trim_rseq{};         // Sequence of last byte trimmed so far + 1.
    ChunkList _chunks;             // Buffered data not yet delivered or trimmed
};
 
} // namespace spicy::rt
 
namespace hilti::rt::detail::adl {
std::string to_string(const spicy::rt::Sink& /* x */, adl::tag /*unused*/);
std::string to_string(const spicy::rt::sink::ReassemblerPolicy& x, adl::tag /*unused*/);
} // namespace hilti::rt::detail::adl