mirror of
https://github.com/jart/cosmopolitan.git
synced 2025-10-28 11:52:56 +00:00
Add tokenbucket module to python.com
This commit is contained in:
Justine Tunney
parent
e6b7c16a53
commit
2a1c588826
5 changed files with 248 additions and 2 deletions
242
third_party/python/Modules/tokenbucket.c
vendored
Normal file
242
third_party/python/Modules/tokenbucket.c
vendored
Normal file
|
|
@ -0,0 +1,242 @@
|
|||
/*-*- mode:c;indent-tabs-mode:nil;c-basic-offset:4;tab-width:8;coding:utf-8 -*-│
|
||||
│vi: set net ft=c ts=4 sts=4 sw=4 fenc=utf-8 :vi│
|
||||
╞══════════════════════════════════════════════════════════════════════════════╡
|
||||
│ Copyright 2021 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/str/str.h"
|
||||
#include "libc/sysv/consts/clock.h"
|
||||
#include "libc/sysv/consts/sig.h"
|
||||
#include "libc/sysv/consts/timer.h"
|
||||
#include "net/http/tokenbucket.h"
|
||||
#include "third_party/libcxx/math.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 "third_party/python/Include/yoink.h"
|
||||
// clang-format off
|
||||
|
||||
PYTHON_PROVIDE("tokenbucket");
|
||||
PYTHON_PROVIDE("tokenbucket.program");
|
||||
PYTHON_PROVIDE("tokenbucket.acquire");
|
||||
PYTHON_PROVIDE("tokenbucket.count");
|
||||
|
||||
struct TokenBucket {
|
||||
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));
|
||||
}
|
||||
|
||||
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},
|
||||
{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;
|
||||
}
|
||||
|
||||
_Section(".rodata.pytab.1") const struct _inittab _PyImport_Inittab_tokenbucket = {
|
||||
"tokenbucket",
|
||||
PyInit_tokenbucket,
|
||||
};
|
||||
2
third_party/python/Python/cosmomodule.c
vendored
2
third_party/python/Python/cosmomodule.c
vendored
|
|
@ -356,7 +356,7 @@ static PyMethodDef cosmo_methods[] = {
|
|||
{"syscount", cosmo_syscount, METH_NOARGS, syscount_doc},
|
||||
{"popcount", cosmo_popcount, METH_VARARGS, popcount_doc},
|
||||
{"decimate", cosmo_decimate, METH_VARARGS, decimate_doc},
|
||||
{"verynice", cosmo_verynice, METH_VARARGS, verynice_doc},
|
||||
{"verynice", cosmo_verynice, METH_NOARGS, verynice_doc},
|
||||
#ifdef __x86_64__
|
||||
{"getcpucore", cosmo_getcpucore, METH_NOARGS, getcpucore_doc},
|
||||
{"getcpunode", cosmo_getcpunode, METH_NOARGS, getcpunode_doc},
|
||||
|
|
|
|||
3
third_party/python/python.c
vendored
3
third_party/python/python.c
vendored
|
|
@ -226,7 +226,6 @@ PYTHON_YOINK("decimal");
|
|||
PYTHON_YOINK("difflib");
|
||||
PYTHON_YOINK("doctest");
|
||||
PYTHON_YOINK("dummy_threading");
|
||||
PYTHON_YOINK("threading");
|
||||
PYTHON_YOINK("enum");
|
||||
PYTHON_YOINK("filecmp");
|
||||
PYTHON_YOINK("fileinput");
|
||||
|
|
@ -306,7 +305,9 @@ PYTHON_YOINK("tabnanny");
|
|||
PYTHON_YOINK("tempfile");
|
||||
PYTHON_YOINK("textwrap");
|
||||
PYTHON_YOINK("this");
|
||||
PYTHON_YOINK("threading");
|
||||
PYTHON_YOINK("token");
|
||||
PYTHON_YOINK("tokenbucket");
|
||||
PYTHON_YOINK("tokenize");
|
||||
PYTHON_YOINK("trace");
|
||||
PYTHON_YOINK("traceback");
|
||||
|
|
|
|||
2
third_party/python/python.mk
vendored
2
third_party/python/python.mk
vendored
|
|
@ -685,6 +685,7 @@ THIRD_PARTY_PYTHON_STAGE2_A_SRCS = \
|
|||
third_party/python/Modules/mmapmodule.c \
|
||||
third_party/python/Modules/parsermodule.c \
|
||||
third_party/python/Modules/posixmodule.c \
|
||||
third_party/python/Modules/tokenbucket.c \
|
||||
third_party/python/Modules/pwdmodule.c \
|
||||
third_party/python/Modules/pyexpat.c \
|
||||
third_party/python/Modules/resource.c \
|
||||
|
|
@ -1177,6 +1178,7 @@ THIRD_PARTY_PYTHON_STAGE2_A_DIRECTDEPS = \
|
|||
LIBC_TINYMATH \
|
||||
LIBC_X \
|
||||
LIBC_ZIPOS \
|
||||
NET_HTTP \
|
||||
NET_HTTPS \
|
||||
THIRD_PARTY_BZIP2 \
|
||||
THIRD_PARTY_GDTOA \
|
||||
|
|
|
|||
|
|
@ -101,6 +101,7 @@
|
|||
"__builtin_ia32_movntdq"
|
||||
"__has_attribute"
|
||||
"__has_builtin"
|
||||
"__has_include"
|
||||
"__has_cpp_attribute"
|
||||
"__has_feature"
|
||||
"__ATOMIC_RELAXED"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue