ip_reassembler.h

/*
 * Copyright (c) 2017, Matias Fontanini
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 * 
 * * Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 * * Redistributions in binary form must reproduce the above
 *   copyright notice, this list of conditions and the following disclaimer
 *   in the documentation and/or other materials provided with the
 *   distribution.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

#ifndef TINS_IP_REASSEMBLER_H
#define TINS_IP_REASSEMBLER_H

#include <vector>
#include <map>
#include <tins/pdu.h>
#include <tins/macros.h>
#include <tins/ip_address.h>
#include <tins/ip.h>

namespace Tins {

namespace Internals {
class IPv4Fragment {
public:
    typedef PDU::serialization_type payload_type;

    IPv4Fragment() : offset_() { }

    template<typename T>
    IPv4Fragment(T* pdu, uint16_t offset)
    : payload_(pdu->serialize()), offset_(offset) {
        
    }
    
    const payload_type& payload() const {
        return payload_;
    }
    
    uint16_t offset() const {
        return offset_;
    }
private:
    payload_type payload_;
    uint16_t offset_;
};

class TINS_API IPv4Stream {
public:
    IPv4Stream();
    
    void add_fragment(IP* ip);
    bool is_complete() const;
    PDU* allocate_pdu() const;
    const IP& first_fragment() const;
private:
    typedef std::vector<IPv4Fragment> fragments_type;
    
    uint16_t extract_offset(const IP* ip);
    bool extract_more_frag(const IP* ip);

    fragments_type fragments_;
    size_t received_size_;
    size_t total_size_;
    IP first_fragment_;
    bool received_end_;
};
} // namespace Internals

class TINS_API IPv4Reassembler {
public:
    enum PacketStatus {
        NOT_FRAGMENTED, 
        FRAGMENTED, 
        REASSEMBLED 
    };

    TINS_DEPRECATED(typedef PacketStatus packet_status);

    enum OverlappingTechnique {
        NONE 
    };

    IPv4Reassembler();

    IPv4Reassembler(OverlappingTechnique technique);

    PacketStatus process(PDU& pdu);

    void clear_streams();

    void remove_stream(uint16_t id, IPv4Address addr1, IPv4Address addr2);
private:
    typedef std::pair<IPv4Address, IPv4Address> address_pair;
    typedef std::pair<uint16_t, address_pair> key_type;
    typedef std::map<key_type, Internals::IPv4Stream> streams_type;

    key_type make_key(const IP* ip) const;
    address_pair make_address_pair(IPv4Address addr1, IPv4Address addr2) const;
    
    streams_type streams_;
    OverlappingTechnique technique_;
};

template<typename Functor>
class IPv4ReassemblerProxy {
public:
    IPv4ReassemblerProxy(Functor func)
    : functor_(func) {

    }

    bool operator()(PDU& pdu) {
        // Forward it unless it's fragmented.
        if (reassembler_.process(pdu) != IPv4Reassembler::FRAGMENTED) {
            return functor_(pdu);
        }
        else {
            return true;
        }
    }
private:
    IPv4Reassembler reassembler_;
    Functor functor_;
};

template<typename Functor>
IPv4ReassemblerProxy<Functor> make_ipv4_reassembler_proxy(Functor func) {
    return IPv4ReassemblerProxy<Functor>(func);
}

} // Tins

#endif // TINS_IP_REASSEMBLER_H