/* SPDX-License-Identifier: GPL-2.0-or-later
 *
 * Copyright (C) 2024  Herbert Valerio Riedel
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of 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.
 *
 * This program 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 a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#ifndef HS_AES_GCM_H
#define HS_AES_GCM_H

#include <openssl/evp.h>
#include <openssl/err.h>
#include <stdint.h>
#include <stddef.h>

/* return codes */
#define HS_AES_GCM_OK       0
#define HS_AES_GCM_AUTHFAIL 1

#define HS_AES_GCM_ERR      (-1)
#define HS_AES_GCM_INVAL    (-2)

typedef enum {
  HS_AES_GCM_ENC = 1,
  HS_AES_GCM_DEC = 2
} hs_aes_gcm_op_t;

static inline int
hs_aes_gcm_ctx_init(EVP_CIPHER_CTX **pctx,
                    const hs_aes_gcm_op_t op,
                    const uint8_t key[], size_t keysize,
                    const uint8_t iv[],  size_t ivsize)

{
  if (!(op == HS_AES_GCM_DEC || op == HS_AES_GCM_ENC))
    return HS_AES_GCM_INVAL;

  if (!key || (!iv || ivsize<1))
    return HS_AES_GCM_INVAL;

  const EVP_CIPHER *cipher = NULL;
  switch(keysize) {
  case 16: cipher = EVP_aes_128_gcm(); break;
  case 24: cipher = EVP_aes_192_gcm(); break;
  case 32: cipher = EVP_aes_256_gcm(); break;
  default: return HS_AES_GCM_INVAL;
  }

  int ret = HS_AES_GCM_ERR;

  EVP_CIPHER_CTX *ctx = *pctx = EVP_CIPHER_CTX_new();
  if (!ctx) goto done;

  switch(op) {
  case HS_AES_GCM_ENC:
    if (ivsize == 12) {
      if (1 != EVP_EncryptInit_ex(ctx, cipher, NULL, key, iv))
        goto done;
    } else {
      if (1 != EVP_EncryptInit_ex(ctx, cipher, NULL, key, NULL))
        goto done;

      if (1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, (int)ivsize, NULL))
        goto done;

      if (1 != EVP_EncryptInit_ex(ctx, NULL, NULL, NULL, iv))
        goto done;
    }
    break;

  case HS_AES_GCM_DEC:
    if (ivsize == 12) {
      if (1 != EVP_DecryptInit_ex(ctx, cipher, NULL, key, iv))
        goto done;
    } else {
      if (1 != EVP_DecryptInit_ex(ctx, cipher, NULL, key, NULL))
        goto done;

      if (1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, (int)ivsize, NULL))
        goto done;

      if (1 != EVP_DecryptInit_ex(ctx, NULL, NULL, NULL, iv))
        goto done;
    }
    break;

  default: /* impossible */
    ret = HS_AES_GCM_INVAL;
    goto done;
  }

  return HS_AES_GCM_OK;

 done:
  if (ctx) {
    EVP_CIPHER_CTX_free(ctx);
    *pctx = NULL;
  }
  if (ret != HS_AES_GCM_OK) ERR_clear_error();
  return ret;
}

static inline int
hs_aes_gcm_check_params(const uint8_t aad[], size_t aadsize,
                        const uint8_t tag[], size_t tagsize,
                        const uint8_t ct[],
                        const uint8_t pt[],  size_t msgsize)
{
  if (!aad && aadsize>0)
    return HS_AES_GCM_INVAL;

  if ((!pt || !ct) && msgsize>0)
    return HS_AES_GCM_INVAL;

  if (!tag || !(tagsize == 4 || tagsize == 8 || (tagsize >= 12 && tagsize <= 16)))
    return HS_AES_GCM_INVAL;

  return HS_AES_GCM_OK;
}

/* set dst = NULL for AAD */
static inline int
hs_aes_gcm_ctx_update(EVP_CIPHER_CTX *ctx,
                      const hs_aes_gcm_op_t op,
                      const uint8_t src[],
                            uint8_t dst[], const size_t msgsize)
{
  if (!(op == HS_AES_GCM_DEC || op == HS_AES_GCM_ENC))
    return HS_AES_GCM_INVAL;

  const size_t blocksize = 1024*1024; /* guaranteed to fit into an 'int' */

  for(size_t ofs = 0; ofs < msgsize; ofs += blocksize) {
    const size_t n = (ofs + blocksize <= msgsize) ? blocksize : msgsize - ofs;

    int written = 0;
    int rc = 0;

    switch(op) {
    case HS_AES_GCM_ENC:
      rc = EVP_EncryptUpdate(ctx, dst ? (dst+ofs) : dst, &written, src+ofs, (int)n);
      break;
    case HS_AES_GCM_DEC:
      rc = EVP_DecryptUpdate(ctx, dst ? (dst+ofs) : dst, &written, src+ofs, (int)n);
      break;
    }

    if (1 != rc) return HS_AES_GCM_ERR;
    if (n != written) return HS_AES_GCM_ERR;
  } /* for */

  return HS_AES_GCM_OK;
}

static inline int
hs_aes_gcm_ctx_finalize_enc(EVP_CIPHER_CTX *ctx, uint8_t tag[], const size_t tagsize)
{
  int written = 0;
  if (1 != EVP_EncryptFinal_ex(ctx, NULL, &written)) return HS_AES_GCM_ERR;
  if (written) return HS_AES_GCM_ERR;

  if (1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, (int)tagsize, (void*)tag))
    return HS_AES_GCM_ERR;

  return HS_AES_GCM_OK;
}

static inline int
hs_aes_gcm_ctx_finalize_dec(EVP_CIPHER_CTX *ctx, const uint8_t tag[], const size_t tagsize)
{
  if (1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, (int)tagsize, (void*)tag))
    return HS_AES_GCM_ERR;

  int written = 0;
  int rc = EVP_DecryptFinal_ex(ctx, NULL, &written);
  if (written) return HS_AES_GCM_ERR;

  switch(rc) {
    case 1:
      return HS_AES_GCM_OK;
    case 0:
      /* the OpenSSL documentation isn't clear about how to
         differentiate between auth failures and other internal
         errors, so we have to trust that when EVP_DecryptFinal_ex
         returns with 0 it means auth failure */
      return HS_AES_GCM_AUTHFAIL;
    default:
      return HS_AES_GCM_ERR;
    }
}

/*****************************************************************************/

static inline int
hs_aes_gcm_enc(const uint8_t key[], const size_t keysize,
               const uint8_t iv[],  const size_t ivsize,
                     uint8_t tag[], const size_t tagsize,
               const uint8_t aad[], const size_t aadsize,
               const uint8_t pt[],
                     uint8_t ct[],  const size_t msgsize)
{
  EVP_CIPHER_CTX *ctx = NULL;

  int ret = hs_aes_gcm_check_params(aad, aadsize, tag, tagsize, ct, pt, msgsize);
  if (HS_AES_GCM_OK != ret) return ret;

  ret = hs_aes_gcm_ctx_init(&ctx, HS_AES_GCM_ENC, key, keysize, iv, ivsize);
  if (HS_AES_GCM_OK != ret) return ret;

  ret = hs_aes_gcm_ctx_update(ctx, HS_AES_GCM_ENC, aad, NULL, aadsize);
  if (HS_AES_GCM_OK != ret) goto done;

  ret = hs_aes_gcm_ctx_update(ctx, HS_AES_GCM_ENC, pt, ct, msgsize);
  if (HS_AES_GCM_OK != ret) goto done;

  hs_aes_gcm_ctx_finalize_enc(ctx, tag, tagsize);

 done:
  if (ctx) EVP_CIPHER_CTX_free(ctx);
  if (ret != HS_AES_GCM_OK) ERR_clear_error();
  return ret;
}

static inline int
hs_aes_gcm_dec(const uint8_t key[], const size_t keysize,
               const uint8_t iv[],  const size_t ivsize,
               const uint8_t tag[], const size_t tagsize,
               const uint8_t aad[], const size_t aadsize,
               const uint8_t ct[],
                     uint8_t pt[],  const size_t msgsize)
{
  EVP_CIPHER_CTX *ctx = NULL;

  int ret = hs_aes_gcm_check_params(aad, aadsize, tag, tagsize, ct, pt, msgsize);
  if (HS_AES_GCM_OK != ret) return ret;

  ret = hs_aes_gcm_ctx_init(&ctx, HS_AES_GCM_DEC, key, keysize, iv, ivsize);
  if (HS_AES_GCM_OK != ret) return ret;

  ret = hs_aes_gcm_ctx_update(ctx, HS_AES_GCM_DEC, aad, NULL, aadsize);
  if (HS_AES_GCM_OK != ret) goto done;

  ret = hs_aes_gcm_ctx_update(ctx, HS_AES_GCM_DEC, ct, pt, msgsize);
  if (HS_AES_GCM_OK != ret) goto done;

  ret = hs_aes_gcm_ctx_finalize_dec(ctx, tag, tagsize);

 done:
  if (ctx) EVP_CIPHER_CTX_free(ctx);
  if (ret != HS_AES_GCM_OK) ERR_clear_error();
  return ret;
}

#endif