diff --git a/lib/isc/include/isc/quota.h b/lib/isc/include/isc/quota.h index 8ca1fdf6da..6f593c0ba9 100644 --- a/lib/isc/include/isc/quota.h +++ b/lib/isc/include/isc/quota.h @@ -40,11 +40,23 @@ ISC_LANG_BEGINDECLS +/*% isc_quota_cb - quota callback structure */ +typedef struct isc_quota_cb isc_quota_cb_t; +typedef void (*isc_quota_cb_func_t)(isc_quota_t *quota, void *data); +struct isc_quota_cb { + isc_quota_cb_func_t cb_func; + void * data; + ISC_LINK(isc_quota_cb_t) link; +}; + /*% isc_quota structure */ struct isc_quota { atomic_uint_fast32_t max; atomic_uint_fast32_t used; atomic_uint_fast32_t soft; + atomic_uint_fast32_t waiting; + isc_mutex_t cblock; + ISC_LIST(isc_quota_cb_t) cbs; }; void @@ -90,41 +102,46 @@ isc_quota_getused(isc_quota_t *quota); */ isc_result_t -isc_quota_reserve(isc_quota_t *quota); +isc_quota_attach(isc_quota_t *quota, isc_quota_t **p); /*%< - * Attempt to reserve one unit of 'quota'. + * + * Attempt to reserve one unit of 'quota', and also attaches '*p' to the quota + * if successful (ISC_R_SUCCESS or ISC_R_SOFTQUOTA). * * Returns: - * \li #ISC_R_SUCCESS Success + * \li #ISC_R_SUCCESS Success + * \li #ISC_R_SOFTQUOTA Success soft quota reached + * \li #ISC_R_QUOTA Quota is full + */ + +isc_result_t +isc_quota_attach_cb(isc_quota_t *quota, isc_quota_t **p, isc_quota_cb_t *cb); +/*%< + * + * Like isc_quota_attach(), but if there's no quota left then cb->cb_func will + * be called when we are attached to quota. + * Note: It's the callee responsibility to make sure that we don't end up with + * extremely huge number of callbacks waiting - making it easy to create a + * resource exhaustion attack. For example in case of TCP listening we simply + * don't accept new connections - so the number of callbacks waiting in the + * queue is limited by listen() backlog. + * + * Returns: + * \li #ISC_R_SUCCESS Success * \li #ISC_R_SOFTQUOTA Success soft quota reached * \li #ISC_R_QUOTA Quota is full */ void -isc_quota_release(isc_quota_t *quota); +isc_quota_cb_init(isc_quota_cb_t *cb, isc_quota_cb_func_t cb_func, void *data); /*%< - * Release one unit of quota. - */ - -isc_result_t -isc_quota_attach(isc_quota_t *quota, isc_quota_t **p); -/*%< - * Like isc_quota_reserve, and also attaches '*p' to the - * quota if successful (ISC_R_SUCCESS or ISC_R_SOFTQUOTA). - */ - -isc_result_t -isc_quota_force(isc_quota_t *quota, isc_quota_t **p); -/*%< - * Like isc_quota_attach, but will attach '*p' to the quota - * even if the hard quota has been exceeded. + * Initialize isc_quota_cb_t - setup the list, set the callback and data. */ void isc_quota_detach(isc_quota_t **p); /*%< - * Like isc_quota_release, and also detaches '*p' from the - * quota. + * Release one unit of quota, and also detaches '*p' from the quota. */ ISC_LANG_ENDDECLS diff --git a/lib/isc/quota.c b/lib/isc/quota.c index 9036dd3f73..137838828f 100644 --- a/lib/isc/quota.c +++ b/lib/isc/quota.c @@ -22,14 +22,20 @@ isc_quota_init(isc_quota_t *quota, unsigned int max) { atomic_init("a->max, max); atomic_init("a->used, 0); atomic_init("a->soft, 0); + atomic_init("a->waiting, 0); + ISC_LIST_INIT(quota->cbs); + isc_mutex_init("a->cblock); } void isc_quota_destroy(isc_quota_t *quota) { INSIST(atomic_load("a->used) == 0); + INSIST(atomic_load("a->waiting) == 0); + INSIST(ISC_LIST_EMPTY(quota->cbs)); atomic_store_release("a->max, 0); atomic_store_release("a->used, 0); atomic_store_release("a->soft, 0); + isc_mutex_destroy("a->cblock); } void @@ -57,43 +63,77 @@ isc_quota_getused(isc_quota_t *quota) { return (atomic_load_relaxed("a->used)); } -isc_result_t -isc_quota_reserve(isc_quota_t *quota) { +static isc_result_t +quota_reserve(isc_quota_t *quota) { isc_result_t result; - uint32_t max = atomic_load_acquire("a->max); - uint32_t soft = atomic_load_acquire("a->soft); - uint32_t used = atomic_fetch_add_relaxed("a->used, 1); - if (max == 0 || used < max) { - if (soft == 0 || used < soft) { - result = ISC_R_SUCCESS; - } else { - result = ISC_R_SOFTQUOTA; + uint_fast32_t max = atomic_load_acquire("a->max); + uint_fast32_t soft = atomic_load_acquire("a->soft); + uint_fast32_t used = atomic_load_acquire("a->used); + do { + if (max != 0 && used >= max) { + return (ISC_R_QUOTA); } - } else { - INSIST(atomic_fetch_sub_release("a->used, 1) > 0); - result = ISC_R_QUOTA; - } + if (soft != 0 && used >= soft) { + result = ISC_R_SOFTQUOTA; + } else { + result = ISC_R_SUCCESS; + } + } while (!atomic_compare_exchange_weak_acq_rel("a->used, &used, + used + 1)); return (result); } -void -isc_quota_release(isc_quota_t *quota) { +/* Must be quota->cbslock locked */ +static void +enqueue(isc_quota_t *quota, isc_quota_cb_t *cb) { + REQUIRE(cb != NULL); + ISC_LIST_ENQUEUE(quota->cbs, cb, link); + atomic_fetch_add_release("a->waiting, 1); +} + +/* Must be quota->cbslock locked */ +static isc_quota_cb_t * +dequeue(isc_quota_t *quota) { + isc_quota_cb_t *cb = ISC_LIST_HEAD(quota->cbs); + INSIST(cb != NULL); + ISC_LIST_DEQUEUE(quota->cbs, cb, link); + atomic_fetch_sub_relaxed("a->waiting, 1); + return (cb); +} + +static void +quota_release(isc_quota_t *quota) { + /* + * This is opportunistic - we might race with a failing quota_attach_cb + * and not detect that something is waiting, but eventually someone will + * be releasing quota and will detect it, so we don't need to worry - + * and we're saving a lot by not locking cblock every time. + */ + + if (atomic_load_acquire("a->waiting) > 0) { + isc_quota_cb_t *cb = NULL; + LOCK("a->cblock); + if (atomic_load_relaxed("a->waiting) > 0) { + cb = dequeue(quota); + } + UNLOCK("a->cblock); + if (cb != NULL) { + cb->cb_func(quota, cb->data); + return; + } + } + INSIST(atomic_fetch_sub_release("a->used, 1) > 0); } static isc_result_t -doattach(isc_quota_t *quota, isc_quota_t **p, bool force) { +doattach(isc_quota_t *quota, isc_quota_t **p) { isc_result_t result; REQUIRE(p != NULL && *p == NULL); - result = isc_quota_reserve(quota); + result = quota_reserve(quota); if (result == ISC_R_SUCCESS || result == ISC_R_SOFTQUOTA) { *p = quota; - } else if (result == ISC_R_QUOTA && force) { - /* attach anyway */ - atomic_fetch_add_relaxed("a->used, 1); - *p = quota; - result = ISC_R_SUCCESS; } return (result); @@ -101,17 +141,30 @@ doattach(isc_quota_t *quota, isc_quota_t **p, bool force) { isc_result_t isc_quota_attach(isc_quota_t *quota, isc_quota_t **p) { - return (doattach(quota, p, false)); + return (isc_quota_attach_cb(quota, p, NULL)); } isc_result_t -isc_quota_force(isc_quota_t *quota, isc_quota_t **p) { - return (doattach(quota, p, true)); +isc_quota_attach_cb(isc_quota_t *quota, isc_quota_t **p, isc_quota_cb_t *cb) { + isc_result_t result = doattach(quota, p); + if (result == ISC_R_QUOTA && cb != NULL) { + LOCK("a->cblock); + enqueue(quota, cb); + UNLOCK("a->cblock); + } + return (result); +} + +void +isc_quota_cb_init(isc_quota_cb_t *cb, isc_quota_cb_func_t cb_func, void *data) { + ISC_LINK_INIT(cb, link); + cb->cb_func = cb_func; + cb->data = data; } void isc_quota_detach(isc_quota_t **p) { INSIST(p != NULL && *p != NULL); - isc_quota_release(*p); + quota_release(*p); *p = NULL; } diff --git a/lib/isc/tests/Makefile.in b/lib/isc/tests/Makefile.in index 71a058a84e..6f70b6d301 100644 --- a/lib/isc/tests/Makefile.in +++ b/lib/isc/tests/Makefile.in @@ -33,7 +33,7 @@ SRCS = isctest.c aes_test.c buffer_test.c \ counter_test.c crc64_test.c errno_test.c file_test.c hash_test.c \ heap_test.c hmac_test.c ht_test.c lex_test.c \ mem_test.c md_test.c netaddr_test.c parse_test.c pool_test.c \ - radix_test.c random_test.c \ + quota_test.c radix_test.c random_test.c \ regex_test.c result_test.c safe_test.c siphash_test.c sockaddr_test.c \ socket_test.c socket_test.c symtab_test.c task_test.c \ taskpool_test.c time_test.c timer_test.c @@ -46,7 +46,7 @@ TARGETS = aes_test@EXEEXT@ buffer_test@EXEEXT@ \ ht_test@EXEEXT@ \ lex_test@EXEEXT@ mem_test@EXEEXT@ md_test@EXEEXT@ \ netaddr_test@EXEEXT@ parse_test@EXEEXT@ pool_test@EXEEXT@ \ - radix_test@EXEEXT@ \ + quota_test@EXEEXT@ radix_test@EXEEXT@ \ random_test@EXEEXT@ regex_test@EXEEXT@ result_test@EXEEXT@ \ safe_test@EXEEXT@ siphash_test@EXEEXT@ sockaddr_test@EXEEXT@ socket_test@EXEEXT@ \ socket_test@EXEEXT@ symtab_test@EXEEXT@ task_test@EXEEXT@ \ @@ -134,6 +134,11 @@ pool_test@EXEEXT@: pool_test.@O@ isctest.@O@ ${ISCDEPLIBS} ${LDFLAGS} -o $@ pool_test.@O@ isctest.@O@ \ ${ISCLIBS} ${LIBS} +quota_test@EXEEXT@: quota_test.@O@ isctest.@O@ ${ISCDEPLIBS} + ${LIBTOOL_MODE_LINK} ${PURIFY} ${CC} ${CFLAGS} \ + ${LDFLAGS} -o $@ quota_test.@O@ isctest.@O@ \ + ${ISCLIBS} ${LIBS} + radix_test@EXEEXT@: radix_test.@O@ isctest.@O@ ${ISCDEPLIBS} ${LIBTOOL_MODE_LINK} ${PURIFY} ${CC} ${CFLAGS} \ ${LDFLAGS} -o $@ radix_test.@O@ isctest.@O@ \ diff --git a/lib/isc/tests/quota_test.c b/lib/isc/tests/quota_test.c new file mode 100644 index 0000000000..997b253fb2 --- /dev/null +++ b/lib/isc/tests/quota_test.c @@ -0,0 +1,343 @@ +/* + * Copyright (C) Internet Systems Consortium, Inc. ("ISC") + * + * This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. + * + * See the COPYRIGHT file distributed with this work for additional + * information regarding copyright ownership. + */ + +#if HAVE_CMOCKA + +#include /* IWYU pragma: keep */ +#include +#include +#include +#include +#include +#include +#include + +#define UNIT_TESTING +#include + +#include +#include +#include +#include + +static void +isc_quota_get_set_test(void **state) { + UNUSED(state); + isc_quota_t quota; + isc_quota_t *quota2 = NULL; + isc_quota_init("a, 100); + + assert_int_equal(isc_quota_getmax("a), 100); + assert_int_equal(isc_quota_getsoft("a), 0); + + isc_quota_max("a, 50); + isc_quota_soft("a, 30); + + assert_int_equal(isc_quota_getmax("a), 50); + assert_int_equal(isc_quota_getsoft("a), 30); + + assert_int_equal(isc_quota_getused("a), 0); + isc_quota_attach("a, "a2); + assert_int_equal(isc_quota_getused("a), 1); + isc_quota_detach("a2); + assert_int_equal(isc_quota_getused("a), 0); + isc_quota_destroy("a); +} + +#define add_quota(quota, quotasp, exp, attached, exp_used) \ + { \ + *quotasp = NULL; \ + isc_result_t result = isc_quota_attach(quota, quotasp); \ + assert_int_equal(result, exp); \ + if (attached) { \ + assert_ptr_equal(*quotasp, quota); \ + } else { \ + assert_null(*quotasp); \ + } \ + assert_int_equal(isc_quota_getused(quota), exp_used); \ + } + +static void +isc_quota_hard_test(void **state) { + isc_quota_t quota; + isc_quota_t *quotas[110]; + int i; + UNUSED(state); + + isc_quota_init("a, 100); + + for (i = 0; i < 100; i++) { + add_quota("a, "as[i], ISC_R_SUCCESS, true, i + 1); + } + + add_quota("a, "as[100], ISC_R_QUOTA, false, 100); + + assert_int_equal(isc_quota_getused("a), 100); + + isc_quota_detach("as[0]); + assert_null(quotas[0]); + + add_quota("a, "as[100], ISC_R_SUCCESS, true, 100); + add_quota("a, "as[101], ISC_R_QUOTA, false, 100); + + for (i = 100; i > 0; i--) { + isc_quota_detach("as[i]); + assert_null(quotas[i]); + assert_int_equal(isc_quota_getused("a), i - 1); + } + assert_int_equal(isc_quota_getused("a), 0); + isc_quota_destroy("a); +} + +static void +isc_quota_soft_test(void **state) { + isc_quota_t quota; + isc_quota_t *quotas[110]; + int i; + UNUSED(state); + + isc_quota_init("a, 100); + isc_quota_soft("a, 50); + + for (i = 0; i < 50; i++) { + add_quota("a, "as[i], ISC_R_SUCCESS, true, i + 1); + } + for (i = 50; i < 100; i++) { + add_quota("a, "as[i], ISC_R_SOFTQUOTA, true, i + 1); + } + + add_quota("a, "as[i], ISC_R_QUOTA, false, 100); + + for (i = 99; i >= 0; i--) { + isc_quota_detach("as[i]); + assert_null(quotas[i]); + assert_int_equal(isc_quota_getused("a), i); + } + assert_int_equal(isc_quota_getused("a), 0); + isc_quota_destroy("a); +} + +static atomic_uint_fast32_t cb_calls = ATOMIC_VAR_INIT(0); +static isc_quota_cb_t cbs[30]; +static isc_quota_t *qp; + +static void +callback(isc_quota_t *quota, void *data) { + int val = *(int *)data; + /* Callback is not called if we get the quota directly */ + assert_int_not_equal(val, -1); + + /* We get the proper quota pointer */ + assert_ptr_equal(quota, qp); + + /* Verify that the callbacks are called in order */ + int v = atomic_fetch_add_relaxed(&cb_calls, 1); + assert_int_equal(v, val); + + /* + * First 5 will be detached by the test function, + * for the last 5 - do a 'chain detach'. + */ + if (v >= 5) { + isc_quota_detach("a); + } +} + +static void +isc_quota_callback_test(void **state) { + isc_result_t result; + isc_quota_t quota; + isc_quota_t *quotas[30]; + qp = "a; + /* + * - 10 calls that end with SUCCESS + * - 10 calls that end with SOFTQUOTA + * - 10 callbacks + */ + int ints[] = { -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 }; + int i; + UNUSED(state); + + isc_quota_init("a, 20); + isc_quota_soft("a, 10); + + for (i = 0; i < 10; i++) { + quotas[i] = NULL; + isc_quota_cb_init(&cbs[i], callback, &ints[i]); + result = isc_quota_attach_cb("a, "as[i], &cbs[i]); + assert_int_equal(result, ISC_R_SUCCESS); + assert_ptr_equal(quotas[i], "a); + assert_int_equal(isc_quota_getused("a), i + 1); + } + for (i = 10; i < 20; i++) { + quotas[i] = NULL; + isc_quota_cb_init(&cbs[i], callback, &ints[i]); + result = isc_quota_attach_cb("a, "as[i], &cbs[i]); + assert_int_equal(result, ISC_R_SOFTQUOTA); + assert_ptr_equal(quotas[i], "a); + assert_int_equal(isc_quota_getused("a), i + 1); + } + + for (i = 20; i < 30; i++) { + quotas[i] = NULL; + isc_quota_cb_init(&cbs[i], callback, &ints[i]); + result = isc_quota_attach_cb("a, "as[i], &cbs[i]); + assert_int_equal(result, ISC_R_QUOTA); + assert_ptr_equal(quotas[i], NULL); + assert_int_equal(isc_quota_getused("a), 20); + } + assert_int_equal(atomic_load(&cb_calls), 0); + + for (i = 0; i < 5; i++) { + isc_quota_detach("as[i]); + assert_null(quotas[i]); + assert_int_equal(isc_quota_getused("a), 20); + assert_int_equal(atomic_load(&cb_calls), i + 1); + } + /* That should cause a chain reaction */ + isc_quota_detach("as[5]); + assert_int_equal(atomic_load(&cb_calls), 10); + + /* Release the quotas that we did not released in the callback */ + for (i = 0; i < 5; i++) { + isc_quota_detach("as[i]); + } + + for (i = 6; i < 20; i++) { + isc_quota_detach("as[i]); + assert_null(quotas[i]); + assert_int_equal(isc_quota_getused("a), 19 - i); + } + assert_int_equal(atomic_load(&cb_calls), 10); + + assert_int_equal(isc_quota_getused("a), 0); + isc_quota_destroy("a); +} + +/* + * Multithreaded quota callback test: + * - quota set to 100 + * - 10 threads, each trying to get 100 quotas. + * - creates a separate thread to release it after 10ms + */ + +typedef struct qthreadinfo { + atomic_uint_fast32_t direct; + atomic_uint_fast32_t callback; + isc_quota_t *quota; + isc_quota_cb_t callbacks[100]; +} qthreadinfo_t; + +static atomic_uint_fast32_t g_tnum = ATOMIC_VAR_INIT(0); +/* at most 10 * 100 quota_detach threads */ +isc_thread_t g_threads[10 * 100]; + +static void * +quota_detach(void *quotap) { + isc_quota_t *quota = (isc_quota_t *)quotap; + usleep(10000); + isc_quota_detach("a); + return ((isc_threadresult_t)0); +} + +static void +quota_callback(isc_quota_t *quota, void *data) { + qthreadinfo_t *qti = (qthreadinfo_t *)data; + atomic_fetch_add_relaxed(&qti->callback, 1); + int tnum = atomic_fetch_add_relaxed(&g_tnum, 1); + isc_thread_create(quota_detach, quota, &g_threads[tnum]); +} + +static isc_threadresult_t +quota_thread(void *qtip) { + qthreadinfo_t *qti = (qthreadinfo_t *)qtip; + for (int i = 0; i < 100; i++) { + isc_quota_cb_init(&qti->callbacks[i], quota_callback, qti); + isc_quota_t *quota = NULL; + isc_result_t result = isc_quota_attach_cb(qti->quota, "a, + &qti->callbacks[i]); + if (result == ISC_R_SUCCESS) { + atomic_fetch_add_relaxed(&qti->direct, 1); + int tnum = atomic_fetch_add_relaxed(&g_tnum, 1); + isc_thread_create(quota_detach, quota, + &g_threads[tnum]); + } + } + return ((isc_threadresult_t)0); +} + +static void +isc_quota_callback_mt_test(void **state) { + UNUSED(state); + isc_quota_t quota; + int i; + + isc_quota_init("a, 100); + static qthreadinfo_t qtis[10]; + isc_thread_t threads[10]; + for (i = 0; i < 10; i++) { + atomic_init(&qtis[i].direct, 0); + atomic_init(&qtis[i].callback, 0); + qtis[i].quota = "a; + isc_thread_create(quota_thread, &qtis[i], &threads[i]); + } + for (i = 0; i < 10; i++) { + isc_thread_join(threads[i], NULL); + } + + for (i = 0; i < (int)atomic_load(&g_tnum); i++) { + isc_thread_join(g_threads[i], NULL); + } + int direct = 0, callback = 0; + + for (i = 0; i < 10; i++) { + direct += atomic_load(&qtis[i].direct); + callback += atomic_load(&qtis[i].callback); + } + /* Total quota gained must be 10 threads * 100 tries */ + assert_int_equal(direct + callback, 10 * 100); + /* + * At least 100 must be direct, the rest is virtually random: + * - in a regular run I'm constantly getting 100:900 ratio + * - under rr - usually around ~120:880 + * - under rr -h - 1000:0 + */ + assert_true(direct >= 100); + + isc_quota_destroy("a); +} + +int +main(void) { + const struct CMUnitTest tests[] = { + cmocka_unit_test(isc_quota_get_set_test), + cmocka_unit_test(isc_quota_hard_test), + cmocka_unit_test(isc_quota_soft_test), + cmocka_unit_test(isc_quota_callback_test), + cmocka_unit_test(isc_quota_callback_mt_test), + }; + + return (cmocka_run_group_tests(tests, NULL, NULL)); +} + +#else /* HAVE_CMOCKA */ + +#include + +int +main(void) { + printf("1..0 # Skipped: cmocka not available\n"); + return (0); +} + +#endif /* if HAVE_CMOCKA */ diff --git a/lib/isc/win32/libisc.def.in b/lib/isc/win32/libisc.def.in index c3dfad5a3e..7cd09bc70f 100644 --- a/lib/isc/win32/libisc.def.in +++ b/lib/isc/win32/libisc.def.in @@ -497,16 +497,15 @@ isc_queue_dequeue isc_queue_destroy isc_queue_new isc_quota_attach +isc_quota_attach_cb +isc_quota_cb_init isc_quota_destroy isc_quota_detach -isc_quota_force isc_quota_getmax isc_quota_getsoft isc_quota_getused isc_quota_init isc_quota_max -isc_quota_release -isc_quota_reserve isc_quota_soft isc_radix_create isc_radix_destroy diff --git a/util/copyrights b/util/copyrights index edbe59451e..2087c90769 100644 --- a/util/copyrights +++ b/util/copyrights @@ -2314,6 +2314,7 @@ ./lib/isc/tests/netaddr_test.c C 2016,2018,2019,2020 ./lib/isc/tests/parse_test.c C 2012,2013,2016,2018,2019,2020 ./lib/isc/tests/pool_test.c C 2013,2016,2018,2019,2020 +./lib/isc/tests/quota_test.c C 2020 ./lib/isc/tests/radix_test.c C 2014,2016,2018,2019,2020 ./lib/isc/tests/random_test.c C 2014,2015,2016,2017,2018,2019,2020 ./lib/isc/tests/regex_test.c C 2013,2015,2016,2018,2019,2020