#include "os.h"
#include <libsec.h>

static void
ccpolyotk(Chachastate *cs, DigestState *ds)
{
       uchar otk[ChachaBsize];

       memset(ds, 0, sizeof(*ds));
       memset(otk, 0, 32);
       chacha_setblock(cs, 0);
       chacha_encrypt(otk, ChachaBsize, cs);
       poly1305(nil, 0, otk, 32, nil, ds);
}

static void
ccpolypad(uchar *buf, ulong nbuf, DigestState *ds)
{
       static uchar zeros[16] = {0};
       ulong npad;

       if(nbuf == 0)
               return;
       poly1305(buf, nbuf, nil, 0, nil, ds);
       npad = nbuf % 16;
       if(npad == 0)
               return;
       poly1305(zeros, 16 - npad, nil, 0, nil, ds);
}

static void
ccpolylen(ulong n, uchar tag[16], DigestState *ds)
{
       uchar info[8];

       info[0] = n;
       info[1] = n>>8;
       info[2] = n>>16;
       info[3] = n>>24;
       info[4] = 0;
       info[5] = 0;
       info[6] = 0;
       info[7] = 0;
       poly1305(info, 8, nil, 0, tag, ds);
}

void
ccpoly_encrypt(uchar *dat, ulong ndat, uchar *aad, ulong naad, uchar tag[16], Chachastate *cs)
{
       DigestState ds;

       ccpolyotk(cs, &ds);
       if(cs->ivwords == 2){
               poly1305(aad, naad, nil, 0, nil, &ds);
               ccpolylen(naad, nil, &ds);
               chacha_encrypt(dat, ndat, cs);
               poly1305(dat, ndat, nil, 0, nil, &ds);
               ccpolylen(ndat, tag, &ds);
       } else {
               ccpolypad(aad, naad, &ds);
               chacha_encrypt(dat, ndat, cs);
               ccpolypad(dat, ndat, &ds);
               ccpolylen(naad, nil, &ds);
               ccpolylen(ndat, tag, &ds);
       }
}

int
ccpoly_decrypt(uchar *dat, ulong ndat, uchar *aad, ulong naad, uchar tag[16], Chachastate *cs)
{
       DigestState ds;
       uchar tmp[16];

       ccpolyotk(cs, &ds);
       if(cs->ivwords == 2){
               poly1305(aad, naad, nil, 0, nil, &ds);
               ccpolylen(naad, nil, &ds);
               poly1305(dat, ndat, nil, 0, nil, &ds);
               ccpolylen(ndat, tmp, &ds);
       } else {
               ccpolypad(aad, naad, &ds);
               ccpolypad(dat, ndat, &ds);
               ccpolylen(naad, nil, &ds);
               ccpolylen(ndat, tmp, &ds);
       }
       if(tsmemcmp(tag, tmp, 16) != 0)
               return -1;
       chacha_encrypt(dat, ndat, cs);
       return 0;
}