Search code examples
c++vectorstructtcpmalloc

Store allocated struct (using the "struct hack") in std::vector


In our assignment, we have to build our own TCP handler (framework is provided). I am now at that point where I have to collect all TCP segments and store them into a vector so I can reassemble the stream after the last segment has been received, and forward it to the next layer (e.g. HTTP). For that purpose we have to implement void handle_packet(). In our framework there are several structs and classes that should "help" us.

This is what I have come up so far:

tcp.h

#include <iostream>
#include <cstdint>
#include <vector>
#include <netinet/tcp.h> // struct tcphdr

using Header = struct tcphdr;

struct Segment {

    const Header * hdr;
    uint8_t payload[];
};

struct Stream {

  std::vector<Segment> segments;

  void add_segment(Segment new_segment);

};

void Stream::add_segment(Segment new_segment) {

    segments.push_back(new_segment);

    // testing: print all sequence numbers of each segment
    for(size_t i = 0; i < segments.size(); i++)
        std::cout << ntohl(segments[i].hdr->seq) << std::endl;
}

class Protocol {
private:
  Stream streams;

public:
  void handle_packet(const Header* tcp_segment_hdr, uint8_t* tcp_payload, size_t tcp_payload_len);
};

tcp.cpp

#include "tcp.h"

void handle_packet(const Header* tcp_segment_hdr, uint8_t* tcp_payload, size_t tcp_payload_len) {

    // "struct hack" to copy content 
    Segment* segment = reinterpret_cast<Segment *>(malloc(sizeof(struct Segment) + buffer_len * sizeof(uint8_t)) );

    // store the header
    segment->hdr = tcp_segment_hdr; // add the header

    // store the payload
    for(size_t i = 0; i < tcp_payload_len; i++)
        segment->payload[i] = tcp_payload[i];

    streams.add_segment(*segment);

    free(segment);
}

I am copying the header and the payload using the "struct hack" into a Segment * struct (this was the only working way I have come up with). After that I store the new segment into streams (member of class Protocol) member variable vector<Segments>. Then I free the memory.

After that I wanted to see if everything is stored in segments (for test purposes I tried to print each sequence number of a segment) - and it's not.. However, the size is still correct.

I guess this is because I free'd the struct in handle_packet() afterwards and this also effects my entry in std::vector<Segment>segments.. But how can I achieve that the entries are still in there without any memory leaks? I googled but did not find a satisyfing solution how I can solve my problem.

Any help is appreciated! Thanks

Edit: I am also getting the following error in AddressSanitizer:

==22658==ERROR: AddressSanitizer: stack-buffer-overflow on address 
0x7fff5288b66c at pc 0x559511b2f7bc bp 0x7fff5288b2b0 sp 0x7fff5288b2a0
WRITE of size 1 at 0x7fff5288b66c thread T0

In IPv4 we had the following struct:

struct Fragment : public std::vector<uint8_t> {
  using std::vector<uint8_t>::vector;
  const Header *getHeader() { return (const Header *)data(); }
};

But I did not do that part since I couldn't figure out how to "fill up" the values and add an entry. Any ideas how I can adopt struct Fragment to struct Segment and fill/add (persistent) values?

Edit 2: As a temporary result, I have given the payload a fixed size (1500) and remove the "struct hack" and edited it back to a simple Segment segment.

Edit 3: Ok it still doesn't work even with a fixed size..

Edit 4: I have found a working solution! However, I will describe my solution on saturday (after deadline) to avoid plagiarism.


Solution

  • As promised, here is my solution. I took the following struct definition:

    /**
    * Container for a segment of a TCP stream.
    */
    struct Segment : public std::vector<uint8_t> {
        using std::vector<uint8_t>::vector; 
    };
    

    To copy the segment data into the struct I did:

    Segment segment;
    
    // fill the segment with the data
    for(unsigned int i = 0; i < buffer_len; i++)
      segment.push_back(buffer[i]);
    

    And to access the data you can use the defined data() member from std::vector (from which the struct inherits). I also had another struct Streamwith a member called Segment allSegments where every segment was added (all segments in total). The resulting access looked like this:

    for(unsigned int i = 0; i < allSegments.size(); i++)
    {
      for(unsigned int j = 0; j < allSegments[i].size(); j++)
        printf( "%x",allSegments[i].data()[j] );
    }
    

    I didn't know that you can inherit std::vector. The struct Segment was pre-defined by our study assistents and not implemented by myself. At first I didn't get how you can access the data and wanted to implement a different way - that was why I did come up with the struct hack.