IPReassembly.h

Go to the documentation of this file.

#pragma once
 
#include "Packet.h"
#include "LRUList.h"
#include "IpAddress.h"
#include "PointerVector.h"
#include <unordered_map>
 
namespace pcpp
{
 
#define PCPP_IP_REASSEMBLY_DEFAULT_MAX_PACKETS_TO_STORE 500000
 
    class IPReassembly
    {
    public:
        class PacketKey
        {
        public:
            virtual ~PacketKey() = default;
 
            virtual uint32_t getHashValue() const = 0;
 
            virtual ProtocolType getProtocolType() const = 0;
 
            virtual PacketKey* clone() const = 0;
 
        protected:
            // private c'tor
            PacketKey() = default;
 
            // private copy c'tor
            PacketKey(const PacketKey& other) = default;
        };
 
        class IPv4PacketKey : public PacketKey
        {
        public:
            IPv4PacketKey() : m_IpID(0), m_SrcIP(IPv4Address::Zero), m_DstIP(IPv4Address::Zero)
            {}
 
            IPv4PacketKey(uint16_t ipid, IPv4Address srcip, IPv4Address dstip)
                : m_IpID(ipid), m_SrcIP(srcip), m_DstIP(dstip)
            {}
 
            IPv4PacketKey(const IPv4PacketKey& other)
                : PacketKey(other), m_IpID(other.m_IpID), m_SrcIP(other.m_SrcIP), m_DstIP(other.m_DstIP)
            {}
 
            IPv4PacketKey& operator=(const IPv4PacketKey& other)
            {
                m_IpID = other.m_IpID;
                m_SrcIP = other.m_SrcIP;
                m_DstIP = other.m_DstIP;
                return *this;
            }
 
            uint16_t getIpID() const
            {
                return m_IpID;
            }
 
            IPv4Address getSrcIP() const
            {
                return m_SrcIP;
            }
 
            IPv4Address getDstIP() const
            {
                return m_DstIP;
            }
 
            void setIpID(uint16_t ipID)
            {
                m_IpID = ipID;
            }
 
            void setSrcIP(const IPv4Address& srcIP)
            {
                m_SrcIP = srcIP;
            }
 
            void setDstIP(const IPv4Address& dstIP)
            {
                m_DstIP = dstIP;
            }
 
            // implement abstract methods
 
            uint32_t getHashValue() const;
 
            ProtocolType getProtocolType() const
            {
                return IPv4;
            }
 
            PacketKey* clone() const
            {
                return new IPv4PacketKey(*this);
            }
 
        private:
            uint16_t m_IpID;
            IPv4Address m_SrcIP;
            IPv4Address m_DstIP;
        };
 
        class IPv6PacketKey : public PacketKey
        {
        public:
            IPv6PacketKey() : m_FragmentID(0), m_SrcIP(IPv6Address::Zero), m_DstIP(IPv6Address::Zero)
            {}
 
            IPv6PacketKey(uint32_t fragmentID, IPv6Address srcip, IPv6Address dstip)
                : m_FragmentID(fragmentID), m_SrcIP(srcip), m_DstIP(dstip)
            {}
 
            IPv6PacketKey(const IPv6PacketKey& other)
                : PacketKey(other), m_FragmentID(other.m_FragmentID), m_SrcIP(other.m_SrcIP), m_DstIP(other.m_DstIP)
            {}
 
            IPv6PacketKey& operator=(const IPv6PacketKey& other)
            {
                m_FragmentID = other.m_FragmentID;
                m_SrcIP = other.m_SrcIP;
                m_DstIP = other.m_DstIP;
                return *this;
            }
 
            uint32_t getFragmentID() const
            {
                return m_FragmentID;
            }
 
            IPv6Address getSrcIP() const
            {
                return m_SrcIP;
            }
 
            IPv6Address getDstIP() const
            {
                return m_DstIP;
            }
 
            void setFragmentID(uint32_t fragID)
            {
                m_FragmentID = fragID;
            }
 
            void setSrcIP(const IPv6Address& srcIP)
            {
                m_SrcIP = srcIP;
            }
 
            void setDstIP(const IPv6Address& dstIP)
            {
                m_DstIP = dstIP;
            }
 
            // implement abstract methods
 
            uint32_t getHashValue() const;
 
            ProtocolType getProtocolType() const
            {
                return IPv6;
            }
 
            PacketKey* clone() const
            {
                return new IPv6PacketKey(*this);
            }
 
        private:
            uint32_t m_FragmentID;
            IPv6Address m_SrcIP;
            IPv6Address m_DstIP;
        };
 
        typedef void (*OnFragmentsClean)(const PacketKey* key, void* userCookie);
 
        enum ReassemblyStatus
        {
            NON_IP_PACKET = 0x00,
            NON_FRAGMENT = 0x01,
            FIRST_FRAGMENT = 0x02,
            FRAGMENT = 0x04,
            OUT_OF_ORDER_FRAGMENT = 0x08,
            MALFORMED_FRAGMENT = 0x10,
            REASSEMBLED = 0x20
        };
 
        explicit IPReassembly(OnFragmentsClean onFragmentsCleanCallback = nullptr, void* callbackUserCookie = nullptr,
                              size_t maxPacketsToStore = PCPP_IP_REASSEMBLY_DEFAULT_MAX_PACKETS_TO_STORE)
            : m_PacketLRU(maxPacketsToStore), m_OnFragmentsCleanCallback(onFragmentsCleanCallback),
              m_CallbackUserCookie(callbackUserCookie)
        {}
 
        ~IPReassembly();
 
        Packet* processPacket(Packet* fragment, ReassemblyStatus& status, ProtocolType parseUntil = UnknownProtocol,
                              OsiModelLayer parseUntilLayer = OsiModelLayerUnknown);
 
        Packet* processPacket(RawPacket* fragment, ReassemblyStatus& status, ProtocolType parseUntil = UnknownProtocol,
                              OsiModelLayer parseUntilLayer = OsiModelLayerUnknown);
 
        Packet* getCurrentPacket(const PacketKey& key);
 
        void removePacket(const PacketKey& key);
 
        size_t getMaxCapacity() const
        {
            return m_PacketLRU.getMaxSize();
        }
 
        size_t getCurrentCapacity() const
        {
            return m_FragmentMap.size();
        }
 
    private:
        struct IPFragment
        {
            uint16_t fragmentOffset;
            bool lastFragment;
            uint8_t* fragmentData;
            size_t fragmentDataLen;
            IPFragment()
            {
                fragmentOffset = 0;
                lastFragment = false;
                fragmentData = nullptr;
                fragmentDataLen = 0;
            }
            ~IPFragment()
            {
                delete[] fragmentData;
            }
        };
 
        struct IPFragmentData
        {
            uint16_t currentOffset;
            RawPacket* data;
            bool deleteData;
            uint32_t fragmentID;
            PacketKey* packetKey;
            PointerVector<IPFragment> outOfOrderFragments;
            IPFragmentData(PacketKey* pktKey, uint32_t fragId)
            {
                currentOffset = 0;
                data = nullptr;
                deleteData = true;
                fragmentID = fragId;
                packetKey = pktKey;
            }
            ~IPFragmentData()
            {
                delete packetKey;
                if (deleteData && data != nullptr)
                {
                    delete data;
                }
            }
        };
 
        LRUList<uint32_t> m_PacketLRU;
        std::unordered_map<uint32_t, IPFragmentData*> m_FragmentMap;
        OnFragmentsClean m_OnFragmentsCleanCallback;
        void* m_CallbackUserCookie;
 
        void addNewFragment(uint32_t hash, IPFragmentData* fragData);
        bool matchOutOfOrderFragments(IPFragmentData* fragData);
    };
 
}  // namespace pcpp