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) : m_IpID(other.m_IpID), m_SrcIP(other.m_SrcIP), m_DstIP(other.m_DstIP) { }

            inline uint16_t getIpID() { return m_IpID; }

            inline IPv4Address getSrcIP() { return m_SrcIP; }

            inline IPv4Address getDstIP() { return m_DstIP; }

            inline void setIpID(uint16_t ipID) { m_IpID = ipID; }

            inline void setSrcIP(const IPv4Address& srcIP) { m_SrcIP = srcIP; }

            inline 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) : m_FragmentID(other.m_FragmentID), m_SrcIP(other.m_SrcIP), m_DstIP(other.m_DstIP) { }

            inline uint32_t getFragmentID() { return m_FragmentID; }

            inline IPv6Address getSrcIP() { return m_SrcIP; }

            inline IPv6Address getDstIP() { return m_DstIP; }

            inline void setFragmentID(uint32_t fragID) { m_FragmentID = fragID; }

            inline void setSrcIP(const IPv6Address& srcIP) { m_SrcIP = srcIP; }

            inline 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);

        ~IPReassembly();

        Packet* processPacket(Packet* fragment, ReassemblyStatus& status);

        Packet* processPacket(RawPacket* fragment, ReassemblyStatus& status);

        Packet* getCurrentPacket(const PacketKey& key);

        void removePacket(const PacketKey& key);

        inline size_t getMaxCapacity() { return (int)m_PacketLRU->getMaxSize(); }

        inline size_t getCurrentCapacity() { 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