IPReassembly.h

Go to the documentation of this file.

#ifndef PACKETPP_IP_REASSEMBLY
#define PACKETPP_IP_REASSEMBLY

#include "Packet.h"
#include "LRUList.h"
#include "IpAddress.h"
#include "PointerVector.h"
#include <map>

namespace pcpp
{

    #define PCPP_IP_REASSEMBLY_DEFAULT_MAX_PACKETS_TO_STORE 500000

    class IPReassembly
    {
    public:

        class PacketKey
        {
        public:

            virtual ~PacketKey() {}

            virtual uint32_t getHashValue() const = 0;

            virtual ProtocolType getProtocolType() const = 0;

            virtual PacketKey* clone() const = 0;

        protected:
            // private c'tor
            PacketKey() {}

            // private copy c'tor
            PacketKey(const PacketKey& other) {}
        };


        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_SrcIP = 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
        };

        IPReassembly(OnFragmentsClean onFragmentsCleanCallback = NULL, void *callbackUserCookie = NULL, 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 = NULL; 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 = NULL; deleteData = true; fragmentID = fragId; packetKey = pktKey; }
            ~IPFragmentData() { delete packetKey; if (deleteData && data != NULL) { delete data; } }
        };

        LRUList<uint32_t> m_PacketLRU;
        std::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

#endif // PACKETPP_IP_REASSEMBLY