1 /*-
2 * Copyright 2016 Vsevolod Stakhov
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16 #include "lua_common.h"
17
18 /***
19 * @module rspamd_mempool
20 * Rspamd memory pool is used to allocate memory attached to specific objects,
21 * namely it was initially used for memory allocation for rspamd_task.
22 *
23 * All memory allocated by the pool is destroyed when the associated object is
24 * destroyed. This allows a sort of controlled garbage collection for memory
25 * allocated from the pool. Memory pools are extensively used by rspamd internal
26 * components and provide some powerful features, such as destructors or
27 * persistent variables.
28 * @example
29 local mempool = require "rspamd_mempool"
30 local pool = mempool.create()
31
32 pool:set_variable('a', 'bcd', 1, 1.01, false)
33 local v1, v2, v3, v4 = pool:get_variable('a', 'string,double,double,bool')
34 pool:destroy()
35 */
36
37 /* Lua bindings */
38 /***
39 * @function mempool.create([size])
40 * Creates a memory pool of a specified `size` or platform dependent optimal size (normally, a page size)
41 * @param {number} size size of a page inside pool
42 * @return {rspamd_mempool} new pool object (that should be removed by explicit call to `pool:destroy()`)
43 */
44 LUA_FUNCTION_DEF (mempool, create);
45 /***
46 * @method mempool:add_destructor(func)
47 * Adds new destructor function to the pool
48 * @param {function} func function to be called when the pool is destroyed
49 */
50 LUA_FUNCTION_DEF (mempool, add_destructor);
51 /***
52 * @method mempool:destroy()
53 * Destroys memory pool cleaning all variables and calling all destructors registered (both C and Lua ones)
54 */
55 LUA_FUNCTION_DEF (mempool, delete);
56 LUA_FUNCTION_DEF (mempool, stat);
57 LUA_FUNCTION_DEF (mempool, suggest_size);
58 /***
59 * @method mempool:set_variable(name, [value1[, value2 ...]])
60 * Sets a variable that's valid during memory pool lifetime. This function allows
61 * to pack multiple values inside a single variable. Currently supported types are:
62 *
63 * - `string`: packed as null terminated C string (so no `\0` are allowed)
64 * - `number`: packed as C double
65 * - `boolean`: packed as bool
66 * @param {string} name variable's name to set
67 */
68 LUA_FUNCTION_DEF (mempool, set_variable);
69 /***
70 * @method mempool:set_bucket(name, num_values, [value1...valuen]|[table])
71 * Stores a variable bucket of numbers where the first number is number of elements to pack
72 * and then there should be either n numeric values or a plain table of numeric values
73 * @param {string} name variable's name to set
74 * @param {number} num_values number of variables in the bucket
75 * @param {table|list} values values
76 */
77 LUA_FUNCTION_DEF (mempool, set_bucket);
78 /***
79 * @method mempool:get_variable(name[, type])
80 * Unpacks mempool variable to lua If `type` is not specified, then a variable is
81 * assumed to be zero-terminated C string. Otherwise, `type` is a comma separated (spaces are ignored)
82 * list of types that should be unpacked from a variable's content. The following types
83 * are supported:
84 *
85 * - `string`: null terminated C string (so no `\0` are allowed)
86 * - `double`: returned as lua number
87 * - `int`: unpack a single integer
88 * - `int64`: unpack 64-bits integer
89 * - `boolean`: unpack boolean
90 * - `bucket`: bucket of numbers represented as a Lua table
91 * - `fstrings`: list of rspamd_fstring_t (GList) represented as a Lua table
92 * @param {string} name variable's name to get
93 * @param {string} type list of types to be extracted
94 * @return {variable list} list of variables extracted (but **not** a table)
95 */
96 LUA_FUNCTION_DEF (mempool, get_variable);
97 /***
98 * @method mempool:has_variable(name)
99 * Checks if the specified variable `name` exists in the memory pool
100 * @param {string} name variable's name to get
101 * @return {boolean} `true` if variable exists and `false` otherwise
102 */
103 LUA_FUNCTION_DEF (mempool, has_variable);
104
105 /***
106 * @method mempool:delete_variable(name)
107 * Removes the specified variable `name` from the memory pool
108 * @param {string} name variable's name to remove
109 * @return {boolean} `true` if variable exists and has been removed
110 */
111 LUA_FUNCTION_DEF (mempool, delete_variable);
112 /**
113 * @method mempool:topointer()
114 *
115 * Returns raw C pointer (lightuserdata) associated with mempool. This might be
116 * broken with luajit and GC64, use with caution.
117 */
118 LUA_FUNCTION_DEF (mempool, topointer);
119
120 static const struct luaL_reg mempoollib_m[] = {
121 LUA_INTERFACE_DEF (mempool, add_destructor),
122 LUA_INTERFACE_DEF (mempool, stat),
123 LUA_INTERFACE_DEF (mempool, suggest_size),
124 LUA_INTERFACE_DEF (mempool, set_variable),
125 LUA_INTERFACE_DEF (mempool, set_bucket),
126 LUA_INTERFACE_DEF (mempool, get_variable),
127 LUA_INTERFACE_DEF (mempool, has_variable),
128 LUA_INTERFACE_DEF (mempool, delete_variable),
129 LUA_INTERFACE_DEF (mempool, topointer),
130 LUA_INTERFACE_DEF (mempool, delete),
131 {"destroy", lua_mempool_delete},
132 {"__tostring", rspamd_lua_class_tostring},
133 {NULL, NULL}
134 };
135
136 static const struct luaL_reg mempoollib_f[] = {
137 LUA_INTERFACE_DEF (mempool, create),
138 {NULL, NULL}
139 };
140
141 /*
142 * Struct for lua destructor
143 */
144
145 struct lua_mempool_udata {
146 lua_State *L;
147 gint cbref;
148 rspamd_mempool_t *mempool;
149 };
150
151 struct memory_pool_s *
rspamd_lua_check_mempool(lua_State * L,gint pos)152 rspamd_lua_check_mempool (lua_State * L, gint pos)
153 {
154 void *ud = rspamd_lua_check_udata (L, pos, "rspamd{mempool}");
155 luaL_argcheck (L, ud != NULL, pos, "'mempool' expected");
156 return ud ? *((struct memory_pool_s **)ud) : NULL;
157 }
158
159
160 static int
lua_mempool_create(lua_State * L)161 lua_mempool_create (lua_State *L)
162 {
163 LUA_TRACE_POINT;
164 struct memory_pool_s *mempool = rspamd_mempool_new (
165 rspamd_mempool_suggest_size (), "lua", 0), **pmempool;
166
167 if (mempool) {
168 pmempool = lua_newuserdata (L, sizeof (struct memory_pool_s *));
169 rspamd_lua_setclass (L, "rspamd{mempool}", -1);
170 *pmempool = mempool;
171 }
172 else {
173 lua_pushnil (L);
174 }
175
176 return 1;
177 }
178
179 static void
lua_mempool_destructor_func(gpointer p)180 lua_mempool_destructor_func (gpointer p)
181 {
182 struct lua_mempool_udata *ud = p;
183
184 lua_rawgeti (ud->L, LUA_REGISTRYINDEX, ud->cbref);
185 if (lua_pcall (ud->L, 0, 0, 0) != 0) {
186 msg_info ("call to destructor failed: %s", lua_tostring (ud->L, -1));
187 lua_pop (ud->L, 1);
188 }
189 luaL_unref (ud->L, LUA_REGISTRYINDEX, ud->cbref);
190 }
191
192 static int
lua_mempool_add_destructor(lua_State * L)193 lua_mempool_add_destructor (lua_State *L)
194 {
195 LUA_TRACE_POINT;
196 struct memory_pool_s *mempool = rspamd_lua_check_mempool (L, 1);
197 struct lua_mempool_udata *ud;
198
199 if (mempool) {
200 if (lua_isfunction (L, 2)) {
201 ud = rspamd_mempool_alloc (mempool,
202 sizeof (struct lua_mempool_udata));
203 lua_pushvalue (L, 2);
204 /* Get a reference */
205 ud->cbref = luaL_ref (L, LUA_REGISTRYINDEX);
206 ud->L = L;
207 ud->mempool = mempool;
208 rspamd_mempool_add_destructor (mempool,
209 lua_mempool_destructor_func,
210 ud);
211 }
212 else {
213 msg_err ("trying to add destructor without function");
214 }
215 }
216 else {
217 lua_pushnil (L);
218 }
219
220 return 1;
221 }
222
223 static int
lua_mempool_delete(lua_State * L)224 lua_mempool_delete (lua_State *L)
225 {
226 LUA_TRACE_POINT;
227 struct memory_pool_s *mempool = rspamd_lua_check_mempool (L, 1);
228
229 if (mempool) {
230 rspamd_mempool_delete (mempool);
231 return 0;
232 }
233 else {
234 lua_pushnil (L);
235 }
236
237 return 1;
238 }
239
240 static int
lua_mempool_stat(lua_State * L)241 lua_mempool_stat (lua_State *L)
242 {
243 LUA_TRACE_POINT;
244 struct memory_pool_s *mempool = rspamd_lua_check_mempool (L, 1);
245
246 if (mempool) {
247
248 }
249 else {
250 lua_pushnil (L);
251 }
252
253 return 1;
254 }
255
256 static int
lua_mempool_suggest_size(lua_State * L)257 lua_mempool_suggest_size (lua_State *L)
258 {
259 LUA_TRACE_POINT;
260 struct memory_pool_s *mempool = rspamd_lua_check_mempool (L, 1);
261
262 if (mempool) {
263 lua_pushinteger (L, rspamd_mempool_suggest_size ());
264 return 0;
265 }
266 else {
267 lua_pushnil (L);
268 }
269
270 return 1;
271 }
272
273 struct lua_numbers_bucket {
274 guint nelts;
275 gdouble elts[0];
276 };
277
278 static int
lua_mempool_set_bucket(lua_State * L)279 lua_mempool_set_bucket (lua_State *L)
280 {
281 LUA_TRACE_POINT;
282 struct memory_pool_s *mempool = rspamd_lua_check_mempool (L, 1);
283 const gchar *var = luaL_checkstring (L, 2);
284 struct lua_numbers_bucket *bucket;
285 gint nelts = luaL_checknumber (L, 3), i;
286
287 if (var && nelts > 0) {
288 bucket = rspamd_mempool_alloc (mempool,
289 sizeof (*bucket) + sizeof (gdouble) * nelts);
290 bucket->nelts = nelts;
291
292 if (lua_type (L, 4) == LUA_TTABLE) {
293 /* Table version */
294 for (i = 1; i <= nelts; i ++) {
295 lua_rawgeti (L, 4, i);
296 bucket->elts[i - 1] = lua_tonumber (L, -1);
297 lua_pop (L, 1);
298 }
299 }
300 else {
301 for (i = 0; i <= nelts; i ++) {
302 bucket->elts[i] = lua_tonumber (L, 4 + i);
303 }
304 }
305
306 rspamd_mempool_set_variable (mempool, var, bucket, NULL);
307 }
308 else {
309 return luaL_error (L, "invalid arguments");
310 }
311
312 return 0;
313 }
314
315 static int
lua_mempool_set_variable(lua_State * L)316 lua_mempool_set_variable (lua_State *L)
317 {
318 LUA_TRACE_POINT;
319 struct memory_pool_s *mempool = rspamd_lua_check_mempool (L, 1);
320 const gchar *var = luaL_checkstring (L, 2);
321 gpointer value;
322 struct lua_numbers_bucket *bucket;
323 gchar *vp;
324 union {
325 gdouble d;
326 const gchar *s;
327 gboolean b;
328 } val;
329 gsize slen;
330 gint i, j, len = 0, type;
331
332 if (mempool && var) {
333
334 for (i = 3; i <= lua_gettop (L); i ++) {
335 type = lua_type (L, i);
336
337 if (type == LUA_TNUMBER) {
338 /* We have some ambiguity here between integer and double */
339 len += sizeof (gdouble);
340 }
341 else if (type == LUA_TBOOLEAN) {
342 len += sizeof (gboolean);
343 }
344 else if (type == LUA_TSTRING) {
345 (void)lua_tolstring (L, i, &slen);
346 len += slen + 1;
347 }
348 else if (type == LUA_TTABLE) {
349 /* We assume it as a bucket of numbers so far */
350 slen = rspamd_lua_table_size (L, i);
351 len += sizeof (gdouble) * slen + sizeof (*bucket);
352 }
353 else {
354 msg_err ("cannot handle lua type %s", lua_typename (L, type));
355 }
356 }
357
358 if (len == 0) {
359 msg_err ("no values specified");
360 }
361 else {
362 value = rspamd_mempool_alloc (mempool, len);
363 vp = value;
364
365 for (i = 3; i <= lua_gettop (L); i ++) {
366 type = lua_type (L, i);
367
368 if (type == LUA_TNUMBER) {
369 val.d = lua_tonumber (L, i);
370 memcpy (vp, &val, sizeof (gdouble));
371 vp += sizeof (gdouble);
372 }
373 else if (type == LUA_TBOOLEAN) {
374 val.b = lua_toboolean (L, i);
375 memcpy (vp, &val, sizeof (gboolean));
376 vp += sizeof (gboolean);
377 }
378 else if (type == LUA_TSTRING) {
379 val.s = lua_tolstring (L, i, &slen);
380 memcpy (vp, val.s, slen + 1);
381 vp += slen + 1;
382 }
383 else if (type == LUA_TTABLE) {
384 slen = rspamd_lua_table_size (L, i);
385 /* XXX: Ret, ret, ret: alignment issues */
386 bucket = (struct lua_numbers_bucket *)vp;
387 bucket->nelts = slen;
388
389 for (j = 0; j < slen; j ++) {
390 lua_rawgeti (L, i, j + 1);
391 bucket->elts[j] = lua_tonumber (L, -1);
392 lua_pop (L, 1);
393 }
394
395 vp += sizeof (gdouble) * slen + sizeof (*bucket);
396 }
397 else {
398 msg_err ("cannot handle lua type %s", lua_typename (L, type));
399 }
400 }
401
402 rspamd_mempool_set_variable (mempool, var, value, NULL);
403 }
404
405 return 0;
406 }
407 else {
408 lua_pushnil (L);
409 }
410
411 return 1;
412 }
413
414
415 static int
lua_mempool_get_variable(lua_State * L)416 lua_mempool_get_variable (lua_State *L)
417 {
418 LUA_TRACE_POINT;
419 struct memory_pool_s *mempool = rspamd_lua_check_mempool (L, 1);
420 const gchar *var = luaL_checkstring (L, 2);
421 const gchar *type = NULL, *pt;
422 struct lua_numbers_bucket bucket;
423 const gchar *value, *pv;
424 guint len, nvar, slen, i;
425
426 if (mempool && var) {
427 value = rspamd_mempool_get_variable (mempool, var);
428
429 if (lua_gettop (L) >= 3) {
430 type = luaL_checkstring (L, 3);
431 }
432
433 if (value) {
434
435 if (type) {
436 pt = type;
437 pv = value;
438 nvar = 0;
439
440 while ((len = strcspn (pt, ", ")) > 0) {
441 if (len == sizeof ("double") - 1 &&
442 g_ascii_strncasecmp (pt, "double", len) == 0) {
443 gdouble num;
444 memcpy (&num, pv, sizeof (gdouble));
445 lua_pushnumber (L, num);
446 pv += sizeof (gdouble);
447 }
448 else if (len == sizeof ("int") - 1 &&
449 g_ascii_strncasecmp (pt, "int", len) == 0) {
450 gint num;
451 memcpy (&num, pv, sizeof (gint));
452 lua_pushinteger (L, num);
453 pv += sizeof (gint);
454 }
455 else if (len == sizeof ("int64") - 1 &&
456 g_ascii_strncasecmp (pt, "int64", len) == 0) {
457 gint64 num;
458 memcpy (&num, pv, sizeof (gint64));
459 lua_pushinteger (L, num);
460 pv += sizeof (gint64);
461 }
462 else if (len == sizeof ("bool") - 1 &&
463 g_ascii_strncasecmp (pt, "bool", len) == 0) {
464 gboolean num;
465 memcpy (&num, pv, sizeof (gboolean));
466 lua_pushboolean (L, num);
467 pv += sizeof (gboolean);
468 }
469 else if (len == sizeof ("string") - 1 &&
470 g_ascii_strncasecmp (pt, "string", len) == 0) {
471 slen = strlen ((const gchar *)pv);
472 lua_pushlstring (L, (const gchar *)pv, slen);
473 pv += slen + 1;
474 }
475 else if (len == sizeof ("gstring") - 1 &&
476 g_ascii_strncasecmp (pt, "gstring", len) == 0) {
477 GString *st = (GString *)pv;
478 lua_pushlstring (L, st->str, st->len);
479 pv += sizeof (GString *);
480 }
481 else if (len == sizeof ("bucket") - 1 &&
482 g_ascii_strncasecmp (pt, "bucket", len) == 0) {
483 memcpy (&bucket, pv, sizeof (bucket));
484 lua_createtable (L, bucket.nelts, 0);
485 pv += sizeof (struct lua_numbers_bucket);
486
487 for (i = 0; i < bucket.nelts; i ++) {
488 gdouble num;
489 memcpy (&num, pv, sizeof (num));
490 lua_pushnumber (L, num);
491 lua_rawseti (L, -2, i + 1);
492 pv += sizeof (num);
493 }
494 }
495 else if (len == sizeof ("fstrings") - 1 &&
496 g_ascii_strncasecmp (pt, "fstrings", len) == 0) {
497 GList *cur;
498 rspamd_fstring_t *fstr;
499
500 cur = (GList *)pv;
501 lua_newtable (L);
502
503 i = 1;
504 while (cur != NULL) {
505 fstr = cur->data;
506 lua_pushlstring (L, fstr->str, fstr->len);
507 lua_rawseti (L, -2, i);
508 i ++;
509 cur = g_list_next (cur);
510 }
511
512 pv += sizeof (GList *);
513 }
514 else {
515 msg_err ("unknown type for get_variable: %s", pt);
516 lua_pushnil (L);
517 }
518
519 pt += len;
520 pt += strspn (pt, ", ");
521
522 nvar ++;
523 }
524
525 return nvar;
526 }
527 else {
528 /* No type specified, return string */
529 lua_pushstring(L, value);
530 }
531 }
532 else {
533 lua_pushnil (L);
534 }
535 }
536 else {
537 lua_pushnil (L);
538 }
539
540 return 1;
541 }
542
543 static int
lua_mempool_has_variable(lua_State * L)544 lua_mempool_has_variable (lua_State *L)
545 {
546 LUA_TRACE_POINT;
547 struct memory_pool_s *mempool = rspamd_lua_check_mempool (L, 1);
548 const gchar *var = luaL_checkstring (L, 2);
549 gboolean ret = FALSE;
550
551 if (mempool && var) {
552 if (rspamd_mempool_get_variable (mempool, var) != NULL) {
553 ret = TRUE;
554 }
555 }
556
557 lua_pushboolean (L, ret);
558
559 return 1;
560 }
561
562 static int
lua_mempool_delete_variable(lua_State * L)563 lua_mempool_delete_variable (lua_State *L)
564 {
565 LUA_TRACE_POINT;
566 struct memory_pool_s *mempool = rspamd_lua_check_mempool (L, 1);
567 const gchar *var = luaL_checkstring (L, 2);
568 gboolean ret = FALSE;
569
570 if (mempool && var) {
571 if (rspamd_mempool_get_variable (mempool, var) != NULL) {
572 ret = TRUE;
573
574 rspamd_mempool_remove_variable (mempool, var);
575 }
576 }
577
578 lua_pushboolean (L, ret);
579
580 return 1;
581 }
582
583 static gint
lua_mempool_topointer(lua_State * L)584 lua_mempool_topointer (lua_State *L)
585 {
586 LUA_TRACE_POINT;
587 rspamd_mempool_t *pool = rspamd_lua_check_mempool (L, 1);
588
589 if (pool) {
590 /* XXX: this might cause issues on arm64 and LuaJIT */
591 lua_pushlightuserdata (L, pool);
592 }
593 else {
594 return luaL_error (L, "invalid arguments");
595 }
596
597 return 1;
598 }
599
600 static gint
lua_load_mempool(lua_State * L)601 lua_load_mempool (lua_State * L)
602 {
603 lua_newtable (L);
604 luaL_register (L, NULL, mempoollib_f);
605
606 return 1;
607 }
608
609 void
luaopen_mempool(lua_State * L)610 luaopen_mempool (lua_State * L)
611 {
612 rspamd_lua_new_class (L, "rspamd{mempool}", mempoollib_m);
613 lua_pop (L, 1);
614 rspamd_lua_add_preload (L, "rspamd_mempool", lua_load_mempool);
615 }
616