#include "pgsodium.h"

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_ietf_keygen);
Datum pgsodium_crypto_aead_ietf_keygen(PG_FUNCTION_ARGS) {
    size_t result_size = VARHDRSZ + crypto_aead_chacha20poly1305_IETF_KEYBYTES;
    bytea *result = _pgsodium_zalloc_bytea(result_size);
    crypto_aead_chacha20poly1305_ietf_keygen(PGSODIUM_UCHARDATA(result));
    PG_RETURN_BYTEA_P(result);
}

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_ietf_noncegen);
Datum pgsodium_crypto_aead_ietf_noncegen(PG_FUNCTION_ARGS) {
    int result_size = VARHDRSZ + crypto_aead_chacha20poly1305_IETF_NPUBBYTES;
    bytea *result = _pgsodium_zalloc_bytea(result_size);
    randombytes_buf(VARDATA(result),
                    crypto_aead_chacha20poly1305_IETF_NPUBBYTES);
    PG_RETURN_BYTEA_P(result);
}

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_ietf_encrypt);
Datum pgsodium_crypto_aead_ietf_encrypt(PG_FUNCTION_ARGS) {
    bytea *message = PG_GETARG_BYTEA_P(0);
    bytea *additional = PG_GETARG_BYTEA_P(1);
    bytea *nonce = PG_GETARG_BYTEA_P(2);
    bytea *key = PG_GETARG_BYTEA_P(3);
    unsigned long long result_size;
    bytea *result;
    ERRORIF(VARSIZE_ANY_EXHDR(nonce) !=
                crypto_aead_chacha20poly1305_IETF_NPUBBYTES,
            "%s: invalid nonce");
    ERRORIF(VARSIZE_ANY_EXHDR(key) !=
                crypto_aead_chacha20poly1305_IETF_KEYBYTES,
            "%s: invalid key");

    result_size =
        VARSIZE_ANY(message) + crypto_aead_chacha20poly1305_IETF_ABYTES;
    result = _pgsodium_zalloc_bytea(result_size);
    crypto_aead_chacha20poly1305_ietf_encrypt(PGSODIUM_UCHARDATA(result),
                                              &result_size,
                                              PGSODIUM_UCHARDATA(message),
                                              VARSIZE_ANY_EXHDR(message),
                                              PGSODIUM_UCHARDATA(additional),
                                              VARSIZE_ANY_EXHDR(additional),
                                              NULL,
                                              PGSODIUM_UCHARDATA(nonce),
                                              PGSODIUM_UCHARDATA(key));
    SET_VARSIZE(result,
                result_size + VARHDRSZ +
                    crypto_aead_chacha20poly1305_IETF_ABYTES);
    PG_RETURN_BYTEA_P(result);
}

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_ietf_decrypt);
Datum pgsodium_crypto_aead_ietf_decrypt(PG_FUNCTION_ARGS) {
    bytea *ciphertext = PG_GETARG_BYTEA_P(0);
    bytea *additional = PG_GETARG_BYTEA_P(1);
    bytea *nonce = PG_GETARG_BYTEA_P(2);
    bytea *key = PG_GETARG_BYTEA_P(3);
    size_t ciphertext_len;
    unsigned long long result_size;
    bytea *result;
    int success;

    ERRORIF(VARSIZE_ANY_EXHDR(ciphertext) <=
                crypto_aead_chacha20poly1305_IETF_ABYTES,
            "%s: invalid message");
    ERRORIF(VARSIZE_ANY_EXHDR(nonce) !=
                crypto_aead_chacha20poly1305_IETF_NPUBBYTES,
            "%s: invalid nonce");
    ERRORIF(VARSIZE_ANY_EXHDR(key) !=
                crypto_aead_chacha20poly1305_IETF_KEYBYTES,
            "%s: invalid key");

    ciphertext_len = VARSIZE_ANY_EXHDR(ciphertext) -
                     crypto_aead_chacha20poly1305_IETF_ABYTES;
    result = _pgsodium_zalloc_bytea(ciphertext_len + VARHDRSZ);

    success = crypto_aead_chacha20poly1305_ietf_decrypt(
        PGSODIUM_UCHARDATA(result),
        &result_size,
        NULL,
        PGSODIUM_UCHARDATA(ciphertext),
        ciphertext_len,
        PGSODIUM_UCHARDATA(additional),
        VARSIZE_ANY_EXHDR(additional),
        PGSODIUM_UCHARDATA(nonce),
        PGSODIUM_UCHARDATA(key));
    ERRORIF(success != 0, "%s: invalid ciphertext");
    SET_VARSIZE(result, VARHDRSZ + result_size);
    PG_RETURN_BYTEA_P(result);
}

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_ietf_encrypt_by_id);
Datum pgsodium_crypto_aead_ietf_encrypt_by_id(PG_FUNCTION_ARGS) {
    bytea *message = PG_GETARG_BYTEA_P(0);
    bytea *additional = PG_GETARG_BYTEA_P(1);
    bytea *nonce = PG_GETARG_BYTEA_P(2);
    unsigned long long key_id = PG_GETARG_INT64(3);
    bytea *context = PG_GETARG_BYTEA_P(4);
    unsigned long long result_size;
    bytea *result;
    bytea *key;
    ERRORIF(VARSIZE_ANY_EXHDR(nonce) !=
                crypto_aead_chacha20poly1305_IETF_NPUBBYTES,
            "%s: invalid nonce");
    key = pgsodium_derive_helper(
        key_id, crypto_aead_chacha20poly1305_IETF_KEYBYTES, context);
    result_size =
        VARSIZE_ANY(message) + crypto_aead_chacha20poly1305_IETF_ABYTES;
    result = _pgsodium_zalloc_bytea(result_size);
    crypto_aead_chacha20poly1305_ietf_encrypt(PGSODIUM_UCHARDATA(result),
                                              &result_size,
                                              PGSODIUM_UCHARDATA(message),
                                              VARSIZE_ANY_EXHDR(message),
                                              PGSODIUM_UCHARDATA(additional),
                                              VARSIZE_ANY_EXHDR(additional),
                                              NULL,
                                              PGSODIUM_UCHARDATA(nonce),
                                              PGSODIUM_UCHARDATA(key));
    SET_VARSIZE(result,
                VARHDRSZ + result_size +
                    crypto_aead_chacha20poly1305_IETF_ABYTES);
    PG_RETURN_BYTEA_P(result);
}

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_ietf_decrypt_by_id);
Datum pgsodium_crypto_aead_ietf_decrypt_by_id(PG_FUNCTION_ARGS) {
    bytea *ciphertext = PG_GETARG_BYTEA_P(0);
    bytea *additional = PG_GETARG_BYTEA_P(1);
    bytea *nonce = PG_GETARG_BYTEA_P(2);
    unsigned long long key_id = PG_GETARG_INT64(3);
    bytea *context = PG_GETARG_BYTEA_P(4);
    size_t ciphertext_len;
    unsigned long long result_size;
    bytea *result;
    bytea *key;
    int success;
    ERRORIF(VARSIZE_ANY_EXHDR(ciphertext) <=
                crypto_aead_chacha20poly1305_IETF_ABYTES,
            "%s: invalid message");
    ERRORIF(VARSIZE_ANY_EXHDR(nonce) !=
                crypto_aead_chacha20poly1305_IETF_NPUBBYTES,
            "%s: invalid nonce");
    key = pgsodium_derive_helper(
        key_id, crypto_aead_chacha20poly1305_IETF_KEYBYTES, context);

    ciphertext_len = VARSIZE_ANY_EXHDR(ciphertext) -
                     crypto_aead_chacha20poly1305_IETF_ABYTES;
    result = _pgsodium_zalloc_bytea(ciphertext_len);

    success = crypto_aead_chacha20poly1305_ietf_decrypt(
        PGSODIUM_UCHARDATA(result),
        &result_size,
        NULL,
        PGSODIUM_UCHARDATA(ciphertext),
        ciphertext_len,
        PGSODIUM_UCHARDATA(additional),
        VARSIZE_ANY_EXHDR(additional),
        PGSODIUM_UCHARDATA(nonce),
        PGSODIUM_UCHARDATA(key));
    ERRORIF(success != 0, "%s: invalid ciphertext");
    SET_VARSIZE(result, VARHDRSZ + result_size);
    PG_RETURN_BYTEA_P(result);
}

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_det_keygen);
Datum pgsodium_crypto_aead_det_keygen(PG_FUNCTION_ARGS) {
    size_t result_size = VARHDRSZ + crypto_aead_det_xchacha20_KEYBYTES;
    bytea *result = _pgsodium_zalloc_bytea(result_size);
    crypto_aead_det_xchacha20_keygen(PGSODIUM_UCHARDATA(result));
    PG_RETURN_BYTEA_P(result);
}

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_det_noncegen);
Datum pgsodium_crypto_aead_det_noncegen(PG_FUNCTION_ARGS) {
    int result_size = VARHDRSZ + crypto_aead_det_xchacha20_NONCEBYTES;
    bytea *result = _pgsodium_zalloc_bytea(result_size);
    randombytes_buf(VARDATA(result),
                    crypto_aead_det_xchacha20_NONCEBYTES);
    PG_RETURN_BYTEA_P(result);
}

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_det_encrypt);
Datum pgsodium_crypto_aead_det_encrypt(PG_FUNCTION_ARGS) {
    bytea *message = PG_GETARG_BYTEA_P(0);
    bytea *additional = PG_GETARG_BYTEA_P(1);
    bytea *key = PG_GETARG_BYTEA_P(2);
    bytea *nonce;
    size_t result_size;
    bytea *result;
    int success;

    ERRORIF(VARSIZE_ANY_EXHDR(key) != crypto_aead_det_xchacha20_KEYBYTES,
            "%s: invalid key");
    if (!PG_ARGISNULL(3)) {
        nonce = PG_GETARG_BYTEA_P(3);
        ERRORIF(VARSIZE_ANY_EXHDR(nonce) !=
                    crypto_aead_det_xchacha20_NONCEBYTES,
                "%s: invalid nonce");
    } else {
        nonce = NULL;
    }
    result_size = VARSIZE_ANY_EXHDR(message) + crypto_aead_det_xchacha20_ABYTES;
    result = _pgsodium_zalloc_bytea(result_size);
    success = crypto_aead_det_xchacha20_encrypt(
        PGSODIUM_UCHARDATA(result),
        PGSODIUM_UCHARDATA(message),
        VARSIZE_ANY_EXHDR(message),
        PGSODIUM_UCHARDATA(additional),
        VARSIZE_ANY_EXHDR(additional),
        nonce != NULL ? PGSODIUM_UCHARDATA(nonce) : NULL,
        PGSODIUM_UCHARDATA(key));
    ERRORIF(success != 0, "%s: crypto_aead_det_xchacha20_encrypt failed");
    SET_VARSIZE(result, VARHDRSZ + result_size);
    PG_RETURN_BYTEA_P(result);
}

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_det_decrypt);
Datum pgsodium_crypto_aead_det_decrypt(PG_FUNCTION_ARGS) {
    bytea *ciphertext = PG_GETARG_BYTEA_P(0);
    bytea *additional = PG_GETARG_BYTEA_P(1);
    bytea *key = PG_GETARG_BYTEA_P(2);
    size_t result_len;
    bytea *result, *nonce;
    int success;
    ERRORIF(VARSIZE_ANY_EXHDR(ciphertext) <= crypto_aead_det_xchacha20_ABYTES,
            "%s: invalid message");
    ERRORIF(VARSIZE_ANY_EXHDR(key) != crypto_aead_det_xchacha20_KEYBYTES,
            "%s: invalid key");
    result_len =
        VARSIZE_ANY_EXHDR(ciphertext) - crypto_aead_det_xchacha20_ABYTES;
    result = _pgsodium_zalloc_bytea(result_len);
    if (!PG_ARGISNULL(3)) {
        nonce = PG_GETARG_BYTEA_P(3);
        ERRORIF(VARSIZE_ANY_EXHDR(nonce) !=
                    crypto_aead_det_xchacha20_NONCEBYTES,
                "%s: invalid nonce");
    } else {
        nonce = NULL;
    }
    success = crypto_aead_det_xchacha20_decrypt(
        PGSODIUM_UCHARDATA(result),
        PGSODIUM_UCHARDATA(ciphertext),
        VARSIZE_ANY_EXHDR(ciphertext),
        PGSODIUM_UCHARDATA(additional),
        VARSIZE_ANY_EXHDR(additional),
        nonce != NULL ? PGSODIUM_UCHARDATA(nonce) : NULL,
        PGSODIUM_UCHARDATA(key));
    ERRORIF(success != 0, "%s: invalid ciphertext");
    SET_VARSIZE(result, VARHDRSZ + result_len);
    PG_RETURN_BYTEA_P(result);
}

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_det_encrypt_by_id);
Datum pgsodium_crypto_aead_det_encrypt_by_id(PG_FUNCTION_ARGS) {
    bytea *message = PG_GETARG_BYTEA_P(0);
    bytea *additional = PG_GETARG_BYTEA_P(1);
    unsigned long long key_id = PG_GETARG_INT64(2);
    bytea *context = PG_GETARG_BYTEA_P(3);
    bytea *key, *nonce;
    size_t result_size;
    bytea *result;
    int success;

    if (!PG_ARGISNULL(4)) {
        nonce = PG_GETARG_BYTEA_P(4);
        ERRORIF(VARSIZE_ANY_EXHDR(nonce) !=
                    crypto_aead_det_xchacha20_NONCEBYTES,
                "%s: invalid nonce");
    } else {
        nonce = NULL;
    }
    result_size = VARSIZE_ANY_EXHDR(message) + crypto_aead_det_xchacha20_ABYTES;
    result = _pgsodium_zalloc_bytea(result_size);

    key = pgsodium_derive_helper(
        key_id, crypto_aead_det_xchacha20_KEYBYTES, context);

    success = crypto_aead_det_xchacha20_encrypt(
        PGSODIUM_UCHARDATA(result),
        PGSODIUM_UCHARDATA(message),
        VARSIZE_ANY_EXHDR(message),
        PGSODIUM_UCHARDATA(additional),
        VARSIZE_ANY_EXHDR(additional),
        nonce != NULL ? PGSODIUM_UCHARDATA(nonce) : NULL,
        PGSODIUM_UCHARDATA(key));
    ERRORIF(success != 0, "%s: failed");
    SET_VARSIZE(result, VARHDRSZ + result_size);
    PG_RETURN_BYTEA_P(result);
}

PG_FUNCTION_INFO_V1(pgsodium_crypto_aead_det_decrypt_by_id);
Datum pgsodium_crypto_aead_det_decrypt_by_id(PG_FUNCTION_ARGS) {
    bytea *ciphertext = PG_GETARG_BYTEA_P(0);
    bytea *additional = PG_GETARG_BYTEA_P(1);
    unsigned long long key_id = PG_GETARG_INT64(2);
    bytea *context = PG_GETARG_BYTEA_P(3);
    size_t result_len;
    bytea *key, *result, *nonce;
    int success;
    if (!PG_ARGISNULL(4)) {
        nonce = PG_GETARG_BYTEA_P(4);
        ERRORIF(VARSIZE_ANY_EXHDR(nonce) !=
                    crypto_aead_det_xchacha20_NONCEBYTES,
                "%s: invalid nonce");
    } else {
        nonce = NULL;
    }
    ERRORIF(VARSIZE_ANY_EXHDR(ciphertext) <= crypto_aead_det_xchacha20_ABYTES,
            "%s: invalid message");
    result_len =
        VARSIZE_ANY_EXHDR(ciphertext) - crypto_aead_det_xchacha20_ABYTES;
    result = _pgsodium_zalloc_bytea(result_len);
    key = pgsodium_derive_helper(
        key_id, crypto_aead_det_xchacha20_KEYBYTES, context);

    success = crypto_aead_det_xchacha20_decrypt(
        PGSODIUM_UCHARDATA(result),
        PGSODIUM_UCHARDATA(ciphertext),
        VARSIZE_ANY_EXHDR(ciphertext),
        PGSODIUM_UCHARDATA(additional),
        VARSIZE_ANY_EXHDR(additional),
        nonce != NULL ? PGSODIUM_UCHARDATA(nonce) : NULL,
        PGSODIUM_UCHARDATA(key));
    ERRORIF(success != 0, "%s: invalid ciphertext");
    SET_VARSIZE(result, VARHDRSZ + result_len);
    PG_RETURN_BYTEA_P(result);
}