/*

   Unix SMB/CIFS implementation.

   PAC Glue between Samba and the KDC

   Copyright (C) Catalyst.Net Ltd 2023

   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 3 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 <http://www.gnu.org/licenses/>.

*/

#include "source4/kdc/pac-blobs.h"

#include "lib/util/debug.h"

#include "lib/util/samba_util.h"

static inline size_t *pac_blobs_get_index(struct pac_blobs *pac_blobs, size_t type)

{

  /* Ensure the type is valid. */

  SMB_ASSERT(type >= PAC_TYPE_BEGIN);

  SMB_ASSERT(type < PAC_TYPE_END);

  return &pac_blobs->type_index[type - PAC_TYPE_BEGIN];

}

static inline struct type_data *pac_blobs_get(struct pac_blobs *pac_blobs, size_t type)

{

  size_t index = *pac_blobs_get_index(pac_blobs, type);

  SMB_ASSERT(index < pac_blobs->num_types);

  return &pac_blobs->type_blobs[index];

}

krb5_error_code pac_blobs_from_krb5_pac(TALLOC_CTX *mem_ctx,

          krb5_context context,

          const krb5_const_pac pac,

          struct pac_blobs **pac_blobs)

{

  krb5_error_code code = 0;

  uint32_t *types = NULL;

  struct pac_blobs *blobs = NULL;

  size_t i;

  SMB_ASSERT(pac_blobs != NULL);

  *pac_blobs = NULL;

  blobs = talloc(mem_ctx, struct pac_blobs);

  if (blobs == NULL) {

    code = ENOMEM;

    goto out;

  }

  *blobs = (struct pac_blobs) {};

  /* Initialize the array indices. */

  for (i = 0; i < ARRAY_SIZE(blobs->type_index); ++i) {

    blobs->type_index[i] = SIZE_MAX;

  }

  code = krb5_pac_get_types(context, pac, &blobs->num_types, &types);

  if (code != 0) {

    DBG_ERR("krb5_pac_get_types failed\n");

    goto out;

  }

  blobs->type_blobs = talloc_array(blobs, struct type_data, blobs->num_types);

  if (blobs->type_blobs == NULL) {

    DBG_ERR("Out of memory\n");

    code = ENOMEM;

    goto out;

  }

  for (i = 0; i < blobs->num_types; ++i) {

    uint32_t type = types[i];

    size_t *type_index = NULL;

    blobs->type_blobs[i] = (struct type_data) {

      .type = type,

      .data = NULL,

    };

    switch (type) {

      /* PAC buffer types that we support. */

    case PAC_TYPE_LOGON_INFO:

    case PAC_TYPE_CREDENTIAL_INFO:

    case PAC_TYPE_SRV_CHECKSUM:

    case PAC_TYPE_KDC_CHECKSUM:

    case PAC_TYPE_LOGON_NAME:

    case PAC_TYPE_CONSTRAINED_DELEGATION:

    case PAC_TYPE_UPN_DNS_INFO:

    case PAC_TYPE_CLIENT_CLAIMS_INFO:

    case PAC_TYPE_DEVICE_INFO:

    case PAC_TYPE_DEVICE_CLAIMS_INFO:

    case PAC_TYPE_TICKET_CHECKSUM:

    case PAC_TYPE_ATTRIBUTES_INFO:

    case PAC_TYPE_REQUESTER_SID:

    case PAC_TYPE_FULL_CHECKSUM:

      type_index = pac_blobs_get_index(blobs, type);

      if (*type_index != SIZE_MAX) {

        DBG_WARNING("PAC buffer type[%"PRIu32"] twice\n", type);

        code = EINVAL;

        goto out;

      }

      *type_index = i;

      break;

    default:

      break;

    }

  }

  *pac_blobs = blobs;

  blobs = NULL;

out:

  SAFE_FREE(types);

  TALLOC_FREE(blobs);

  return code;

}

krb5_error_code _pac_blobs_ensure_exists(struct pac_blobs *pac_blobs,

           const uint32_t type,

           const char *name,

           const char *location,

           const char *function)

{

  if (*pac_blobs_get_index(pac_blobs, type) == SIZE_MAX) {

    DEBUGLF(DBGLVL_ERR, ("%s: %s missing\n", function, name), location, function);

    return EINVAL;

  }

  return 0;

}

krb5_error_code _pac_blobs_replace_existing(struct pac_blobs *pac_blobs,

              const uint32_t type,

              const char *name,

              const DATA_BLOB *blob,

              const char *location,

              const char *function)

{

  krb5_error_code code;

  code = _pac_blobs_ensure_exists(pac_blobs,

          type,

          name,

          location,

          function);

  if (code != 0) {

    return code;

  }

  pac_blobs_get(pac_blobs, type)->data = blob;

  return 0;

}

krb5_error_code pac_blobs_add_blob(struct pac_blobs *pac_blobs,

           const uint32_t type,

           const DATA_BLOB *blob)

{

  size_t *index = NULL;

  if (blob == NULL) {

    return 0;

  }

  index = pac_blobs_get_index(pac_blobs, type);

  if (*index == SIZE_MAX) {

    struct type_data *type_blobs = NULL;

    type_blobs = talloc_realloc(pac_blobs,

              pac_blobs->type_blobs,

              struct type_data,

              pac_blobs->num_types + 1);

    if (type_blobs == NULL) {

      DBG_ERR("Out of memory\n");

      return ENOMEM;

    }

    pac_blobs->type_blobs = type_blobs;

    *index = pac_blobs->num_types++;

  }

  *pac_blobs_get(pac_blobs, type) = (struct type_data) {

    .type = type,

    .data = blob,

  };

  return 0;

}

void pac_blobs_remove_blob(struct pac_blobs *pac_blobs,

         const uint32_t type)

{

  struct type_data *type_blobs = NULL;

  size_t found_index;

  size_t i;

  /* Get the index of this PAC buffer type. */

  found_index = *pac_blobs_get_index(pac_blobs, type);

  if (found_index == SIZE_MAX) {

    /* We don't have a PAC buffer of this type, so we're done. */

    return;

  }

  /* Since the PAC buffer is present, there will be at least one type in the array. */

  SMB_ASSERT(pac_blobs->num_types > 0);

  /* The index should be valid. */

  SMB_ASSERT(found_index < pac_blobs->num_types);

  /*

   * Even though a consistent ordering of PAC buffers is not to be relied

   * upon, we must still maintain the ordering we are given.

   */

  for (i = found_index; i < pac_blobs->num_types - 1; ++i) {

    size_t moved_type;

    /* Shift each following element backwards by one. */

    pac_blobs->type_blobs[i] = pac_blobs->type_blobs[i + 1];

    /* Mark the new position of the moved element in the index. */

    moved_type = pac_blobs->type_blobs[i].type;

    if (moved_type >= PAC_TYPE_BEGIN && moved_type < PAC_TYPE_END) {

      *pac_blobs_get_index(pac_blobs, moved_type) = i;

    }

  }

  /* Mark the removed element as no longer present. */

  *pac_blobs_get_index(pac_blobs, type) = SIZE_MAX;

  /* We do not free the removed data blob, as it may be statically allocated (e.g., a null blob). */

  /* Remove the last element from the array. */

  type_blobs = talloc_realloc(pac_blobs,

            pac_blobs->type_blobs,

            struct type_data,

            --pac_blobs->num_types);

  if (type_blobs != NULL) {

    pac_blobs->type_blobs = type_blobs;

  }

}