/*
 * A tool to check checksum integrity and gather simple traffic statistics.
 *
 * Fragment reassembly.
 *
 * James Bulpin, October 2004.
 *
 * (C) University of Cambridge 2004.
 *
 * $Id: frag.c,v 1.1 2004/11/09 16:15:47 jrb44 Exp $
 *
 * $Log: frag.c,v $
 * Revision 1.1  2004/11/09 16:15:47  jrb44
 * Initial revision
 *
 *
 */

#include <stdio.h>
#include <unistd.h>
#include <stdint.h>
#include <stdlib.h>
#include <stdarg.h>
#include <string.h>
#include <netinet/in.h>
#include <netdb.h>
#include <time.h>

#include <assert.h>

#include "pcap.h"
#include "headers.h"
#include "frag.h"
#include "cksumck.h"

#define DEBUG

struct frag {
    struct frag       *prev;
    struct frag       *next;
    struct frag       *more_prev;
    struct frag       *more_next;
    unsigned long long age;
    uint16_t           offset;
    struct pcap_pkthdr pkthdr;
    u_char            *user;
    u_char            *sp;
};

/*
 * handle_frag: if this frag completes a datagram then pass that back to
 * the main packet handler. Otherwise, store.
 */

struct frag *frag_head = NULL;
struct frag *frag_tail = NULL;

static unsigned long long agecount = 0;

int handle_frag(u_char                   *user,
                const struct pcap_pkthdr *h,
                const u_char             *sp,
                unsigned int              ll_hdr_len)
{
    struct frag *ptr;
    uint16_t offset   = ntohs(EXTRACT_IPHDR(sp)->flags_frag) & IP_OFFSET;
    uint16_t ident_ne = EXTRACT_IPHDR(sp)->ident;
    addr_t   src_ne   = EXTRACT_IPHDR(sp)->src;
    struct frag *f;

    agecount++;

    /* Remove old frags?
     */
    if (!(agecount&0xff)) {
        for (ptr = frag_head; ptr;) {
            assert(ptr);
            if ((ptr->age + 256) < agecount) {
                struct frag *more;
                assert(ptr->sp);
                fprintf(stderr,
                        "FRAGDB: discarding %llu %08x/%04x offset=%u "
                        "length=%u(%u/%u)%s%s\n",
                        ptr->age,
                        (unsigned int)(EXTRACT_IPHDR(ptr->sp)->src),
                        (unsigned int)(EXTRACT_IPHDR(ptr->sp)->ident),
                        (unsigned int)(ntohs(EXTRACT_IPHDR(ptr->sp)->flags_frag)
                                       & IP_OFFSET),
                        (unsigned int)
                        ntohs(EXTRACT_IPHDR(ptr->sp)->total_length),
                        (unsigned int)IP_HDR_LEN(EXTRACT_IPHDR(ptr->sp)),
                        (unsigned int)
                        (ntohs(EXTRACT_IPHDR(ptr->sp)->total_length) -
                         IP_HDR_LEN(EXTRACT_IPHDR(ptr->sp))),
                        ((ntohs(EXTRACT_IPHDR(ptr->sp)->flags_frag)&IP_MF))?
                        " MF":"",
                        ((ntohs(EXTRACT_IPHDR(ptr->sp)->flags_frag)&IP_MF))?
                        " DF":"");
                
                if (ptr->next)
                    ptr->next->prev = ptr->prev;
                else
                    frag_tail = ptr->prev;
                if (ptr->prev)
                    ptr->prev->next = ptr->next;
                else
                    frag_head = ptr->next;
                
                more = ptr;
                ptr = ptr->next;
                for (; more;) {
                    struct frag *df;
                    assert(more);
                    df = more->more_next;
                    
                    if (more->sp)
                        free(more->sp);
                    free(more);
                    
                    more = df;
                }

            }
            else
                ptr = ptr->next;
        }
    }

#ifdef DEBUG
    fprintf(stderr,
            "FRAGDB Frag %08x/%04x offset=%u length=%u(%u/%u)%s\n",
            (unsigned int)src_ne,
            (unsigned int)ident_ne,
            (unsigned int)offset,
            (unsigned int)ntohs(EXTRACT_IPHDR(sp)->total_length),
            (unsigned int)IP_HDR_LEN(EXTRACT_IPHDR(sp)),
            (unsigned int)(ntohs(EXTRACT_IPHDR(sp)->total_length) -
                           IP_HDR_LEN(EXTRACT_IPHDR(sp))),
            ((ntohs(EXTRACT_IPHDR(sp)->flags_frag)&IP_MF))?" MF":"");
#endif

    f = (struct frag *)malloc(sizeof(struct frag));
    if (!f) {
        fprintf(stderr, "Could not malloc for frag storage.\n");
        return -1;
    }
    assert(f);
    f->prev = NULL;
    f->next = NULL;
    f->more_prev = NULL;
    f->more_next = NULL;
    f->age = agecount;

    assert(h);
    memcpy((void *)&(f->pkthdr), h, sizeof(struct pcap_pkthdr));
    f->user = user;
    f->offset = offset;

    f->sp = (u_char *)malloc(h->caplen);
    if (!f->sp) {
        fprintf(stderr,
                "Could not malloc for frag data (%u bytes).\n",
                h->caplen);
        free(f);
        return -1;
    }
    assert(f->sp);
    assert(sp);
    memcpy(f->sp, sp, h->caplen);

    /* Trawl list to see if we have other frags with the same ident and
     * source address. Stick our new frag into the correct chain. 
     */   
    for (ptr = frag_head; ptr; ptr = ptr->next) {
        assert(ptr);
        if ((EXTRACT_IPHDR(ptr->sp)->ident == ident_ne) &&
            (EXTRACT_IPHDR(ptr->sp)->src == src_ne)) {
            
            struct frag *more;
            unsigned int reached = 0;

            /* Insert into chain.
             */
            for (more = ptr; more; more = more->more_next) {
                assert(more);
                if (offset < more->offset) {
                    /* Insert before this entry.
                     */
                    f->more_next = more;
                    f->more_prev = more->more_prev;
                    if (more->more_prev) {
                        more->more_prev->more_next = f;
                    }
                    else {
                        /* Replacing the head of the chain,
                         * need to chnage linked list.
                         */
                        f->next = more->next;
                        f->prev = more->prev;
                        if (more->next)
                            more->next->prev = f;
                        else
                            frag_tail = f;
                        if (more->prev)
                            more->prev->next = f;
                        else
                            frag_head = f;
                        more->next = NULL;
                        more->prev = NULL;
                    }

                    more->more_prev = f;
                    break;
                }
                else if (!more->more_next) {
                    /* Append.
                     */
                    more->more_next = f;
                    f->more_prev = more;
                    break;
                }
            }

            /* Check for complete datagram. Wind back up the chain in case
             * we prepended.
             */
            while (ptr->more_prev)
                ptr = ptr->more_prev;
            more = ptr;
            assert(more);
            if (more->offset) {
                /* Missing first one.
                 */
#ifdef DEBUG
                fprintf(stderr, "FRAGDB    Missing first frag.\n");
#endif
#ifdef DEBUG
                fprintf(stderr, "FRAGDB EXIT (missing first frag)\n");
#endif
                return 0;
            }
            reached =
                (unsigned int)(ntohs(EXTRACT_IPHDR(more->sp)->total_length) -
                               IP_HDR_LEN(EXTRACT_IPHDR(more->sp)));
            //fprintf(stderr, "XXX %u\n", reached);
            for (; more; more = more->more_next) {
                assert(more);
                if ((unsigned int)(more->offset << 3) > reached) {
                    /* There has been a gap.
                     */
#ifdef DEBUG
                    fprintf(stderr,
                            "FRAGDB    Sequence gap %u %u.\n",
                            reached,
                            (unsigned int)(more->offset << 3));
#endif
#ifdef DEBUG
                    fprintf(stderr, "FRAGDB EXIT (gap)\n");
#endif
                    return 0;
                }
                reached = (unsigned int)
                    (ntohs(EXTRACT_IPHDR(more->sp)->total_length) -
                     IP_HDR_LEN(EXTRACT_IPHDR(more->sp)) +
                     (more->offset << 3));
                if (!((ntohs(EXTRACT_IPHDR(more->sp)->flags_frag)&IP_MF))) {
                    /* Reached end, all good. Make the datagram.
                     */
                    u_char *dgsp;                    
                    struct pcap_pkthdr *dgh;
                    unsigned int dglen = reached +
                        IP_HDR_LEN(EXTRACT_IPHDR(ptr->sp));
#ifdef DEBUG
                    fprintf(stderr,
                            "FRAGDB    Got datagram length %u.\n",
                            dglen);
#endif
                    dgsp = (u_char *)malloc(dglen + LL_HDR_LEN);
                    if (!dgsp) {
                        fprintf(stderr,
                                "Could not malloc for reassembled datagram "
                                "length %u.\n",
                                dglen + LL_HDR_LEN);
                        return -1;
                    }
                    assert(dgsp);
                    memset(dgsp, 0, dglen + LL_HDR_LEN);
                    reached = 0;
                    for (more = ptr; more; more = more->more_next) {
                        assert(more);
                        if (!more->offset) {
                            /* First frag, copy header too.
                             */
                            assert(more->sp);
                            assert((ntohs(EXTRACT_IPHDR(more->sp)->total_length)
                                    + LL_HDR_LEN) <= more->pkthdr.caplen);
                            memcpy(dgsp,
                                   more->sp,
                                   ntohs(EXTRACT_IPHDR(more->sp)
                                         ->total_length) + LL_HDR_LEN);
                            
                        }
                        else {
                            /* Just copy payload
                             */
                            unsigned int destoffset = (more->offset << 3) +
                                IP_HDR_LEN(EXTRACT_IPHDR(dgsp)) + LL_HDR_LEN;
                            unsigned int copylen =
                                ntohs(EXTRACT_IPHDR(more->sp)->total_length) -
                                IP_HDR_LEN(EXTRACT_IPHDR(more->sp));
                            unsigned int srcoffset = 
                                IP_HDR_LEN(EXTRACT_IPHDR(more->sp)) +
                                LL_HDR_LEN;
                            assert(more->sp);
                            if ((destoffset+copylen)>(dglen+LL_HDR_LEN)) {
                                fprintf(stderr, "FRAGDB: copylen error "
                                        "%u %u %u %u\n", destoffset,
                                        copylen, dglen, LL_HDR_LEN);
                            }
                            else {
                                assert((srcoffset+copylen)<=more->pkthdr.caplen);
                                assert((destoffset+copylen)<=(dglen+LL_HDR_LEN));
                                memcpy(&dgsp[destoffset],
                                       &more->sp[srcoffset],
                                       copylen);
                            }
                        }
                    }

                    /* Make new pkthdr and edit IP header.
                     */
                    EXTRACT_IPHDR(dgsp)->total_length = htons(dglen);
                    EXTRACT_IPHDR(dgsp)->flags_frag &=
                        ~htons(IP_OFFSET | IP_MF);
                    dgh = (struct pcap_pkthdr *)
                        malloc(sizeof(struct pcap_pkthdr));
                    if (!dgh) {
                        fprintf(stderr, "Could not malloc for pkthdr.\n");
                        free(dgsp);
                        return -1;
                    }
                    assert(dgh);
                    dgh->caplen = dglen + LL_HDR_LEN;
                    dgh->len = dglen + LL_HDR_LEN;
                    assert(ptr);
                    assert(&ptr->pkthdr.ts);
                    memcpy(&dgh->ts, &ptr->pkthdr.ts, sizeof(struct timeval));
                    EXTRACT_IPHDR(dgsp)->hdr_chksum = 0;
                    EXTRACT_IPHDR(dgsp)->hdr_chksum =
                        ip_fast_csum((unsigned char *)EXTRACT_IPHDR(dgsp),
                                     IP_HDR_LEN(EXTRACT_IPHDR(dgsp))>>2);
                    //fprintf(stderr,
                    //      "XXX %04x\n",
                    //      (unsigned int)EXTRACT_IPHDR(dgsp)->hdr_chksum);
                    /* Remove this chain.
                     */
                    if (ptr->next)
                        ptr->next->prev = ptr->prev;
                    else
                        frag_tail = ptr->prev;
                    if (ptr->prev)
                        ptr->prev->next = ptr->next;
                    else
                        frag_head = ptr->next;
                    
                    for (more = ptr; more;) {
                        assert(more);
                        ptr = more->more_next;
                        
                        if (more->sp)
                            free(more->sp);
                        free(more);
                        
                        more = ptr;
                    }

                    /* Handle the packet
                     */
#ifdef DEBUG
                    fprintf(stderr, "FRAGDB    Handle datagram...)\n");
#endif
                    assert(dgh);
                    assert(dgsp);
                    cksumck_callback((u_char *)REASSEMBLED_PACKET,
                                     dgh, dgsp);
                    free(dgh);
                    free(dgsp);
                    
#ifdef DEBUG
                    fprintf(stderr, "FRAGDB EXIT (handle frag)\n");
#endif
                    return 0;
                }
            }
#ifdef DEBUG
            fprintf(stderr,
                    "FRAGDB    Waiting for final frag, reached %u.\n",
                    reached);
#endif

            
#ifdef DEBUG
            fprintf(stderr, "FRAGDB EXIT (waiting)\n");
#endif
            return 0;
        }
    }

    if (!ptr) {
        /* New datagram
         */
        f->prev = NULL;
        f->next = frag_head;
        if (frag_head)
            frag_head->prev = f;            
        else
            frag_tail = f;
        frag_head = f;
#ifdef DEBUG
        fprintf(stderr, "FRAGDB EXIT (new datagram)\n");
#endif
        return 0;
    }
   

#ifdef DEBUG
    fprintf(stderr, "FRAGDB EXIT (end of function)\n");
#endif
    return 0;
}

/* End of $RCSfile: frag.c,v $ */
