// See the file "COPYING" in the main distribution directory for copyright.

#pragma once

#include <sys/types.h> // for u_char
#include <tuple>

#include "zeek/IPAddr.h"
#include "zeek/Reassem.h"
#include "zeek/Timer.h"
#include "zeek/util-types.h"

namespace zeek {

class IP_Hdr;

namespace session {
class Manager;
}

namespace detail {

class FragReassembler;
class FragTimer;

using FragReassemblerKey = std::tuple<IPAddr, IPAddr, zeek_uint_t>;

class FragReassembler : public Reassembler {
public:
    FragReassembler(session::Manager* s, const std::shared_ptr<IP_Hdr>& ip, const u_char* pkt,
                    const FragReassemblerKey& k, double t);
    ~FragReassembler() override;

    void AddFragment(double t, const std::shared_ptr<IP_Hdr>& ip, const u_char* pkt);

    void Expire(double t);
    void DeleteTimer();
    void ClearTimer() { expire_timer = nullptr; }

    std::shared_ptr<IP_Hdr> ReassembledPkt() { return std::move(reassembled_pkt); }
    const FragReassemblerKey& Key() const { return key; }

protected:
    void BlockInserted(DataBlockMap::const_iterator it) override;
    void Overlap(const u_char* b1, const u_char* b2, uint64_t n) override;
    void Weird(const char* name) const;

    u_char* proto_hdr;
    std::shared_ptr<IP_Hdr> reassembled_pkt;
    session::Manager* s;
    uint64_t frag_size; // size of fully reassembled fragment
    FragReassemblerKey key;
    uint16_t next_proto; // first IPv6 fragment header's next proto field
    uint16_t proto_hdr_len;

    FragTimer* expire_timer;
};

class FragTimer final : public Timer {
public:
    FragTimer(FragReassembler* arg_f, double arg_t) : Timer(arg_t, TIMER_FRAG) { f = arg_f; }
    ~FragTimer() override;

    void Dispatch(double t, bool is_expire) override;

    // Break the association between this timer and its creator.
    void ClearReassembler() { f = nullptr; }

protected:
    FragReassembler* f;
};

class FragmentManager {
public:
    FragmentManager() = default;
    ~FragmentManager();

    FragReassembler* NextFragment(double t, const std::shared_ptr<IP_Hdr>& ip, const u_char* pkt);
    void Clear();
    void Remove(detail::FragReassembler* f);

    size_t Size() const { return fragments.size(); }
    size_t MaxFragments() const { return max_fragments; }

private:
    using FragmentMap = std::map<detail::FragReassemblerKey, detail::FragReassembler*>;
    FragmentMap fragments;
    size_t max_fragments = 0;
};

extern FragmentManager* fragment_mgr;

class FragReassemblerTracker {
public:
    FragReassemblerTracker(FragReassembler* f) : frag_reassembler(f) {}

    ~FragReassemblerTracker() { fragment_mgr->Remove(frag_reassembler); }

private:
    FragReassembler* frag_reassembler;
};

} // namespace detail
} // namespace zeek