/*      $OpenBSD: bn_isqrt.c,v 1.4 2023/08/03 18:53:56 tb Exp $ */
/*
 * Copyright (c) 2022 Theo Buehler <tb@openbsd.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <err.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <openssl/bn.h>

#include "bn_local.h"

#define N_TESTS         100

/* Sample squares between 2^128 and 2^4096. */
#define LOWER_BITS      128
#define UPPER_BITS      4096

extern const uint8_t is_square_mod_11[];
extern const uint8_t is_square_mod_63[];
extern const uint8_t is_square_mod_64[];
extern const uint8_t is_square_mod_65[];

static void
hexdump(const unsigned char *buf, size_t len)
{
        size_t i;

        for (i = 1; i <= len; i++)
                fprintf(stderr, " 0x%02hhx,%s", buf[i - 1], i % 8 ? "" : "\n");

        if (len % 8)
                fprintf(stderr, "\n");
}

static const uint8_t *
get_table(int modulus)
{
        switch (modulus) {
        case 11:
                return is_square_mod_11;
        case 63:
                return is_square_mod_63;
        case 64:
                return is_square_mod_64;
        case 65:
                return is_square_mod_65;
        default:
                return NULL;
        }
}

static int
check_tables(int print)
{
        int fill[] = {11, 63, 64, 65};
        const uint8_t *table;
        uint8_t q[65];
        size_t i;
        int j;
        int failed = 0;

        for (i = 0; i < sizeof(fill) / sizeof(fill[0]); i++) {
                memset(q, 0, sizeof(q));

                for (j = 0; j < fill[i]; j++)
                        q[(j * j) % fill[i]] = 1;

                if ((table = get_table(fill[i])) == NULL) {
                        fprintf(stderr, "failed to get table %d\n", fill[i]);
                        failed |= 1;
                        continue;
                }

                if (memcmp(table, q, fill[i]) != 0) {
                        fprintf(stderr, "table %d does not match:\n", fill[i]);
                        fprintf(stderr, "want:\n");
                        hexdump(table, fill[i]);
                        fprintf(stderr, "got:\n");
                        hexdump(q, fill[i]);
                        failed |= 1;
                        continue;
                }

                if (!print)
                        continue;

                printf("const uint8_t is_square_mod_%d[] = {\n\t", fill[i]);
                for (j = 0; j < fill[i]; j++) {
                        const char *end = " ";

                        if (j % 16 == 15)
                                end = "\n\t";
                        if (j + 1 == fill[i])
                                end = "";

                        printf("%d,%s", q[j], end);
                }
                printf("\n};\nCTASSERT(sizeof(is_square_mod_%d) == %d);\n\n",
                    fill[i], fill[i]);
        }

        return failed;
}

static int
validate_tables(void)
{
        int fill[] = {11, 63, 64, 65};
        const uint8_t *table;
        size_t i;
        int j, k;
        int failed = 0;

        for (i = 0; i < sizeof(fill) / sizeof(fill[0]); i++) {
                if ((table = get_table(fill[i])) == NULL) {
                        fprintf(stderr, "failed to get table %d\n", fill[i]);
                        failed |= 1;
                        continue;
                }

                for (j = 0; j < fill[i]; j++) {
                        for (k = 0; k < fill[i]; k++) {
                                if (j == (k * k) % fill[i])
                                        break;
                        }

                        if (table[j] == 0 && k < fill[i]) {
                                fprintf(stderr, "%d == %d^2 (mod %d)", j, k,
                                    fill[i]);
                                failed |= 1;
                        }
                        if (table[j] == 1 && k == fill[i]) {
                                fprintf(stderr, "%d not a square (mod %d)", j,
                                    fill[i]);
                                failed |= 1;
                        }
                }
        }

        return failed;
}

/*
 * Choose a random number n of bit length between LOWER_BITS and UPPER_BITS and
 * check that n == isqrt(n^2). Random numbers n^2 <= testcase < (n + 1)^2 are
 * checked to have isqrt(testcase) == n.
 */
static int
isqrt_test(void)
{
        BN_CTX *ctx;
        BIGNUM *n, *n_sqr, *lower, *upper, *testcase, *isqrt;
        int cmp, i, is_perfect_square;
        int failed = 0;

        if ((ctx = BN_CTX_new()) == NULL)
                errx(1, "BN_CTX_new");

        BN_CTX_start(ctx);

        if ((lower = BN_CTX_get(ctx)) == NULL)
                errx(1, "lower = BN_CTX_get(ctx)");
        if ((upper = BN_CTX_get(ctx)) == NULL)
                errx(1, "upper = BN_CTX_get(ctx)");
        if ((n = BN_CTX_get(ctx)) == NULL)
                errx(1, "n = BN_CTX_get(ctx)");
        if ((n_sqr = BN_CTX_get(ctx)) == NULL)
                errx(1, "n = BN_CTX_get(ctx)");
        if ((isqrt = BN_CTX_get(ctx)) == NULL)
                errx(1, "result = BN_CTX_get(ctx)");
        if ((testcase = BN_CTX_get(ctx)) == NULL)
                errx(1, "testcase = BN_CTX_get(ctx)");

        /* lower = 2^LOWER_BITS, upper = 2^UPPER_BITS. */
        if (!BN_set_bit(lower, LOWER_BITS))
                errx(1, "BN_set_bit(lower, %d)", LOWER_BITS);
        if (!BN_set_bit(upper, UPPER_BITS))
                errx(1, "BN_set_bit(upper, %d)", UPPER_BITS);

        if (!bn_rand_in_range(n, lower, upper))
                errx(1, "bn_rand_in_range n");

        /* n_sqr = n^2 */
        if (!BN_sqr(n_sqr, n, ctx))
                errx(1, "BN_sqr");

        if (!bn_isqrt(isqrt, &is_perfect_square, n_sqr, ctx))
                errx(1, "bn_isqrt n_sqr");

        if ((cmp = BN_cmp(n, isqrt)) != 0 || !is_perfect_square) {
                fprintf(stderr, "n = ");
                BN_print_fp(stderr, n);
                fprintf(stderr, "\nn^2 is_perfect_square: %d, cmp: %d\n",
                    is_perfect_square, cmp);
                failed = 1;
        }

        /* upper = 2 * n + 1 */
        if (!BN_lshift1(upper, n))
                errx(1, "BN_lshift1(upper, n)");
        if (!BN_add_word(upper, 1))
                errx(1, "BN_sub_word(upper, 1)");

        /* upper = (n + 1)^2 = n^2 + upper */
        if (!BN_add(upper, n_sqr, upper))
                errx(1, "BN_add");

        /*
         * Check that isqrt((n + 1)^2) - 1 == n.
         */

        if (!bn_isqrt(isqrt, &is_perfect_square, upper, ctx))
                errx(1, "bn_isqrt(upper)");

        if (!BN_sub_word(isqrt, 1))
                errx(1, "BN_add_word(isqrt, 1)");

        if ((cmp = BN_cmp(n, isqrt)) != 0 || !is_perfect_square) {
                fprintf(stderr, "n = ");
                BN_print_fp(stderr, n);
                fprintf(stderr, "\n(n + 1)^2 is_perfect_square: %d, cmp: %d\n",
                    is_perfect_square, cmp);
                failed = 1;
        }

        /*
         * Test N_TESTS random numbers n^2 <= testcase < (n + 1)^2 and check
         * that their isqrt is n.
         */

        for (i = 0; i < N_TESTS; i++) {
                if (!bn_rand_in_range(testcase, n_sqr, upper))
                        errx(1, "bn_rand_in_range testcase");

                if (!bn_isqrt(isqrt, &is_perfect_square, testcase, ctx))
                        errx(1, "bn_isqrt testcase");

                if ((cmp = BN_cmp(n, isqrt)) != 0 ||
                    (is_perfect_square && BN_cmp(n_sqr, testcase) != 0)) {
                        fprintf(stderr, "n = ");
                        BN_print_fp(stderr, n);
                        fprintf(stderr, "\ntestcase = ");
                        BN_print_fp(stderr, testcase);
                        fprintf(stderr,
                            "\ntestcase is_perfect_square: %d, cmp: %d\n",
                            is_perfect_square, cmp);
                        failed = 1;
                }
        }

        /*
         * Finally check that isqrt(n^2 - 1) + 1 == n.
         */

        if (!BN_sub(testcase, n_sqr, BN_value_one()))
                errx(1, "BN_sub(testcase, n_sqr, 1)");

        if (!bn_isqrt(isqrt, &is_perfect_square, testcase, ctx))
                errx(1, "bn_isqrt(n_sqr - 1)");

        if (!BN_add_word(isqrt, 1))
                errx(1, "BN_add_word(isqrt, 1)");

        if ((cmp = BN_cmp(n, isqrt)) != 0 || is_perfect_square) {
                fprintf(stderr, "n = ");
                BN_print_fp(stderr, n);
                fprintf(stderr, "\nn_sqr - 1 is_perfect_square: %d, cmp: %d\n",
                    is_perfect_square, cmp);
                failed = 1;
        }

        BN_CTX_end(ctx);
        BN_CTX_free(ctx);

        return failed;
}

static void
usage(void)
{
        fprintf(stderr, "usage: bn_isqrt [-C]\n");
        exit(1);
}

int
main(int argc, char *argv[])
{
        size_t i;
        int ch;
        int failed = 0, print = 0;

        while ((ch = getopt(argc, argv, "C")) != -1) {
                switch (ch) {
                case 'C':
                        print = 1;
                        break;
                default:
                        usage();
                        break;
                }
        }

        if (print)
                return check_tables(1);

        for (i = 0; i < N_TESTS; i++)
                failed |= isqrt_test();

        failed |= check_tables(0);
        failed |= validate_tables();

        return failed;
}