Subversion Repositories planix.SVN

#include <u.h>
#include <libc.h>
#include <oventi.h>
#include "packet.h"

static Frag *fragAlloc(Packet*, int n, int pos, Frag *next);
static Frag *fragDup(Packet*, Frag*);
static void fragFree(Frag*);

static Mem *memAlloc(int, int);
static void memFree(Mem*);
static int memHead(Mem *m, uchar *rp, int n);
static int memTail(Mem *m, uchar *wp, int n);

static char EPacketSize[] = "bad packet size";
static char EPacketOffset[] = "bad packet offset";
static char EBadSize[] = "bad size";

static struct {
        Lock lk;
        Packet *packet;
        int npacket;
        Frag *frag;
        int nfrag;
        Mem *bigMem;
        int nbigMem;
        Mem *smallMem;
        int nsmallMem;
} freeList;

#define FRAGSIZE(f) ((f)->wp - (f)->rp)
#define FRAGASIZE(f) ((f)->mem->ep - (f)->mem->bp)

#define NOTFREE(p) assert((p)->size>=0)

Packet *
packetAlloc(void)
{
        Packet *p;
        
        lock(&freeList.lk);
        p = freeList.packet;
        if(p != nil)
                freeList.packet = p->next;
        else
                freeList.npacket++;
        unlock(&freeList.lk);

        if(p == nil)
                p = vtMemBrk(sizeof(Packet));
        else
                assert(p->size == -1);
        p->size = 0;
        p->asize = 0;
        p->first = nil;
        p->last = nil;
        p->next = nil;

        return p;
}

void
packetFree(Packet *p)
{
        Frag *f, *ff;

if(0)fprint(2, "packetFree %p\n", p);

        NOTFREE(p);
        p->size = -1;

        for(f=p->first; f!=nil; f=ff) {
                ff = f->next;
                fragFree(f);
        }
        p->first = nil;
        p->last = nil;

        lock(&freeList.lk);
        p->next = freeList.packet;
        freeList.packet = p;
        unlock(&freeList.lk);
}

Packet *
packetDup(Packet *p, int offset, int n)
{       
        Frag *f, *ff;
        Packet *pp;

        NOTFREE(p);
        if(offset < 0 || n < 0 || offset+n > p->size) {
                vtSetError(EBadSize);
                return nil;
        }

        pp = packetAlloc();
        if(n == 0)
                return pp;

        pp->size = n;

        /* skip offset */
        for(f=p->first; offset >= FRAGSIZE(f); f=f->next)
                offset -= FRAGSIZE(f);
        
        /* first frag */
        ff = fragDup(pp, f);
        ff->rp += offset;
        pp->first = ff;
        n -= FRAGSIZE(ff);
        pp->asize += FRAGASIZE(ff);

        /* the remaining */
        while(n > 0) {
                f = f->next;
                ff->next = fragDup(pp, f);
                ff = ff->next;
                n -= FRAGSIZE(ff);
                pp->asize += FRAGASIZE(ff);
        }
        
        /* fix up last frag: note n <= 0 */
        ff->wp += n;
        ff->next = nil;
        pp->last = ff;

        return pp;
}

Packet *
packetSplit(Packet *p, int n)
{
        Packet *pp;
        Frag *f, *ff;

        NOTFREE(p);
        if(n < 0 || n > p->size) {
                vtSetError(EPacketSize);
                return nil;
        }

        pp = packetAlloc();
        if(n == 0)
                return pp;

        pp->size = n;
        p->size -= n;
        ff = nil;
        for(f=p->first; n > 0 && n >= FRAGSIZE(f); f=f->next) {
                n -= FRAGSIZE(f);
                p->asize -= FRAGASIZE(f);
                pp->asize += FRAGASIZE(f);
                ff = f;
        }

        /* split shared frag */
        if(n > 0) {
                ff = f;
                f = fragDup(pp, ff);
                pp->asize += FRAGASIZE(ff);
                ff->next = nil;
                ff->wp = ff->rp + n;
                f->rp += n;
        }

        pp->first = p->first;
        pp->last = ff;
        p->first = f;
        return pp;
}

int
packetConsume(Packet *p, uchar *buf, int n)
{
        NOTFREE(p);
        if(buf && !packetCopy(p, buf, 0, n))
                return 0;
        return packetTrim(p, n, p->size-n);
}

int
packetTrim(Packet *p, int offset, int n)
{
        Frag *f, *ff;

        NOTFREE(p);
        if(offset < 0 || offset > p->size) {
                vtSetError(EPacketOffset);
                return 0;
        }

        if(n < 0 || offset + n > p->size) {
                vtSetError(EPacketOffset);
                return 0;
        }

        p->size = n;

        /* easy case */
        if(n == 0) {
                for(f=p->first; f != nil; f=ff) {
                        ff = f->next;
                        fragFree(f);
                }
                p->first = p->last = nil;
                p->asize = 0;
                return 1;
        }
        
        /* free before offset */
        for(f=p->first; offset >= FRAGSIZE(f); f=ff) {
                p->asize -= FRAGASIZE(f);
                offset -= FRAGSIZE(f);
                ff = f->next;
                fragFree(f);
        }

        /* adjust frag */
        f->rp += offset;
        p->first = f;

        /* skip middle */
        for(; n > 0 && n > FRAGSIZE(f); f=f->next)
                n -= FRAGSIZE(f);

        /* adjust end */
        f->wp = f->rp + n;
        p->last = f;
        ff = f->next;
        f->next = nil;

        /* free after */
        for(f=ff; f != nil; f=ff) {
                p->asize -= FRAGASIZE(f);
                ff = f->next;
                fragFree(f);
        }
        return 1;
}

uchar *
packetHeader(Packet *p, int n)
{
        Frag *f;
        Mem *m;

        NOTFREE(p);
        if(n <= 0 || n > MaxFragSize) {
                vtSetError(EPacketSize);
                return 0;
        }

        p->size += n;
        
        /* try and fix in current frag */
        f = p->first;
        if(f != nil) {
                m = f->mem;
                if(n <= f->rp - m->bp)
                if(m->ref == 1 || memHead(m, f->rp, n)) {
                        f->rp -= n;
                        return f->rp;
                }
        }

        /* add frag to front */
        f = fragAlloc(p, n, PEnd, p->first);
        p->asize += FRAGASIZE(f);
        if(p->first == nil)
                p->last = f;
        p->first = f;
        return f->rp;
}

uchar *
packetTrailer(Packet *p, int n)
{
        Mem *m;
        Frag *f;

        NOTFREE(p);
        if(n <= 0 || n > MaxFragSize) {
                vtSetError(EPacketSize);
                return 0;
        }

        p->size += n;
        
        /* try and fix in current frag */
        if(p->first != nil) {
                f = p->last;
                m = f->mem;
                if(n <= m->ep - f->wp)
                if(m->ref == 1 || memTail(m, f->wp, n)) {
                        f->wp += n;
                        return f->wp - n;
                }
        }

        /* add frag to end */
        f = fragAlloc(p, n, (p->first == nil)?PMiddle:PFront, nil);
        p->asize += FRAGASIZE(f);
        if(p->first == nil)
                p->first = f;
        else
                p->last->next = f;
        p->last = f;
        return f->rp;
}

int
packetPrefix(Packet *p, uchar *buf, int n)
{
        Frag *f;
        int nn;
        Mem *m;

        NOTFREE(p);
        if(n <= 0)
                return 1;

        p->size += n;

        /* try and fix in current frag */
        f = p->first;
        if(f != nil) {
                m = f->mem;
                nn = f->rp - m->bp;
                if(nn > n)
                        nn = n;
                if(m->ref == 1 || memHead(m, f->rp, nn)) {
                        f->rp -= nn;
                        n -= nn;
                        memmove(f->rp, buf+n, nn);
                }
        }

        while(n > 0) {
                nn = n;
                if(nn > MaxFragSize)
                        nn = MaxFragSize;
                f = fragAlloc(p, nn, PEnd, p->first);   
                p->asize += FRAGASIZE(f);
                if(p->first == nil)
                        p->last = f;
                p->first = f;
                n -= nn;
                memmove(f->rp, buf+n, nn);
        }
        return 1;
}

int
packetAppend(Packet *p, uchar *buf, int n)
{
        Frag *f;
        int nn;
        Mem *m;

        NOTFREE(p);
        if(n <= 0)
                return 1;

        p->size += n;
        /* try and fix in current frag */
        if(p->first != nil) {
                f = p->last;
                m = f->mem;
                nn = m->ep - f->wp;
                if(nn > n)
                        nn = n;
                if(m->ref == 1 || memTail(m, f->wp, nn)) {
                        memmove(f->wp, buf, nn);
                        f->wp += nn;
                        buf += nn;
                        n -= nn;
                }
        }
        
        while(n > 0) {
                nn = n;
                if(nn > MaxFragSize)
                        nn = MaxFragSize;
                f = fragAlloc(p, nn, (p->first == nil)?PMiddle:PFront, nil);
                p->asize += FRAGASIZE(f);
                if(p->first == nil)
                        p->first = f;
                else
                        p->last->next = f;
                p->last = f;
                memmove(f->rp, buf, nn);
                buf += nn;
                n -= nn;
        }
        return 1;
}

int
packetConcat(Packet *p, Packet *pp)
{
        NOTFREE(p);
        NOTFREE(pp);
        if(pp->size == 0)
                return 1;
        p->size += pp->size;
        p->asize += pp->asize;

        if(p->first != nil)
                p->last->next = pp->first;
        else
                p->first = pp->first;
        p->last = pp->last;
        pp->size = 0;
        pp->asize = 0;
        pp->first = nil;
        pp->last = nil;
        return 1;
}

uchar *
packetPeek(Packet *p, uchar *buf, int offset, int n)
{
        Frag *f;
        int nn;
        uchar *b;

        NOTFREE(p);
        if(n == 0)
                return buf;

        if(offset < 0 || offset >= p->size) {
                vtSetError(EPacketOffset);
                return 0;
        }

        if(n < 0 || offset + n > p->size) {
                vtSetError(EPacketSize);
                return 0;
        }
        
        /* skip up to offset */
        for(f=p->first; offset >= FRAGSIZE(f); f=f->next)
                offset -= FRAGSIZE(f);

        /* easy case */
        if(offset + n <= FRAGSIZE(f))
                return f->rp + offset;

        for(b=buf; n>0; n -= nn) {
                nn = FRAGSIZE(f) - offset;
                if(nn > n)
                        nn = n;
                memmove(b, f->rp+offset, nn);
                offset = 0;
                f = f->next;
                b += nn;
        }

        return buf;
}

int
packetCopy(Packet *p, uchar *buf, int offset, int n)
{
        uchar *b;

        NOTFREE(p);
        b = packetPeek(p, buf, offset, n);
        if(b == nil)
                return 0;
        if(b != buf)
                memmove(buf, b, n);
        return 1;
}

int
packetFragments(Packet *p, IOchunk *io, int nio, int offset)
{
        Frag *f;
        int size;
        IOchunk *eio;

        NOTFREE(p);
        if(p->size == 0 || nio <= 0)
                return 0;
        
        if(offset < 0 || offset > p->size) {
                vtSetError(EPacketOffset);
                return -1;
        }

        for(f=p->first; offset >= FRAGSIZE(f); f=f->next)
                offset -= FRAGSIZE(f);

        size = 0;
        eio = io + nio;
        for(; f != nil && io < eio; f=f->next) {
                io->addr = f->rp + offset;
                io->len = f->wp - (f->rp + offset);     
                offset = 0;
                size += io->len;
                io++;
        }

        return size;
}

void
packetStats(void)
{
        Packet *p;
        Frag *f;
        Mem *m;

        int np, nf, nsm, nbm;

        lock(&freeList.lk);
        np = 0;
        for(p=freeList.packet; p; p=p->next)
                np++;
        nf = 0;
        for(f=freeList.frag; f; f=f->next)
                nf++;
        nsm = 0;
        for(m=freeList.smallMem; m; m=m->next)
                nsm++;
        nbm = 0;
        for(m=freeList.bigMem; m; m=m->next)
                nbm++;
        
        fprint(2, "packet: %d/%d frag: %d/%d small mem: %d/%d big mem: %d/%d\n",
                np, freeList.npacket,
                nf, freeList.nfrag,
                nsm, freeList.nsmallMem,
                nbm, freeList.nbigMem);

        unlock(&freeList.lk);
}


int
packetSize(Packet *p)
{
        NOTFREE(p);
        if(0) {
                Frag *f;
                int size = 0;
        
                for(f=p->first; f; f=f->next)
                        size += FRAGSIZE(f);
                if(size != p->size)
                        fprint(2, "packetSize %d %d\n", size, p->size);
                assert(size == p->size);
        }
        return p->size;
}

int
packetAllocatedSize(Packet *p)
{
        NOTFREE(p);
        if(0) {
                Frag *f;
                int asize = 0;
        
                for(f=p->first; f; f=f->next)
                        asize += FRAGASIZE(f);
                if(asize != p->asize)
                        fprint(2, "packetAllocatedSize %d %d\n", asize, p->asize);
                assert(asize == p->asize);
        }
        return p->asize;
}

void
packetSha1(Packet *p, uchar sha1[VtScoreSize])
{
        Frag *f;
        VtSha1 *s;
        int size;

        NOTFREE(p);
        s = vtSha1Alloc();
        size = p->size;
        for(f=p->first; f; f=f->next) {
                vtSha1Update(s, f->rp, FRAGSIZE(f));
                size -= FRAGSIZE(f);
        }
        assert(size == 0);
        vtSha1Final(s, sha1);
        vtSha1Free(s);
}

int
packetCmp(Packet *pkt0, Packet *pkt1)
{
        Frag *f0, *f1;
        int n0, n1, x;

        NOTFREE(pkt0);
        NOTFREE(pkt1);
        f0 = pkt0->first;
        f1 = pkt1->first;

        if(f0 == nil)
                return (f1 == nil)?0:-1;
        if(f1 == nil)
                return 1;
        n0 = FRAGSIZE(f0);
        n1 = FRAGSIZE(f1);

        for(;;) {
                if(n0 < n1) {
                        x = memcmp(f0->wp - n0, f1->wp - n1, n0);
                        if(x != 0)
                                return x;
                        n1 -= n0;
                        f0 = f0->next;
                        if(f0 == nil)
                                return -1;
                        n0 = FRAGSIZE(f0);
                } else if (n0 > n1) {
                        x = memcmp(f0->wp - n0, f1->wp - n1, n1);
                        if(x != 0)
                                return x;
                        n0 -= n1;
                        f1 = f1->next;
                        if(f1 == nil)
                                return 1;
                        n1 = FRAGSIZE(f1);
                } else { /* n0 == n1 */
                        x = memcmp(f0->wp - n0, f1->wp - n1, n0);
                        if(x != 0)
                                return x;
                        f0 = f0->next;
                        f1 = f1->next;
                        if(f0 == nil)
                                return (f1 == nil)?0:-1;
                        if(f1 == nil)
                                return 1;
                        n0 = FRAGSIZE(f0);
                        n1 = FRAGSIZE(f1);
                }
        }
}
        

static Frag *
fragAlloc(Packet *p, int n, int pos, Frag *next)
{
        Frag *f, *ef;
        Mem *m;

        /* look for local frag */
        f = &p->local[0];
        ef = &p->local[NLocalFrag];
        for(;f<ef; f++) {
                if(f->state == FragLocalFree) {
                        f->state = FragLocalAlloc;
                        goto Found;
                }
        }
        lock(&freeList.lk);     
        f = freeList.frag;
        if(f != nil)
                freeList.frag = f->next;
        else
                freeList.nfrag++;
        unlock(&freeList.lk);

        if(f == nil) {
                f = vtMemBrk(sizeof(Frag));
                f->state = FragGlobal;
        }

Found:
        if(n == 0)
                return f;

        if(pos == PEnd && next == nil)
                pos = PMiddle;
        m = memAlloc(n, pos);
        f->mem = m;
        f->rp = m->rp;
        f->wp = m->wp;
        f->next = next;

        return f;
}

static Frag *
fragDup(Packet *p, Frag *f)
{
        Frag *ff;
        Mem *m;

        m = f->mem;     

        /*
         * m->rp && m->wp can be out of date when ref == 1
         * also, potentially reclaims space from previous frags
         */
        if(m->ref == 1) {
                m->rp = f->rp;
                m->wp = f->wp;
        }

        ff = fragAlloc(p, 0, 0, nil);
        *ff = *f;
        lock(&m->lk);
        m->ref++;
        unlock(&m->lk);
        return ff;
}


static void
fragFree(Frag *f)
{
        memFree(f->mem);

        if(f->state == FragLocalAlloc) {
                f->state = FragLocalFree;
                return;
        }

        lock(&freeList.lk);
        f->next = freeList.frag;
        freeList.frag = f;
        unlock(&freeList.lk);   
}

static Mem *
memAlloc(int n, int pos)
{
        Mem *m;
        int nn;

        if(n < 0 || n > MaxFragSize) {
                vtSetError(EPacketSize);
                return 0;
        }
        if(n <= SmallMemSize) {
                lock(&freeList.lk);
                m = freeList.smallMem;
                if(m != nil)
                        freeList.smallMem = m->next;
                else
                        freeList.nsmallMem++;
                unlock(&freeList.lk);
                nn = SmallMemSize;
        } else {
                lock(&freeList.lk);
                m = freeList.bigMem;
                if(m != nil)
                        freeList.bigMem = m->next;
                else
                        freeList.nbigMem++;
                unlock(&freeList.lk);
                nn = BigMemSize;
        }

        if(m == nil) {
                m = vtMemBrk(sizeof(Mem));
                m->bp = vtMemBrk(nn);
                m->ep = m->bp + nn;
        }
        assert(m->ref == 0);    
        m->ref = 1;

        switch(pos) {
        default:
                assert(0);
        case PFront:
                m->rp = m->bp;
                break;
        case PMiddle:
                /* leave a little bit at end */
                m->rp = m->ep - n - 32;
                break;
        case PEnd:
                m->rp = m->ep - n;
                break; 
        }
        /* check we did not blow it */
        if(m->rp < m->bp)
                m->rp = m->bp;
        m->wp = m->rp + n;
        assert(m->rp >= m->bp && m->wp <= m->ep);
        return m;
}

static void
memFree(Mem *m)
{
        lock(&m->lk);
        m->ref--;
        if(m->ref > 0) {
                unlock(&m->lk);
                return;
        }
        unlock(&m->lk);
        assert(m->ref == 0);

        switch(m->ep - m->bp) {
        default:
                assert(0);
        case SmallMemSize:
                lock(&freeList.lk);
                m->next = freeList.smallMem;
                freeList.smallMem = m;
                unlock(&freeList.lk);
                break;
        case BigMemSize:
                lock(&freeList.lk);
                m->next = freeList.bigMem;
                freeList.bigMem = m;
                unlock(&freeList.lk);
                break;
        }
}

static int
memHead(Mem *m, uchar *rp, int n)
{
        lock(&m->lk);
        if(m->rp != rp) {
                unlock(&m->lk);
                return 0;
        }
        m->rp -= n;
        unlock(&m->lk);
        return 1;
}

static int
memTail(Mem *m, uchar *wp, int n)
{
        lock(&m->lk);
        if(m->wp != wp) {
                unlock(&m->lk);
                return 0;
        }
        m->wp += n;
        unlock(&m->lk);
        return 1;
}