/*
 * Copyright (c) 2022 Fastly, Kazuho Oku
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 */
#include <stdlib.h>
#include "picotls.h"

#define HPKE_V1_LABEL "HPKE-v1"

static int build_suite_id(ptls_buffer_t *buf, ptls_hpke_kem_t *kem, ptls_hpke_cipher_suite_t *cipher)
{
    int ret;

    if (cipher == NULL) {
        ptls_buffer_pushv(buf, "KEM", 3);
        ptls_buffer_push16(buf, kem->id);
    } else {
        ptls_buffer_pushv(buf, "HPKE", 4);
        ptls_buffer_push16(buf, kem->id);
        ptls_buffer_push16(buf, cipher->id.kdf);
        ptls_buffer_push16(buf, cipher->id.aead);
    }

    ret = 0;

Exit:
    return ret;
}

static int labeled_extract(ptls_hpke_kem_t *kem, ptls_hpke_cipher_suite_t *cipher, void *output, ptls_iovec_t salt,
                           const char *label, ptls_iovec_t ikm)
{
    ptls_buffer_t labeled_ikm;
    uint8_t labeled_ikm_smallbuf[64];
    int ret;

    ptls_buffer_init(&labeled_ikm, labeled_ikm_smallbuf, sizeof(labeled_ikm_smallbuf));

    ptls_buffer_pushv(&labeled_ikm, HPKE_V1_LABEL, strlen(HPKE_V1_LABEL));
    if ((ret = build_suite_id(&labeled_ikm, kem, cipher)) != 0)
        goto Exit;
    ptls_buffer_pushv(&labeled_ikm, label, strlen(label));
    ptls_buffer_pushv(&labeled_ikm, ikm.base, ikm.len);

    ret = ptls_hkdf_extract(cipher != NULL ? cipher->hash : kem->hash, output, salt,
                            ptls_iovec_init(labeled_ikm.base, labeled_ikm.off));

Exit:
    ptls_buffer_dispose(&labeled_ikm);
    return ret;
}

static int labeled_expand(ptls_hpke_kem_t *kem, ptls_hpke_cipher_suite_t *cipher, void *output, size_t outlen, ptls_iovec_t prk,
                          const char *label, ptls_iovec_t info)
{
    ptls_buffer_t labeled_info;
    uint8_t labeled_info_smallbuf[64];
    int ret;

    assert(outlen < UINT16_MAX);

    ptls_buffer_init(&labeled_info, labeled_info_smallbuf, sizeof(labeled_info_smallbuf));

    ptls_buffer_push16(&labeled_info, (uint16_t)outlen);
    ptls_buffer_pushv(&labeled_info, HPKE_V1_LABEL, strlen(HPKE_V1_LABEL));
    if ((ret = build_suite_id(&labeled_info, kem, cipher)) != 0)
        goto Exit;
    ptls_buffer_pushv(&labeled_info, label, strlen(label));
    ptls_buffer_pushv(&labeled_info, info.base, info.len);

    ret = ptls_hkdf_expand(cipher != NULL ? cipher->hash : kem->hash, output, outlen, prk,
                           ptls_iovec_init(labeled_info.base, labeled_info.off));

Exit:
    ptls_buffer_dispose(&labeled_info);
    return ret;
}

static int extract_and_expand(ptls_hpke_kem_t *kem, void *secret, size_t secret_len, ptls_iovec_t pk_s, ptls_iovec_t pk_r,
                              ptls_iovec_t dh)
{
    ptls_buffer_t kem_context;
    uint8_t kem_context_smallbuf[128], eae_prk[PTLS_MAX_DIGEST_SIZE];
    int ret;

    ptls_buffer_init(&kem_context, kem_context_smallbuf, sizeof(kem_context_smallbuf));

    ptls_buffer_pushv(&kem_context, pk_s.base, pk_s.len);
    ptls_buffer_pushv(&kem_context, pk_r.base, pk_r.len);

    if ((ret = labeled_extract(kem, NULL, eae_prk, ptls_iovec_init("", 0), "eae_prk", dh)) != 0)
        goto Exit;
    if ((ret = labeled_expand(kem, NULL, secret, secret_len, ptls_iovec_init(eae_prk, kem->hash->digest_size), "shared_secret",
                              ptls_iovec_init(kem_context.base, kem_context.off))) != 0)
        goto Exit;

Exit:
    ptls_buffer_dispose(&kem_context);
    ptls_clear_memory(eae_prk, sizeof(eae_prk));
    return ret;
}

static int dh_derive(ptls_hpke_kem_t *kem, void *secret, ptls_iovec_t pk_s, ptls_iovec_t pk_r, ptls_iovec_t dh)
{
    return extract_and_expand(kem, secret, kem->hash->digest_size, pk_s, pk_r, dh);
}

static int dh_encap(ptls_hpke_kem_t *kem, void *secret, ptls_iovec_t *pk_s, ptls_iovec_t pk_r)
{
    ptls_iovec_t dh = {NULL};
    int ret;

    *pk_s = ptls_iovec_init(NULL, 0);

    if ((ret = kem->keyex->exchange(kem->keyex, pk_s, &dh, pk_r)) != 0) {
        assert(pk_s->base == NULL);
        assert(dh.base == NULL);
        goto Exit;
    }

    if ((ret = dh_derive(kem, secret, *pk_s, pk_r, dh)) != 0)
        goto Exit;

Exit:
    if (dh.base != NULL) {
        ptls_clear_memory(dh.base, dh.len);
        free(dh.base);
    }
    if (ret != 0) {
        free(pk_s->base);
        *pk_s = ptls_iovec_init(NULL, 0);
    }
    return ret;
}

static int dh_decap(ptls_hpke_kem_t *kem, void *secret, ptls_key_exchange_context_t *keyex, ptls_iovec_t pk_s, ptls_iovec_t pk_r)
{
    ptls_iovec_t dh = {NULL};
    int ret;

    if ((ret = keyex->on_exchange(&keyex, 0, &dh, pk_s)) != 0) {
        assert(dh.base == NULL);
        goto Exit;
    }

    if ((ret = dh_derive(kem, secret, pk_s, pk_r, dh)) != 0)
        goto Exit;

Exit:
    if (dh.base != NULL) {
        ptls_clear_memory(dh.base, dh.len);
        free(dh.base);
    }
    return ret;
}

#include <stdio.h>

static int key_schedule(ptls_hpke_kem_t *kem, ptls_hpke_cipher_suite_t *cipher, ptls_aead_context_t **ctx, int is_enc,
                        const void *shared_secret, ptls_iovec_t info)
{
    ptls_buffer_t key_schedule_context;
    uint8_t key_schedule_context_smallbuf[128], secret[PTLS_MAX_DIGEST_SIZE], key[PTLS_MAX_SECRET_SIZE],
        base_nonce[PTLS_MAX_IV_SIZE];
    int ret;

    *ctx = NULL;

    ptls_buffer_init(&key_schedule_context, key_schedule_context_smallbuf, sizeof(key_schedule_context_smallbuf));

    /* key_schedule_context = concat(mode, LabeledExtract("", "psk_id_hash", psk_id), LabeledExtract("", "info_hash", info)) */
    ptls_buffer_push(&key_schedule_context, PTLS_HPKE_MODE_BASE);
    if ((ret = ptls_buffer_reserve(&key_schedule_context, cipher->hash->digest_size)) != 0 ||
        (ret = labeled_extract(kem, cipher, key_schedule_context.base + key_schedule_context.off, ptls_iovec_init(NULL, 0),
                               "psk_id_hash", ptls_iovec_init(NULL, 0))) != 0)
        goto Exit;
    key_schedule_context.off += cipher->hash->digest_size;
    if ((ret = ptls_buffer_reserve(&key_schedule_context, cipher->hash->digest_size)) != 0 ||
        (ret = labeled_extract(kem, cipher, key_schedule_context.base + key_schedule_context.off, ptls_iovec_init(NULL, 0),
                               "info_hash", info)) != 0)
        goto Exit;
    key_schedule_context.off += cipher->hash->digest_size;

    /* secret = LabeledExtract(shared_secret, "secret", psk) */
    if ((ret = labeled_extract(kem, cipher, secret, ptls_iovec_init(shared_secret, kem->hash->digest_size), "secret",
                               ptls_iovec_init("", 0))) != 0)
        goto Exit;

    /* key, base_nonce */
    if ((ret = labeled_expand(kem, cipher, key, cipher->aead->key_size, ptls_iovec_init(secret, cipher->hash->digest_size), "key",
                              ptls_iovec_init(key_schedule_context.base, key_schedule_context.off))) != 0)
        goto Exit;
    if ((ret = labeled_expand(kem, cipher, base_nonce, cipher->aead->iv_size, ptls_iovec_init(secret, cipher->hash->digest_size),
                              "base_nonce", ptls_iovec_init(key_schedule_context.base, key_schedule_context.off))) != 0)
        goto Exit;

    *ctx = ptls_aead_new_direct(cipher->aead, is_enc, key, base_nonce);

Exit:
    ptls_buffer_dispose(&key_schedule_context);
    ptls_clear_memory(secret, sizeof(secret));
    ptls_clear_memory(key, sizeof(key));
    ptls_clear_memory(base_nonce, sizeof(base_nonce));
    return ret;
}

int ptls_hpke_setup_base_s(ptls_hpke_kem_t *kem, ptls_hpke_cipher_suite_t *cipher, ptls_iovec_t *pk_s, ptls_aead_context_t **ctx,
                           ptls_iovec_t pk_r, ptls_iovec_t info)
{
    uint8_t secret[PTLS_MAX_DIGEST_SIZE];
    int ret;

    *pk_s = ptls_iovec_init(NULL, 0);

    if ((ret = dh_encap(kem, secret, pk_s, pk_r)) != 0)
        goto Exit;

    if ((ret = key_schedule(kem, cipher, ctx, 1, secret, info)) != 0)
        goto Exit;

Exit:
    if (ret != 0 && pk_s->len != 0) {
        ptls_clear_memory(pk_s->base, pk_s->len);
        free(pk_s->base);
        *pk_s = ptls_iovec_init(NULL, 0);
    }
    ptls_clear_memory(secret, sizeof(secret));
    return ret;
}

int ptls_hpke_setup_base_r(ptls_hpke_kem_t *kem, ptls_hpke_cipher_suite_t *cipher, ptls_key_exchange_context_t *keyex,
                           ptls_aead_context_t **ctx, ptls_iovec_t pk_s, ptls_iovec_t info)
{
    uint8_t secret[PTLS_MAX_DIGEST_SIZE];
    int ret;

    if ((ret = dh_decap(kem, secret, keyex, pk_s, keyex->pubkey)) != 0)
        goto Exit;

    if ((ret = key_schedule(kem, cipher, ctx, 0, secret, info)) != 0)
        goto Exit;

Exit:
    ptls_clear_memory(secret, sizeof(secret));
    return ret;
}