/* ocb.c

   OCB AEAD mode, RFC 7253

   Copyright (C) 2021 Niels Möller

   This file is part of GNU Nettle.

   GNU Nettle is free software: you can redistribute it and/or

   modify it under the terms of either:

     * the GNU Lesser General Public License as published by the Free

       Software Foundation; either version 3 of the License, or (at your

       option) any later version.

   or

     * the GNU General Public License as published by the Free

       Software Foundation; either version 2 of the License, or (at your

       option) any later version.

   or both in parallel, as here.

   GNU Nettle is distributed in the hope that it will be useful,

   but WITHOUT ANY WARRANTY; without even the implied warranty of

   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU

   General Public License for more details.

   You should have received copies of the GNU General Public License and

   the GNU Lesser General Public License along with this program.  If

   not, see http://www.gnu.org/licenses/.

*/

#if HAVE_CONFIG_H

# include "config.h"

#endif

#include <string.h>

#include "ocb.h"

#include "block-internal.h"

#include "bswap-internal.h"

#include "memops.h"

#define OCB_MAX_BLOCKS 16

/* Returns 64 bits from the concatenation (u0, u1), starting from bit offset. */

static inline uint64_t

extract(uint64_t u0, uint64_t u1, unsigned offset)

{

  if (offset == 0)

    return u0;

  u0 = bswap64_if_le(u0);

  u1 = bswap64_if_le(u1);

  return bswap64_if_le((u0 << offset) | (u1 >> (64 - offset)));

}

void

ocb_set_key (struct ocb_key *key, const void *cipher, nettle_cipher_func *f)

{

  static const union nettle_block16 zero_block;

  f (cipher, OCB_BLOCK_SIZE, key->L[0].b, zero_block.b);

  block16_mulx_be (&key->L[1], &key->L[0]);

  block16_mulx_be (&key->L[2], &key->L[1]);

}

/* Add x^k L[2], where k is the number of trailing zero bits in i. */

static void

update_offset(const struct ocb_key *key,

        union nettle_block16 *offset, size_t i)

{

  if (i & 1)

    block16_xor (offset, &key->L[2]);

  else

    {

      assert (i > 0);

      union nettle_block16 diff;

      block16_mulx_be (&diff, &key->L[2]);

      for (i >>= 1; !(i&1); i >>= 1)

  block16_mulx_be (&diff, &diff);

      block16_xor (offset, &diff);

    }

}

static void

pad_block (union nettle_block16 *block, size_t length, const uint8_t *data)

{

  memcpy (block->b, data, length);

  block->b[length] = 0x80;

  memset (block->b + length + 1, 0, OCB_BLOCK_SIZE - 1 - length);

}

void

ocb_set_nonce (struct ocb_ctx *ctx,

         const void *cipher, nettle_cipher_func *f,

         size_t tag_length,

         size_t nonce_length, const uint8_t *nonce)

{

  union nettle_block16 top;

  uint64_t stretch;

  unsigned bottom;

  assert (nonce_length < 16);

  assert (tag_length > 0);

  assert (tag_length <= 16);

  /* Bit size, or zero for tag_length == 16 */

  top.b[0] = (tag_length & 15) << 4;

  memset (top.b + 1, 0, 15 - nonce_length);

  top.b[15 - nonce_length] |= 1;

  memcpy (top.b + 16 - nonce_length, nonce, nonce_length);

  bottom = top.b[15] & 0x3f;

  top.b[15] &= 0xc0;

  f (cipher, OCB_BLOCK_SIZE, top.b, top.b);

  stretch = top.u64[0];

#if WORDS_BIGENDIAN

  stretch ^= (top.u64[0] << 8) | (top.u64[1] >> 56);

#else

  stretch ^= (top.u64[0] >> 8) | (top.u64[1] << 56);

#endif

  ctx->initial.u64[0] = extract(top.u64[0], top.u64[1], bottom);

  ctx->initial.u64[1] = extract(top.u64[1], stretch, bottom);

  ctx->sum.u64[0] = ctx->sum.u64[1] = 0;

  ctx->checksum.u64[0] = ctx->checksum.u64[1] = 0;

  ctx->data_count = ctx->message_count = 0;

}

static void

ocb_fill_n (const struct ocb_key *key,

      union nettle_block16 *offset, size_t count,

      size_t n, union nettle_block16 *o)

{

  assert (n > 0);

  union nettle_block16 *prev;

  if (count & 1)

    prev = offset;

  else

    {

      /* Do a single block to align block count. */

      count++; /* Always odd. */

      block16_xor (offset, &key->L[2]);

      block16_set (&o[0], offset);

      prev = o;

      n--; o++;

    }

  for (; n >= 2; n -= 2, o += 2)

    {

      size_t i;

      count += 2; /* Always odd. */

      /* Based on trailing zeros of ctx->message_count - 1, the

         initial shift below discards a one bit. */

      block16_mulx_be (&o[0], &key->L[2]);

      for (i = count >> 1; !(i&1); i >>= 1)

  block16_mulx_be (&o[0], &o[0]);

      block16_xor (&o[0], prev);

      block16_xor3 (&o[1], &o[0], &key->L[2]);

      prev = &o[1];

    }

  block16_set(offset, prev);

  if (n > 0)

    {

      update_offset (key, offset, ++count);

      block16_set (o, offset);

    }

}

void

ocb_update (struct ocb_ctx *ctx, const struct ocb_key *key,

      const void *cipher, nettle_cipher_func *f,

      size_t length, const uint8_t *data)

{

  union nettle_block16 block[OCB_MAX_BLOCKS];

  size_t n = length / OCB_BLOCK_SIZE;

  assert (ctx->message_count == 0);

  if (ctx->data_count == 0)

    ctx->offset.u64[0] = ctx->offset.u64[1] = 0;

  while (n > 0)

    {

      size_t size, i;

      size_t blocks = (n <= OCB_MAX_BLOCKS) ? n

  : OCB_MAX_BLOCKS - 1 + (ctx->data_count & 1);

      ocb_fill_n (key, &ctx->offset, ctx->data_count, blocks, block);

      ctx->data_count += blocks;

      size = blocks * OCB_BLOCK_SIZE;

      memxor (block[0].b, data, size);

      f (cipher, size, block[0].b, block[0].b);

      for (i = 0; i < blocks; i++)

  block16_xor(&ctx->sum, &block[i]);

      n -= blocks; data += size;

    }

  length &= 15;

  if (length > 0)

    {

      union nettle_block16 block;

      pad_block (&block, length, data);

      block16_xor (&ctx->offset, &key->L[0]);

      block16_xor (&block, &ctx->offset);

      f (cipher, OCB_BLOCK_SIZE, block.b, block.b);

      block16_xor (&ctx->sum, &block);

    }

}

static void

ocb_crypt_n (struct ocb_ctx *ctx, const struct ocb_key *key,

       const void *cipher, nettle_cipher_func *f,

       size_t n, uint8_t *dst, const uint8_t *src)

{

  union nettle_block16 o[OCB_MAX_BLOCKS], block[OCB_MAX_BLOCKS];

  size_t size;

  while (n > 0)

    {

      size_t blocks = (n <= OCB_MAX_BLOCKS) ? n

  : OCB_MAX_BLOCKS - 1 + (ctx->message_count & 1);

      ocb_fill_n (key, &ctx->offset, ctx->message_count, blocks, o);

      ctx->message_count += blocks;

      size = blocks * OCB_BLOCK_SIZE;

      memxor3 (block[0].b, o[0].b, src, size);

      f (cipher, size, block[0].b, block[0].b);

      memxor3 (dst, block[0].b, o[0].b, size);

      n -= blocks; src += size; dst += size;

    }

}

/* Rotate bytes c positions to the right, in memory order. */

#if WORDS_BIGENDIAN

# define MEM_ROTATE_RIGHT(c, s0, s1) do {       \

    uint64_t __rotate_t = ((s0) >> (8*(c))) | ((s1) << (64-8*(c))); \

    (s1) = ((s1) >> (8*(c))) | ((s0) << (64-8*(c)));      \

    (s0) = __rotate_t;              \

  } while (0)

#else

# define MEM_ROTATE_RIGHT(c, s0, s1) do {       \

    uint64_t __rotate_t = ((s0) << (8*(c))) | ((s1) >> (64-8*(c))); \

    (s1) = ((s1) << (8*(c))) | ((s0) >> (64-8*(c)));      \

    (s0) = __rotate_t;              \

  } while (0)

#endif

/* Mask for the first c bytes in memory */

#if WORDS_BIGENDIAN

# define MEM_MASK(c) (-((uint64_t) 1 << (64 - 8*(c))))

#else

# define MEM_MASK(c) (((uint64_t) 1 << (8*(c))) - 1)

#endif

/* Checksum of n complete blocks. */

static void

ocb_checksum_n (union nettle_block16 *checksum,

    size_t n, const uint8_t *src)

{

  unsigned initial;

  uint64_t edge_word = 0;

  uint64_t s0, s1;

  if (n == 1)

    {

      memxor (checksum->b, src, OCB_BLOCK_SIZE);

      return;

    }

  /* Initial unaligned bytes. */

  initial = -(uintptr_t) src & 7;

  if (initial > 0)

    {

      /* Input not 64-bit aligned. Read initial bytes. */

      unsigned i;

      /* Edge word is read in big-endian order */

      for (i = initial; i > 0; i--)

  edge_word = (edge_word << 8) + *src++;

      n--;

    }

  /* Now src is 64-bit aligned, so do 64-bit reads. */

  for (s0 = s1 = 0 ; n > 0; n--, src += OCB_BLOCK_SIZE)

    {

      s0 ^= ((const uint64_t *) src)[0];

      s1 ^= ((const uint64_t *) src)[1];

    }

  if (initial > 0)

    {

      unsigned i;

      uint64_t mask;

      s0 ^= ((const uint64_t *) src)[0];

      for (i = 8 - initial, src += 8; i > 0; i--)

  edge_word = (edge_word << 8) + *src++;

      /* Rotate [s0, s1] right initial bytes. */

      MEM_ROTATE_RIGHT(initial, s0, s1);

      /* Add in the edge bytes.  */

      mask = MEM_MASK(initial);

      edge_word = bswap64_if_le (edge_word);

      s0 ^= (edge_word & mask);

      s1 ^= (edge_word & ~mask);

    }

  checksum->u64[0] ^= s0;

  checksum->u64[1] ^= s1;

}

void

ocb_encrypt (struct ocb_ctx *ctx, const struct ocb_key *key,

       const void *cipher, nettle_cipher_func *f,

       size_t length, uint8_t *dst, const uint8_t *src)

{

  size_t n = length / OCB_BLOCK_SIZE;

  if (ctx->message_count == 0)

    ctx->offset = ctx->initial;

  if (n > 0)

    {

      ocb_checksum_n (&ctx->checksum, n, src);

      ocb_crypt_n (ctx, key, cipher, f, n, dst, src);

      length &= 15;

    }

  if (length > 0)

    {

      union nettle_block16 block;

      src += n*OCB_BLOCK_SIZE; dst += n*OCB_BLOCK_SIZE;

      pad_block (&block, length, src);

      block16_xor (&ctx->checksum, &block);

      block16_xor (&ctx->offset, &key->L[0]);

      f (cipher, OCB_BLOCK_SIZE, block.b, ctx->offset.b);

      memxor3 (dst, block.b, src, length);

      ctx->message_count++;

    }

}

void

ocb_decrypt (struct ocb_ctx *ctx, const struct ocb_key *key,

       const void *encrypt_ctx, nettle_cipher_func *encrypt,

       const void *decrypt_ctx, nettle_cipher_func *decrypt,

       size_t length, uint8_t *dst, const uint8_t *src)

{

  size_t n = length / OCB_BLOCK_SIZE;

  if (ctx->message_count == 0)

    ctx->offset = ctx->initial;

  if (n > 0)

    {

      ocb_crypt_n (ctx, key, decrypt_ctx, decrypt, n, dst, src);

      ocb_checksum_n (&ctx->checksum, n, dst);

      length &= 15;

    }

  if (length > 0)

    {

      union nettle_block16 block;

      src += n*OCB_BLOCK_SIZE; dst += n*OCB_BLOCK_SIZE;

      block16_xor (&ctx->offset, &key->L[0]);

      encrypt (encrypt_ctx, OCB_BLOCK_SIZE, block.b, ctx->offset.b);

      memxor3 (dst, block.b, src, length);

      pad_block (&block, length, dst);

      block16_xor (&ctx->checksum, &block);

      ctx->message_count++;

    }

}

void

ocb_digest (const struct ocb_ctx *ctx, const struct ocb_key *key,

      const void *cipher, nettle_cipher_func *f,

      size_t length, uint8_t *digest)

{

  union nettle_block16 block;

  assert (length <= OCB_DIGEST_SIZE);

  block16_xor3 (&block,  &key->L[1],

    (ctx->message_count > 0) ? &ctx->offset : &ctx->initial);

  block16_xor (&block, &ctx->checksum);

  f (cipher, OCB_BLOCK_SIZE, block.b, block.b);

  memxor3 (digest, block.b, ctx->sum.b, length);

}

void

ocb_encrypt_message (const struct ocb_key *key,

         const void *cipher, nettle_cipher_func *f,

         size_t nlength, const uint8_t *nonce,

         size_t alength, const uint8_t *adata,

         size_t tlength,

         size_t clength, uint8_t *dst, const uint8_t *src)

{

  struct ocb_ctx ctx;

  assert (clength >= tlength);

  ocb_set_nonce (&ctx, cipher, f, tlength, nlength, nonce);

  ocb_update (&ctx, key, cipher, f, alength, adata);

  ocb_encrypt (&ctx, key, cipher, f,  clength - tlength, dst, src);

  ocb_digest (&ctx, key, cipher, f, tlength, dst + clength - tlength);

}

int

ocb_decrypt_message (const struct ocb_key *key,

         const void *encrypt_ctx, nettle_cipher_func *encrypt,

         const void *decrypt_ctx, nettle_cipher_func *decrypt,

         size_t nlength, const uint8_t *nonce,

         size_t alength, const uint8_t *adata,

         size_t tlength,

         size_t mlength, uint8_t *dst, const uint8_t *src)

{

  struct ocb_ctx ctx;

  union nettle_block16 digest;

  ocb_set_nonce (&ctx, encrypt_ctx, encrypt, tlength, nlength, nonce);

  ocb_update (&ctx, key, encrypt_ctx, encrypt, alength, adata);

  ocb_decrypt (&ctx, key, encrypt_ctx, encrypt, decrypt_ctx, decrypt,

         mlength, dst, src);

  ocb_digest (&ctx, key, encrypt_ctx, encrypt, tlength, digest.b);

  return memeql_sec(digest.b, src + mlength, tlength);

}