/*-
 * SPDX-License-Identifier: MIT
 *
 * Copyright (C) 2015-2021 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
 * Copyright (C) 2019-2021 Matt Dunwoodie <ncon@noconroy.net>
 *
 * 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.
 */

#define T_MESSAGE_LEN 64
#define T_FAILED_ITER(test) do {                                \
        kprintf("%s %s: FAIL, iter: %d\n", __func__, test, i);  \
        goto cleanup;                                           \
} while (0)
#define T_FAILED(test) do {                             \
        kprintf("%s %s: FAIL\n", __func__, test);       \
        goto cleanup;                                   \
} while (0)
#define T_PASSED kprintf("%s: pass\n", __func__)

static const struct expected_results {
        int             result;
        uint64_t        sleep_time; /* nanoseconds */
} rl_expected[] = {
        /* [0 ... INITIATIONS_BURSTABLE-1] entries are implied zero. */
        [INITIATIONS_BURSTABLE] = { ECONNREFUSED, 0 },
        [INITIATIONS_BURSTABLE + 1] = { 0, INITIATION_COST },
        [INITIATIONS_BURSTABLE + 2] = { ECONNREFUSED, 0 },
        [INITIATIONS_BURSTABLE + 3] = { 0, INITIATION_COST * 2 },
        [INITIATIONS_BURSTABLE + 4] = { 0, 0 },
        [INITIATIONS_BURSTABLE + 5] = { ECONNREFUSED, 0 },
};

static struct ratelimit rl_test;

static bool
cookie_ratelimit_timings_test(void)
{
        struct sockaddr_in sin = { .sin_family = AF_INET };
#ifdef INET6
        struct sockaddr_in6 sin6 = { .sin6_family = AF_INET6 };
#endif
        uint64_t t;
        int i;
        bool ret = false;

        ratelimit_init(&rl_test);

        for (i = 0; i < nitems(rl_expected); i++) {
                if ((t = rl_expected[i].sleep_time) != 0) {
                        tsleep(&rl_test, 0, "rl_test",
                               (int)(t * hz / NSEC_PER_SEC));
                }

                /*
                 * The first v4 ratelimit_allow is against a constant address,
                 * and should be indifferent to the port.
                 */
                sin.sin_addr.s_addr = 0x01020304;
                sin.sin_port = karc4random();

                if (ratelimit_allow(&rl_test, sintosa(&sin))
                    != rl_expected[i].result)
                        T_FAILED_ITER("malicious v4");

                /*
                 * The second ratelimit_allow is to test that an arbitrary
                 * address is still allowed.
                 */
                sin.sin_addr.s_addr += i + 1;
                sin.sin_port = karc4random();

                if (ratelimit_allow(&rl_test, sintosa(&sin)) != 0)
                        T_FAILED_ITER("non-malicious v4");

#ifdef INET6
                /*
                 * The first v6 ratelimit_allow is against a constant address,
                 * and should be indifferent to the port.  We also mutate the
                 * lower 64 bits of the address as we want to ensure ratelimit
                 * occurs against the higher 64 bits (/64 network).
                 */
                sin6.sin6_addr.s6_addr32[0] = 0x01020304;
                sin6.sin6_addr.s6_addr32[1] = 0x05060708;
                sin6.sin6_addr.s6_addr32[2] = i;
                sin6.sin6_addr.s6_addr32[3] = i;
                sin6.sin6_port = karc4random();

                if (ratelimit_allow(&rl_test, sin6tosa(&sin6))
                    != rl_expected[i].result)
                        T_FAILED_ITER("malicious v6");

                /*
                 * Again, test that an address different to above is still
                 * allowed.
                 */
                sin6.sin6_addr.s6_addr32[0] += i + 1;
                sin6.sin6_port = karc4random();

                if (ratelimit_allow(&rl_test, sintosa(&sin)) != 0)
                        T_FAILED_ITER("non-malicious v6");
#endif
        }
        T_PASSED;
        ret = true;

cleanup:
        ratelimit_deinit(&rl_test);
        return (ret);
}

static bool
cookie_ratelimit_capacity_test(void)
{
        struct sockaddr_in sin;
        int i;
        bool ret = false;

        ratelimit_init(&rl_test);

        sin.sin_family = AF_INET;
        sin.sin_port = 1234;

        /*
         * Test that the ratelimiter has an upper bound on the number of
         * addresses to be limited.
         */
        for (i = 0; i <= RATELIMIT_SIZE_MAX; i++) {
                sin.sin_addr.s_addr = i;
                if (i == RATELIMIT_SIZE_MAX) {
                        if (ratelimit_allow(&rl_test, sintosa(&sin))
                            != ECONNREFUSED)
                                T_FAILED_ITER("reject");
                } else {
                        if (ratelimit_allow(&rl_test, sintosa(&sin)) != 0)
                                T_FAILED_ITER("allow");
                }
        }
        T_PASSED;
        ret = true;

cleanup:
        ratelimit_deinit(&rl_test);
        return (ret);
}

static bool
cookie_ratelimit_gc_test(void)
{
        struct sockaddr_in sin;
        int i;
        bool ret = false;

        ratelimit_init(&rl_test);

        sin.sin_family = AF_INET;
        sin.sin_port = 1234;

        /* Test that the garbage collect routine will run. */
        if (rl_test.rl_table_num != 0)
                T_FAILED("init not empty");

        for (i = 0; i < RATELIMIT_SIZE_MAX / 2; i++) {
                sin.sin_addr.s_addr = i;
                if (ratelimit_allow(&rl_test, sintosa(&sin)) != 0)
                        T_FAILED_ITER("insert");
        }

        if (rl_test.rl_table_num != RATELIMIT_SIZE_MAX / 2)
                T_FAILED("insert 1 not full");

        tsleep(&rl_test, 0, "rl_test", ELEMENT_TIMEOUT * hz / 2);

        for (i = 0; i < RATELIMIT_SIZE_MAX / 2; i++) {
                sin.sin_addr.s_addr = i;
                if (ratelimit_allow(&rl_test, sintosa(&sin)) != 0)
                        T_FAILED_ITER("insert");
        }

        if (rl_test.rl_table_num != RATELIMIT_SIZE_MAX / 2)
                T_FAILED("insert 2 not full");

        tsleep(&rl_test, 0, "rl_test", ELEMENT_TIMEOUT * hz * 2);

        if (rl_test.rl_table_num != 0)
                T_FAILED("gc");

        T_PASSED;
        ret = true;

cleanup:
        ratelimit_deinit(&rl_test);
        return (ret);
}

static bool
cookie_mac_test(void)
{
        struct cookie_checker *checker;
        struct cookie_maker *maker;
        struct cookie_macs cm;
        struct sockaddr_in sin;
        uint8_t nonce[COOKIE_NONCE_SIZE];
        uint8_t cookie[COOKIE_ENCRYPTED_SIZE];
        uint8_t shared[COOKIE_INPUT_SIZE];
        uint8_t message[T_MESSAGE_LEN];
        int res, i;
        bool ret = false;

        karc4random_buf(shared, COOKIE_INPUT_SIZE);
        karc4random_buf(message, T_MESSAGE_LEN);

        /* Init cookie_maker. */
        maker = cookie_maker_alloc(shared);

        checker = cookie_checker_alloc();
        cookie_checker_update(checker, shared);

        /* Create dummy sockaddr. */
        sin.sin_family = AF_INET;
        sin.sin_len = sizeof(sin);
        sin.sin_addr.s_addr = 1;
        sin.sin_port = 51820;

        /* MAC message. */
        cookie_maker_mac(maker, &cm, message, T_MESSAGE_LEN);

        /* Check we have a null mac2. */
        for (i = 0; i < sizeof(cm.mac2); i++) {
                if (cm.mac2[i] != 0)
                        T_FAILED("validate_macs_noload_mac2_zeroed");
        }

        /* Validate all bytes are checked in mac1. */
        for (i = 0; i < sizeof(cm.mac1); i++) {
                cm.mac1[i] = ~cm.mac1[i];
                if (cookie_checker_validate_macs(checker, &cm, message,
                                                 T_MESSAGE_LEN, 0,
                                                 sintosa(&sin)) != EINVAL)
                        T_FAILED("validate_macs_noload_munge");
                cm.mac1[i] = ~cm.mac1[i];
        }

        /* Check mac2 is zeroed. */
        res = 0;
        for (i = 0; i < sizeof(cm.mac2); i++)
                res |= cm.mac2[i];
        if (res != 0)
                T_FAILED("validate_macs_mac2_checkzero");


        /* Check we can successfully validate the MAC. */
        if (cookie_checker_validate_macs(checker, &cm, message, T_MESSAGE_LEN,
                                         0, sintosa(&sin)) != 0)
                T_FAILED("validate_macs_noload_normal");

        /* Check we get a EAGAIN if no mac2 and under load. */
        if (cookie_checker_validate_macs(checker, &cm, message, T_MESSAGE_LEN,
                                         1, sintosa(&sin)) != EAGAIN)
                T_FAILED("validate_macs_load_normal");

        /* Simulate a cookie message. */
        cookie_checker_create_payload(checker, &cm, nonce, cookie,
                                      sintosa(&sin));

        /* Validate all bytes are checked in cookie. */
        for (i = 0; i < sizeof(cookie); i++) {
                cookie[i] = ~cookie[i];
                if (cookie_maker_consume_payload(maker, nonce, cookie)
                    != EINVAL)
                        T_FAILED("consume_payload_munge");
                cookie[i] = ~cookie[i];
        }

        /* Check we can actually consume the payload. */
        if (cookie_maker_consume_payload(maker, nonce, cookie) != 0)
                T_FAILED("consume_payload_normal");

        /* Check replay isn't allowed. */
        if (cookie_maker_consume_payload(maker, nonce, cookie) != ETIMEDOUT)
                T_FAILED("consume_payload_normal_replay");

        /* MAC message again, with MAC2. */
        cookie_maker_mac(maker, &cm, message, T_MESSAGE_LEN);

        /* Check we added a mac2. */
        res = 0;
        for (i = 0; i < sizeof(cm.mac2); i++)
                res |= cm.mac2[i];
        if (res == 0)
                T_FAILED("validate_macs_make_mac2");

        /* Check we get OK if mac2 and under load */
        if (cookie_checker_validate_macs(checker, &cm, message, T_MESSAGE_LEN,
                                         1, sintosa(&sin)) != 0)
                T_FAILED("validate_macs_load_normal_mac2");

        /* Check we get EAGAIN if we munge the source IP. */
        sin.sin_addr.s_addr = ~sin.sin_addr.s_addr;
        if (cookie_checker_validate_macs(checker, &cm, message, T_MESSAGE_LEN,
                                         1, sintosa(&sin)) != EAGAIN)
                T_FAILED("validate_macs_load_spoofip_mac2");
        sin.sin_addr.s_addr = ~sin.sin_addr.s_addr;

        /* Check we get OK if mac2 and under load */
        if (cookie_checker_validate_macs(checker, &cm, message, T_MESSAGE_LEN,
                                         1, sintosa(&sin)) != 0)
                T_FAILED("validate_macs_load_normal_mac2_retry");

        T_PASSED;
        ret = true;

cleanup:
        cookie_checker_free(checker);
        cookie_maker_free(maker);
        return (ret);
}

bool
cookie_selftest(void)
{
        bool ret = true;

        ret &= cookie_ratelimit_timings_test();
        ret &= cookie_ratelimit_capacity_test();
        ret &= cookie_ratelimit_gc_test();
        ret &= cookie_mac_test();

        kprintf("%s: %s\n", __func__, ret ? "pass" : "FAIL");
        return (ret);
}

#undef T_MESSAGE_LEN
#undef T_FAILED_ITER
#undef T_FAILED
#undef T_PASSED