<commit_before>#include <stdio.h>
#include <stdlib.h>

#include "vrt.h"
#include "cache/cache.h"
#include "vsha256.h"

#include "vtree.h"
#include <sys/time.h>

#include "vcc_if.h"


/* Represents a token bucket for a specific key. */
struct tbucket {
	unsigned		magic;
#define TBUCKET_MAGIC		0x53345eb9
	unsigned char		digest[DIGEST_LEN];
	double			last_used;
	double			period;
	long			tokens;
	long			capacity;
	VRB_ENTRY(tbucket)	tree;
};

static int
keycmp(const struct tbucket *b1, const struct tbucket *b2) {
	return (memcmp(b1->digest, b2->digest, sizeof b1->digest));
}

VRB_HEAD(tbtree, tbucket);
VRB_PROTOTYPE_STATIC(tbtree, tbucket, tree, keycmp);
VRB_GENERATE_STATIC(tbtree, tbucket, tree, keycmp);

/* To lessen potential mutex contention, we partition the buckets into
   N_PART partitions.  */
#define N_PART		16 /* must be 2^n */
#define N_PART_MASK	(N_PART - 1)

static unsigned n_init;
static pthread_mutex_t init_mtx = PTHREAD_MUTEX_INITIALIZER;

/* GC_INTVL: How often (in #calls per partition) we invoke the garbage
   collector. */
#define GC_INTVL	1000
static void run_gc(double now, unsigned part);

static struct vsthrottle {
	unsigned		magic;
#define VSTHROTTLE_MAGIC	0x99fdbef8
	pthread_mutex_t		mtx;
	struct tbtree		buckets;
	unsigned		gc_count;
} vsthrottle[N_PART];

static struct tbucket *
tb_alloc(const unsigned char *digest, long limit, double period, double now) {
	struct tbucket *tb = malloc(sizeof *tb);
	AN(tb);

	memcpy(tb->digest, digest, sizeof tb->digest);
	tb->magic = TBUCKET_MAGIC;
	tb->last_used = now;
	tb->period = period;
	tb->tokens = limit;
	tb->capacity = limit;

	return (tb);
}

static struct tbucket *
get_bucket(const unsigned char *digest, long limit, double period, double now) {
	struct tbucket *b;
	struct tbucket k = { 0 };
	unsigned part = digest[0] & N_PART_MASK;
	struct vsthrottle *v = &vsthrottle[part];

	memcpy(&k.digest, digest, sizeof k.digest);
	b = VRB_FIND(tbtree, &v->buckets, &k);
	if (b) {
		CHECK_OBJ_NOTNULL(b, TBUCKET_MAGIC);
	} else {
		b = tb_alloc(digest, limit, period, now);
		AZ(VRB_INSERT(tbtree, &v->buckets, b));
	}
	return (b);
}

static void
calc_tokens(struct tbucket *b, double now) {
	double delta = now - b->last_used;

	b->tokens += (long) ((delta / b->period) * b->capacity);
	if (b->tokens > b->capacity)
		b->tokens = b->capacity;
	/* VSL(SLT_VCL_Log, 0, "tokens: %ld", b->tokens); */
}

static double
get_ts_now(const struct vrt_ctx *ctx) {
	double now;

	CHECK_OBJ_NOTNULL(ctx, VRT_CTX_MAGIC);
	if (ctx->req) {
		CHECK_OBJ_NOTNULL(ctx->req, REQ_MAGIC);
		now = ctx->req->t_prev;
	} else {
		CHECK_OBJ_NOTNULL(ctx->bo, BUSYOBJ_MAGIC);
		now = ctx->bo->t_prev;
	}

	return (now);
}

VCL_BOOL
vmod_is_denied(const struct vrt_ctx *ctx, VCL_STRING key, VCL_INT limit,
    VCL_DURATION period) {
	unsigned ret = 1;
	struct tbucket *b;
	double now = get_ts_now(ctx);
	SHA256_CTX sctx;
	struct vsthrottle *v;
	unsigned char digest[DIGEST_LEN];
	unsigned part;

	if (!key)
		return (1);

	SHA256_Init(&sctx);
	SHA256_Update(&sctx, key, strlen(key));
	SHA256_Final(digest, &sctx);

	part = digest[0] & N_PART_MASK;
	v = &vsthrottle[part];<commit_msg>Use monotonic clocks for timestamps.<commit_after>#include <stdio.h>
#include <stdlib.h>

#include "vrt.h"
#include "cache/cache.h"
#include "vsha256.h"

#include "vtree.h"
#include <sys/time.h>

#include "vcc_if.h"


/* Represents a token bucket for a specific key. */
struct tbucket {
	unsigned		magic;
#define TBUCKET_MAGIC		0x53345eb9
	unsigned char		digest[DIGEST_LEN];
	double			last_used;
	double			period;
	long			tokens;
	long			capacity;
	VRB_ENTRY(tbucket)	tree;
};

static int
keycmp(const struct tbucket *b1, const struct tbucket *b2) {
	return (memcmp(b1->digest, b2->digest, sizeof b1->digest));
}

VRB_HEAD(tbtree, tbucket);
VRB_PROTOTYPE_STATIC(tbtree, tbucket, tree, keycmp);
VRB_GENERATE_STATIC(tbtree, tbucket, tree, keycmp);

/* To lessen potential mutex contention, we partition the buckets into
   N_PART partitions.  */
#define N_PART		16 /* must be 2^n */
#define N_PART_MASK	(N_PART - 1)

static unsigned n_init;
static pthread_mutex_t init_mtx = PTHREAD_MUTEX_INITIALIZER;

/* GC_INTVL: How often (in #calls per partition) we invoke the garbage
   collector. */
#define GC_INTVL	1000
static void run_gc(double now, unsigned part);

static struct vsthrottle {
	unsigned		magic;
#define VSTHROTTLE_MAGIC	0x99fdbef8
	pthread_mutex_t		mtx;
	struct tbtree		buckets;
	unsigned		gc_count;
} vsthrottle[N_PART];

static struct tbucket *
tb_alloc(const unsigned char *digest, long limit, double period, double now) {
	struct tbucket *tb = malloc(sizeof *tb);
	AN(tb);

	memcpy(tb->digest, digest, sizeof tb->digest);
	tb->magic = TBUCKET_MAGIC;
	tb->last_used = now;
	tb->period = period;
	tb->tokens = limit;
	tb->capacity = limit;

	return (tb);
}

static struct tbucket *
get_bucket(const unsigned char *digest, long limit, double period, double now) {
	struct tbucket *b;
	struct tbucket k = { 0 };
	unsigned part = digest[0] & N_PART_MASK;
	struct vsthrottle *v = &vsthrottle[part];

	memcpy(&k.digest, digest, sizeof k.digest);
	b = VRB_FIND(tbtree, &v->buckets, &k);
	if (b) {
		CHECK_OBJ_NOTNULL(b, TBUCKET_MAGIC);
	} else {
		b = tb_alloc(digest, limit, period, now);
		AZ(VRB_INSERT(tbtree, &v->buckets, b));
	}
	return (b);
}

static void
calc_tokens(struct tbucket *b, double now) {
	double delta = now - b->last_used;
	assert(delta >= 0);

	b->tokens += (long) ((delta / b->period) * b->capacity);
	if (b->tokens > b->capacity)
		b->tokens = b->capacity;
	/* VSL(SLT_VCL_Log, 0, "tokens: %ld", b->tokens); */
}

static double
get_ts_mono(const struct vrt_ctx *ctx) {
	struct timespec ts;
	AZ(clock_gettime(CLOCK_MONOTONIC, &ts));
	return (ts.tv_sec + 1e-9 * ts.tv_nsec);
}

VCL_BOOL
vmod_is_denied(const struct vrt_ctx *ctx, VCL_STRING key, VCL_INT limit,
    VCL_DURATION period) {
	unsigned ret = 1;
	struct tbucket *b;
	double now = get_ts_mono(ctx);
	SHA256_CTX sctx;
	struct vsthrottle *v;
	unsigned char digest[DIGEST_LEN];
	unsigned part;

	if (!key)
		return (1);

	SHA256_Init(&sctx);
	SHA256_Update(&sctx, key, strlen(key));
	SHA256_Final(digest, &sctx);

	part = digest[0] & N_PART_MASK;
	v = &vsthrottle[part];
{'fraction_non_alphanumeric': 0.1317403939296093, 'fraction_numerical': 0.015014530190506942, 'mean_word_length': 4.6369426751592355, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 61, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3671930', 'n_tokens_mistral': 2629, 'n_tokens_neox': 2486, 'n_words': 627}