/*-*- mode:c;indent-tabs-mode:nil;c-basic-offset:4;tab-width:8;coding:utf-8 -*-│ │ vi: set et ft=c ts=4 sts=4 sw=4 fenc=utf-8 :vi │ ╞══════════════════════════════════════════════════════════════════════════════╡ │ Copyright 2023 Justine Alexandra Roberts Tunney │ │ │ │ Permission to use, copy, modify, and/or 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. │ ╚─────────────────────────────────────────────────────────────────────────────*/ #define PY_SSIZE_T_CLEAN #include "net/http/tokenbucket.h" #include "libc/atomic.h" #include "libc/calls/calls.h" #include "libc/calls/struct/sigaction.h" #include "libc/calls/struct/timespec.h" #include "libc/errno.h" #include "libc/limits.h" #include "libc/macros.internal.h" #include "libc/runtime/runtime.h" #include "libc/sock/sock.h" #include "libc/sock/struct/sockaddr.h" #include "libc/str/str.h" #include "libc/sysv/consts/af.h" #include "libc/sysv/consts/clock.h" #include "libc/sysv/consts/f.h" #include "libc/sysv/consts/sig.h" #include "libc/sysv/consts/sock.h" #include "libc/sysv/consts/timer.h" #include "net/http/tokenbucket.h" #include "third_party/python/Include/bytesobject.h" #include "third_party/python/Include/import.h" #include "third_party/python/Include/intrcheck.h" #include "third_party/python/Include/longobject.h" #include "third_party/python/Include/modsupport.h" #include "third_party/python/Include/object.h" #include "third_party/python/Include/pyerrors.h" #include "third_party/python/Include/pymacro.h" #include "libc/serialize.h" #include "third_party/python/Include/yoink.h" PYTHON_PROVIDE("tokenbucket"); PYTHON_PROVIDE("tokenbucket.acquire"); PYTHON_PROVIDE("tokenbucket.blackhole"); PYTHON_PROVIDE("tokenbucket.count"); PYTHON_PROVIDE("tokenbucket.program"); struct TokenBucket { int fd; int pid; signed char cidr; struct timespec replenish; union { atomic_schar *b; atomic_uint_fast64_t *w; }; } g_tokenbucket; PyDoc_STRVAR(tokenbucket_doc, "Token Bucket Module\n\ \n\ This module implements the token bucket algorithm, which can help\n\ keep you safe from denial of service attacks."); wontreturn static void tokenbucket_onsignal(int sig) { _Exit(0); } static void tokenbucket_atexit(void) { kill(g_tokenbucket.pid, SIGTERM); } wontreturn static void tokenbucket_replenisher(void) { struct timespec ts; signal(SIGINT, tokenbucket_onsignal); signal(SIGHUP, tokenbucket_onsignal); signal(SIGTERM, tokenbucket_onsignal); signal(SIGUSR1, SIG_IGN); signal(SIGUSR2, SIG_IGN); ts = timespec_real(); for (;;) { if (!clock_nanosleep(CLOCK_REALTIME, TIMER_ABSTIME, &ts, 0)) { ReplenishTokens(g_tokenbucket.w, (1ul << g_tokenbucket.cidr) / 8); ts = timespec_add(ts, g_tokenbucket.replenish); } } } PyDoc_STRVAR(config_doc, "config($module, replenish=1.0, cidr=24)\n\ --\n\n\ Configures global token bucket.\n\ \n\ This function launches a background process. This function can only be\n\ called once. The replenish interval is in seconds and specifies how\n\ often the background process will add a token to each bucket. The\n\ cidr specifies the network bit granularity."); static PyObject * tokenbucket_config(PyObject *self, PyObject *args) { double replenish = 1.0; unsigned cidr = 24; if (g_tokenbucket.pid > 0) { PyErr_SetString(PyExc_ValueError, "tokenbucket.config() can only be called once"); return 0; } if (!PyArg_ParseTuple(args, "|dI:config", &replenish, &cidr)) { return 0; } if (!(8 <= cidr && cidr <= 32)) { PyErr_SetString(PyExc_ValueError, "require 8 <= cidr <= 32"); return 0; } if (!(0 < replenish && replenish <= LONG_MAX)) { PyErr_SetString(PyExc_ValueError, "need 0 < replenish <= LONG_MAX"); return 0; } if (!(g_tokenbucket.b = _mapshared(1ul << cidr))) { PyErr_NoMemory(); return 0; } memset(g_tokenbucket.b, 127, 1ul << cidr); g_tokenbucket.cidr = cidr; g_tokenbucket.replenish = timespec_frommicros(replenish * 1e6); _PyImport_AcquireLock(); g_tokenbucket.pid = fork(); if (!g_tokenbucket.pid) { PyOS_AfterFork(); } else { _PyImport_ReleaseLock(); } if (g_tokenbucket.pid == -1) { munmap(g_tokenbucket.b, 1ul << cidr); return PyErr_SetFromErrno(PyExc_OSError); } atexit(tokenbucket_atexit); if (!g_tokenbucket.pid) { tokenbucket_replenisher(); __builtin_unreachable(); } Py_RETURN_NONE; } PyDoc_STRVAR(acquire_doc, "acquire($module, ip)\n\ --\n\n\ Acquires token for IP address.\n\ \n\ This removes a token from the bucket associated with `ip` and then\n\ returns the number of tokens that were in the bucket beforehand.\n\ \n\ Return values greater than zero mean a token was atomically acquired.\n\ Values less than or equal zero means the bucket is empty."); static PyObject * tokenbucket_acquire(PyObject *self, PyObject *args) { uint32_t ip; const char *ipstr; if (g_tokenbucket.pid <= 0) { PyErr_SetString(PyExc_ValueError, "tokenbucket.config() needs to be called first"); return 0; } if (!PyArg_ParseTuple(args, "s:acquire", &ipstr)) { return 0; } if ((ip = ntohl(inet_addr(ipstr))) == -1u) { PyErr_SetString(PyExc_ValueError, "bad ipv4 address"); return 0; } return PyLong_FromLong(AcquireToken(g_tokenbucket.b, ip, g_tokenbucket.cidr)); } PyDoc_STRVAR(count_doc, "count($module, ip)\n\ --\n\n\ Counts token for IP address.\n\ \n\ Return values greater than zero mean a token was atomically countd.\n\ Values less than or equal zero means the bucket is empty."); static PyObject * tokenbucket_count(PyObject *self, PyObject *args) { uint32_t ip; const char *ipstr; if (g_tokenbucket.pid <= 0) { PyErr_SetString(PyExc_ValueError, "tokenbucket.config() needs to be called first"); return 0; } if (!PyArg_ParseTuple(args, "s:count", &ipstr)) { return 0; } if ((ip = ntohl(inet_addr(ipstr))) == -1u) { PyErr_SetString(PyExc_ValueError, "bad ipv4 address"); return 0; } return PyLong_FromLong(CountTokens(g_tokenbucket.b, ip, g_tokenbucket.cidr)); } PyDoc_STRVAR(blackhole_doc, "blackhole($module, ip)\n\ --\n\n\ Blackholes token for IP address.\n\ \n\ Bans IP address by sending it to blackholed. Returns 0 on success\n\ or errno on error. To test if blackholed is running, ban 0.0.0.0."); static PyObject * tokenbucket_blackhole(PyObject *self, PyObject *args) { int fd, fd2; char buf[4]; uint32_t ip; const char *ipstr; struct sockaddr_un addr = {AF_UNIX, "/var/run/blackhole.sock"}; if (!PyArg_ParseTuple(args, "s:blackhole", &ipstr)) { return 0; } if ((ip = ntohl(inet_addr(ipstr))) == -1u) { PyErr_SetString(PyExc_ValueError, "bad ipv4 address"); return 0; } if (!(fd = g_tokenbucket.fd)) { if ((fd = socket(AF_UNIX, SOCK_DGRAM, 0)) == -1) { return PyLong_FromLong(errno); } fd2 = fcntl(fd, F_DUPFD_CLOEXEC, 10); close(fd); if (fd2 == -1) { return PyLong_FromLong(errno); } g_tokenbucket.fd = fd = fd2; } WRITE32BE(buf, ip); if (sendto(fd, buf, 4, 0, (struct sockaddr *)&addr, sizeof(addr)) == -1) { return PyLong_FromLong(errno); } return PyLong_FromLong(0); } static PyMethodDef tokenbucket_methods[] = { {"config", tokenbucket_config, METH_VARARGS, config_doc}, {"acquire", tokenbucket_acquire, METH_VARARGS, acquire_doc}, {"count", tokenbucket_count, METH_VARARGS, count_doc}, {"blackhole", tokenbucket_blackhole, METH_VARARGS, blackhole_doc}, {0}, }; static struct PyModuleDef tokenbucketmodule = { PyModuleDef_HEAD_INIT, "tokenbucket", tokenbucket_doc, -1, tokenbucket_methods, }; PyMODINIT_FUNC PyInit_tokenbucket(void) { PyObject *m; if (!(m = PyModule_Create(&tokenbucketmodule))) return 0; return !PyErr_Occurred() ? m : 0; } #ifdef __aarch64__ _Section(".rodata.pytab.1 //") #else _Section(".rodata.pytab.1") #endif const struct _inittab _PyImport_Inittab_tokenbucket = { "tokenbucket", PyInit_tokenbucket, };