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

#pragma once

#include <sys/types.h> // for u_char
#include <cassert>
#include <cstdint>
#include <cstring>
#include <map>

#include "zeek/Obj.h"

namespace zeek
	{

// Whenever subclassing the Reassembler class
// you should add to this for known subclasses.
enum ReassemblerType
	{
	REASSEM_UNKNOWN,
	REASSEM_TCP,
	REASSEM_FRAG,
	REASSEM_FILE,

	// Terminal value. Add new above.
	REASSEM_NUM,
	};

class Reassembler;

/**
 * A block/segment of data for use in the reassembly process.
 */
class DataBlock
	{
public:
	/**
	 * Create a data block/segment with associated sequence numbering.
	 */
	DataBlock(const u_char* data, uint64_t size, uint64_t seq);

	DataBlock(const DataBlock& other)
		{
		seq = other.seq;
		upper = other.upper;
		auto size = other.Size();
		block = new u_char[size];
		memcpy(block, other.block, size);
		}

	DataBlock(DataBlock&& other)
		{
		seq = other.seq;
		upper = other.upper;
		block = other.block;
		other.block = nullptr;
		}

	DataBlock& operator=(const DataBlock& other)
		{
		if ( this == &other )
			return *this;

		seq = other.seq;
		upper = other.upper;
		auto size = other.Size();
		delete[] block;
		block = new u_char[size];
		memcpy(block, other.block, size);
		return *this;
		}

	DataBlock& operator=(DataBlock&& other)
		{
		if ( this == &other )
			return *this;

		seq = other.seq;
		upper = other.upper;
		delete[] block;
		block = other.block;
		other.block = nullptr;
		return *this;
		}

	~DataBlock() { delete[] block; }

	/**
	 * @return length of the data block
	 */
	uint64_t Size() const { return upper - seq; }

	uint64_t seq;
	uint64_t upper;
	u_char* block;
	};

using DataBlockMap = std::map<uint64_t, DataBlock>;

/**
 * The data structure used for reassembling arbitrary sequences of data
 * blocks/segments.  It internally uses an ordered map (std::map).
 */
class DataBlockList
	{
public:
	DataBlockList() { }

	DataBlockList(Reassembler* r) : reassembler(r) { }

	~DataBlockList() { Clear(); }

	/**
	 * @return iterator to start of the block list.
	 */
	DataBlockMap::const_iterator Begin() const { return block_map.begin(); }

	/**
	 * @return iterator to end of the block list (one past last element).
	 */
	DataBlockMap::const_iterator End() const { return block_map.end(); }

	/**
	 * @return reference to the first data block in the list.
	 * Must not be called when the list is empty.
	 */
	const DataBlock& FirstBlock() const
		{
		assert(block_map.size());
		return block_map.begin()->second;
		}

	/**
	 * @return reference to the last data block in the list.
	 * Must not be called when the list is empty.
	 */
	const DataBlock& LastBlock() const
		{
		assert(block_map.size());
		return block_map.rbegin()->second;
		}

	/**
	 * @return whether the list is empty.
	 */
	bool Empty() const { return block_map.empty(); };

	/**
	 * @return the number of blocks in the list.
	 */
	size_t NumBlocks() const { return block_map.size(); };

	/**
	 * @return the total size, in bytes, of all blocks in the list.
	 */
	size_t DataSize() const { return total_data_size; }

	/**
	 * Counts the total size of all data contained in list elements
	 * partitioned by some cutoff.
	 * WARNING: this is an O(n) operation and potentially very slow.
	 * @param seq_cutoff  the sequence number used to partition
	 * element sizes returned via "below" and "above" parameters
	 * @param below  the size in bytes of all data below "seq_cutoff"
	 * @param above  the size in bytes of all data above "seq_cutoff"
	 */
	void DataSize(uint64_t seq_cutoff, uint64_t* below, uint64_t* above) const;

	/**
	 * Remove all elements from the list
	 */
	void Clear();

	/**
	 * Insert a new data block into the list.
	 * @param seq  lower sequence number of the data block
	 * @param upper  highest sequence number of the data block
	 * @param data  points to the data block contents
	 * @param hint  a suggestion of the node from which to start searching
	 * for an insertion point or null to search from the beginning of the list
	 * @return an iterator to the element that was inserted
	 */
	DataBlockMap::const_iterator Insert(uint64_t seq, uint64_t upper, const u_char* data,
	                                    DataBlockMap::const_iterator* hint = nullptr);

	/**
	 * Insert a new data block at the end of the list and remove blocks
	 * from the beginning of the list to keep the list size under a limit.
	 * @param block  the block to append
	 * @param limit  the max number of blocks allowed (list is pruned from
	 * starting from the beginning after the insertion takes place).
	 */
	void Append(DataBlock block, uint64_t limit);

	/**
	 * Remove all elements below a given sequence number.
	 * @param seq  blocks below this number are discarded (removed/deleted)
	 * @param max_old  if non-zero instead of deleting the underlying block,
	 * move it to "old_list"
	 * @param old_list  another list to move discarded blocks into
	 * @return the amount of data (in bytes) that was not part of any
	 * discarded block (the total size of all bypassed gaps).
	 */
	uint64_t Trim(uint64_t seq, uint64_t max_old, DataBlockList* old_list);

	/**
	 * @return an iterator pointing to the first element with a segment whose
	 * starting sequence number is less than or equal to "seq".  If no such
	 * element exists, returns an iterator denoting one-past the end of the
	 * list.
	 */
	DataBlockMap::const_iterator FirstBlockAtOrBefore(uint64_t seq) const;

private:
	/**
	 * Insert a new data block into the list.
	 * @param seq  lower sequence number of the data block
	 * @param upper  highest sequence number of the data block
	 * @param data  points to the data block contents
	 * @param hint  a suggestion of the node from which to start searching
	 * for an insertion point
	 * @return an iterator to the element that was inserted
	 */
	DataBlockMap::const_iterator Insert(uint64_t seq, uint64_t upper, const u_char* data,
	                                    DataBlockMap::const_iterator hint);

	/**
	 * Removes a block from the list and updates other state which keeps
	 * track of total size of blocks.
	 * @param it  the element to remove
	 */
	void Delete(DataBlockMap::const_iterator it);

	/**
	 * Removes a block from the list and returns it, assuming it will
	 * immediately be appended to another list.
	 * @param it  the element to remove
	 * @return the removed block
	 */
	DataBlock Remove(DataBlockMap::const_iterator it);

	Reassembler* reassembler = nullptr;
	size_t total_data_size = 0;
	DataBlockMap block_map;
	};

class Reassembler : public Obj
	{
public:
	Reassembler(uint64_t init_seq, ReassemblerType reassem_type = REASSEM_UNKNOWN);
	~Reassembler() override { }

	void NewBlock(double t, uint64_t seq, uint64_t len, const u_char* data);

	// Throws away all blocks up to seq.  Returns number of bytes
	// if not all in-sequence, 0 if they were.
	uint64_t TrimToSeq(uint64_t seq);

	// Delete all held blocks.
	void ClearBlocks();
	void ClearOldBlocks();

	bool HasBlocks() const { return ! block_list.Empty(); }

	uint64_t LastReassemSeq() const { return last_reassem_seq; }

	uint64_t TrimSeq() const { return trim_seq; }

	void SetTrimSeq(uint64_t seq)
		{
		if ( seq > trim_seq )
			trim_seq = seq;
		}

	uint64_t TotalSize() const; // number of bytes buffered up

	void Describe(ODesc* d) const override;

	static uint64_t TotalMemoryAllocation() { return total_size; }

	// Data buffered by type of reassembler.
	static uint64_t MemoryAllocation(ReassemblerType rtype);

	void SetMaxOldBlocks(uint32_t count) { max_old_blocks = count; }

protected:
	friend class DataBlockList;

	virtual void Undelivered(uint64_t up_to_seq);

	virtual void BlockInserted(DataBlockMap::const_iterator it) = 0;
	virtual void Overlap(const u_char* b1, const u_char* b2, uint64_t n) = 0;

	void CheckOverlap(const DataBlockList& list, uint64_t seq, uint64_t len, const u_char* data);

	DataBlockList block_list;
	DataBlockList old_block_list;

	uint64_t last_reassem_seq = 0;
	uint64_t trim_seq = 0; // how far we've trimmed
	uint32_t max_old_blocks = 0;

	ReassemblerType rtype = REASSEM_UNKNOWN;

	static uint64_t total_size;
	static uint64_t sizes[REASSEM_NUM];
	};

	} // namespace zeek
