You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

lua_worker.c 21KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888
  1. /*-
  2. * Copyright 2019 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. #include "unix-std.h"
  18. #include "worker_util.h"
  19. #include "rspamd_control.h"
  20. #include "ottery.h"
  21. #ifdef WITH_JEMALLOC
  22. #include <jemalloc/jemalloc.h>
  23. #endif
  24. #include <sys/wait.h>
  25. #include <src/libserver/rspamd_control.h>
  26. /***
  27. * @module rspamd_worker
  28. * This module provides methods to access worker related functions in various
  29. * places, such as periodic events or on_load events.
  30. */
  31. LUA_FUNCTION_DEF (worker, get_name);
  32. LUA_FUNCTION_DEF (worker, get_stat);
  33. LUA_FUNCTION_DEF (worker, get_index);
  34. LUA_FUNCTION_DEF (worker, get_count);
  35. LUA_FUNCTION_DEF (worker, get_pid);
  36. LUA_FUNCTION_DEF (worker, is_scanner);
  37. LUA_FUNCTION_DEF (worker, is_primary_controller);
  38. LUA_FUNCTION_DEF (worker, spawn_process);
  39. LUA_FUNCTION_DEF (worker, get_mem_stats);
  40. LUA_FUNCTION_DEF (worker, add_control_handler);
  41. const luaL_reg worker_reg[] = {
  42. LUA_INTERFACE_DEF (worker, get_name),
  43. {"get_type", lua_worker_get_name},
  44. LUA_INTERFACE_DEF (worker, get_stat),
  45. LUA_INTERFACE_DEF (worker, get_index),
  46. LUA_INTERFACE_DEF (worker, get_count),
  47. LUA_INTERFACE_DEF (worker, get_pid),
  48. LUA_INTERFACE_DEF (worker, spawn_process),
  49. LUA_INTERFACE_DEF (worker, is_scanner),
  50. LUA_INTERFACE_DEF (worker, is_primary_controller),
  51. LUA_INTERFACE_DEF (worker, get_mem_stats),
  52. LUA_INTERFACE_DEF (worker, add_control_handler),
  53. {"__tostring", rspamd_lua_class_tostring},
  54. {NULL, NULL}
  55. };
  56. static struct rspamd_worker *
  57. lua_check_worker (lua_State *L, gint pos)
  58. {
  59. void *ud = rspamd_lua_check_udata (L, pos, "rspamd{worker}");
  60. luaL_argcheck (L, ud != NULL, pos, "'worker' expected");
  61. return ud ? *((struct rspamd_worker **)ud) : NULL;
  62. }
  63. static gint
  64. lua_worker_get_stat (lua_State *L)
  65. {
  66. struct rspamd_worker *w = lua_check_worker (L, 1);
  67. if (w) {
  68. rspamd_mempool_stat_t mem_st;
  69. struct rspamd_stat *stat, stat_copy;
  70. ucl_object_t *top, *sub;
  71. gint i;
  72. guint64 spam = 0, ham = 0;
  73. memset (&mem_st, 0, sizeof (mem_st));
  74. rspamd_mempool_stat (&mem_st);
  75. memcpy (&stat_copy, w->srv->stat, sizeof (stat_copy));
  76. stat = &stat_copy;
  77. top = ucl_object_typed_new (UCL_OBJECT);
  78. ucl_object_insert_key (top, ucl_object_fromint (
  79. stat->messages_scanned), "scanned", 0, false);
  80. ucl_object_insert_key (top, ucl_object_fromint (
  81. stat->messages_learned), "learned", 0, false);
  82. if (stat->messages_scanned > 0) {
  83. sub = ucl_object_typed_new (UCL_OBJECT);
  84. for (i = METRIC_ACTION_REJECT; i <= METRIC_ACTION_NOACTION; i++) {
  85. ucl_object_insert_key (sub,
  86. ucl_object_fromint (stat->actions_stat[i]),
  87. rspamd_action_to_str (i), 0, false);
  88. if (i < METRIC_ACTION_GREYLIST) {
  89. spam += stat->actions_stat[i];
  90. }
  91. else {
  92. ham += stat->actions_stat[i];
  93. }
  94. }
  95. ucl_object_insert_key (top, sub, "actions", 0, false);
  96. }
  97. else {
  98. sub = ucl_object_typed_new (UCL_OBJECT);
  99. for (i = METRIC_ACTION_REJECT; i <= METRIC_ACTION_NOACTION; i++) {
  100. ucl_object_insert_key (sub,
  101. 0,
  102. rspamd_action_to_str (i), 0, false);
  103. }
  104. ucl_object_insert_key (top, sub, "actions", 0, false);
  105. }
  106. ucl_object_insert_key (top, ucl_object_fromint (
  107. spam), "spam_count", 0, false);
  108. ucl_object_insert_key (top, ucl_object_fromint (
  109. ham), "ham_count", 0, false);
  110. ucl_object_insert_key (top,
  111. ucl_object_fromint (stat->connections_count), "connections", 0, false);
  112. ucl_object_insert_key (top,
  113. ucl_object_fromint (stat->control_connections_count),
  114. "control_connections", 0, false);
  115. ucl_object_insert_key (top,
  116. ucl_object_fromint (mem_st.pools_allocated), "pools_allocated", 0,
  117. false);
  118. ucl_object_insert_key (top,
  119. ucl_object_fromint (mem_st.pools_freed), "pools_freed", 0, false);
  120. ucl_object_insert_key (top,
  121. ucl_object_fromint (mem_st.bytes_allocated), "bytes_allocated", 0,
  122. false);
  123. ucl_object_insert_key (top,
  124. ucl_object_fromint (
  125. mem_st.chunks_allocated), "chunks_allocated", 0, false);
  126. ucl_object_insert_key (top,
  127. ucl_object_fromint (mem_st.shared_chunks_allocated),
  128. "shared_chunks_allocated", 0, false);
  129. ucl_object_insert_key (top,
  130. ucl_object_fromint (mem_st.chunks_freed), "chunks_freed", 0, false);
  131. ucl_object_insert_key (top,
  132. ucl_object_fromint (
  133. mem_st.oversized_chunks), "chunks_oversized", 0, false);
  134. ucl_object_push_lua (L, top, true);
  135. ucl_object_unref (top);
  136. }
  137. else {
  138. return luaL_error (L, "invalid arguments");
  139. }
  140. return 1;
  141. }
  142. static gint
  143. lua_worker_get_name (lua_State *L)
  144. {
  145. struct rspamd_worker *w = lua_check_worker (L, 1);
  146. if (w) {
  147. lua_pushstring (L, g_quark_to_string (w->type));
  148. }
  149. else {
  150. return luaL_error (L, "invalid arguments");
  151. }
  152. return 1;
  153. }
  154. static gint
  155. lua_worker_get_index (lua_State *L)
  156. {
  157. struct rspamd_worker *w = lua_check_worker (L, 1);
  158. if (w) {
  159. lua_pushinteger (L, w->index);
  160. }
  161. else {
  162. return luaL_error (L, "invalid arguments");
  163. }
  164. return 1;
  165. }
  166. static gint
  167. lua_worker_get_count (lua_State *L)
  168. {
  169. struct rspamd_worker *w = lua_check_worker (L, 1);
  170. if (w) {
  171. lua_pushinteger (L, w->cf->count);
  172. }
  173. else {
  174. return luaL_error (L, "invalid arguments");
  175. }
  176. return 1;
  177. }
  178. static gint
  179. lua_worker_get_pid (lua_State *L)
  180. {
  181. struct rspamd_worker *w = lua_check_worker (L, 1);
  182. if (w) {
  183. lua_pushinteger (L, w->pid);
  184. }
  185. else {
  186. return luaL_error (L, "invalid arguments");
  187. }
  188. return 1;
  189. }
  190. static gint
  191. lua_worker_is_scanner (lua_State *L)
  192. {
  193. struct rspamd_worker *w = lua_check_worker (L, 1);
  194. if (w) {
  195. lua_pushboolean (L, rspamd_worker_is_scanner (w));
  196. }
  197. else {
  198. return luaL_error (L, "invalid arguments");
  199. }
  200. return 1;
  201. }
  202. static gint
  203. lua_worker_is_primary_controller (lua_State *L)
  204. {
  205. struct rspamd_worker *w = lua_check_worker (L, 1);
  206. if (w) {
  207. lua_pushboolean (L, rspamd_worker_is_primary_controller (w));
  208. }
  209. else {
  210. return luaL_error (L, "invalid arguments");
  211. }
  212. return 1;
  213. }
  214. struct rspamd_control_cbdata {
  215. lua_State *L;
  216. rspamd_mempool_t *pool;
  217. struct rspamd_worker *w;
  218. struct rspamd_config *cfg;
  219. struct ev_loop *event_loop;
  220. struct rspamd_async_session *session;
  221. enum rspamd_control_type cmd;
  222. gint cbref;
  223. gint fd;
  224. };
  225. static gboolean
  226. lua_worker_control_fin_session (void *ud)
  227. {
  228. struct rspamd_control_reply rep;
  229. struct rspamd_control_cbdata *cbd = (struct rspamd_control_cbdata *)ud;
  230. rspamd_mempool_t *pool;
  231. pool = cbd->pool;
  232. memset (&rep, 0, sizeof (rep));
  233. rep.type = cbd->cmd;
  234. if (write (cbd->fd, &rep, sizeof (rep)) != sizeof (rep)) {
  235. msg_err_pool ("cannot write reply to the control socket: %s",
  236. strerror (errno));
  237. }
  238. return TRUE;
  239. }
  240. static void
  241. lua_worker_control_session_dtor (void *ud)
  242. {
  243. struct rspamd_control_cbdata *cbd = (struct rspamd_control_cbdata *)ud;
  244. rspamd_mempool_delete (cbd->pool);
  245. }
  246. static gboolean
  247. lua_worker_control_handler (struct rspamd_main *rspamd_main,
  248. struct rspamd_worker *worker,
  249. gint fd,
  250. gint attached_fd,
  251. struct rspamd_control_command *cmd,
  252. gpointer ud)
  253. {
  254. struct rspamd_async_session *session, **psession;
  255. struct rspamd_control_cbdata *cbd = (struct rspamd_control_cbdata *)ud;
  256. rspamd_mempool_t *pool;
  257. lua_State *L;
  258. gint err_idx, status;
  259. L = cbd->L;
  260. pool = cbd->pool;
  261. session = rspamd_session_create (cbd->pool,
  262. lua_worker_control_fin_session,
  263. NULL,
  264. lua_worker_control_session_dtor,
  265. cbd);
  266. cbd->session = session;
  267. cbd->fd = fd;
  268. lua_pushcfunction (L, &rspamd_lua_traceback);
  269. err_idx = lua_gettop (L);
  270. lua_rawgeti (L, LUA_REGISTRYINDEX, cbd->cbref);
  271. psession = lua_newuserdata (L, sizeof (*psession));
  272. rspamd_lua_setclass (L, "rspamd{session}", -1);
  273. *psession = session;
  274. /* Command name */
  275. lua_pushstring (L, rspamd_control_command_to_string (cmd->type));
  276. /* Command's extras */
  277. lua_newtable (L);
  278. switch (cmd->type) {
  279. case RSPAMD_CONTROL_CHILD_CHANGE:
  280. lua_pushinteger (L, cmd->cmd.child_change.pid);
  281. lua_setfield (L, -2, "pid");
  282. switch (cmd->cmd.child_change.what) {
  283. case rspamd_child_offline:
  284. lua_pushstring (L, "offline");
  285. lua_setfield (L, -2, "what");
  286. break;
  287. case rspamd_child_online:
  288. lua_pushstring (L, "online");
  289. lua_setfield (L, -2, "what");
  290. break;
  291. case rspamd_child_terminated:
  292. lua_pushstring (L, "terminated");
  293. lua_setfield (L, -2, "what");
  294. status = cmd->cmd.child_change.additional;
  295. if (WIFEXITED (status)) {
  296. lua_pushinteger (L, WEXITSTATUS (status));
  297. lua_setfield (L, -2, "exit_code");
  298. }
  299. if (WIFSIGNALED (status)) {
  300. lua_pushinteger (L, WTERMSIG (status));
  301. lua_setfield (L, -2, "signal");
  302. #ifdef WCOREDUMP
  303. lua_pushboolean (L, WCOREDUMP (status));
  304. lua_setfield (L, -2, "core");
  305. #endif
  306. }
  307. break;
  308. }
  309. break;
  310. case RSPAMD_CONTROL_MONITORED_CHANGE:
  311. lua_pushinteger (L, cmd->cmd.monitored_change.sender);
  312. lua_setfield (L, -2, "sender");
  313. lua_pushboolean (L, cmd->cmd.monitored_change.alive);
  314. lua_setfield (L, -2, "alive");
  315. lua_pushlstring (L, cmd->cmd.monitored_change.tag,
  316. sizeof (cmd->cmd.monitored_change.tag));
  317. lua_setfield (L, -2, "tag");
  318. break;
  319. case RSPAMD_CONTROL_HYPERSCAN_LOADED:
  320. lua_pushstring (L, cmd->cmd.hs_loaded.cache_dir);
  321. lua_setfield (L, -2, "cache_dir");
  322. lua_pushboolean (L, cmd->cmd.hs_loaded.forced);
  323. lua_setfield (L, -2, "forced");
  324. break;
  325. case RSPAMD_CONTROL_STAT:
  326. case RSPAMD_CONTROL_RELOAD:
  327. case RSPAMD_CONTROL_RERESOLVE:
  328. case RSPAMD_CONTROL_RECOMPILE:
  329. case RSPAMD_CONTROL_LOG_PIPE:
  330. case RSPAMD_CONTROL_FUZZY_STAT:
  331. case RSPAMD_CONTROL_FUZZY_SYNC:
  332. default:
  333. break;
  334. }
  335. if (lua_pcall (L, 3, 0, err_idx) != 0) {
  336. msg_err_pool ("cannot init lua parser script: %s", lua_tostring (L, -1));
  337. lua_settop (L, err_idx - 1);
  338. struct rspamd_control_reply rep;
  339. memset (&rep, 0, sizeof (rep));
  340. rep.type = cbd->cmd;
  341. rep.reply.monitored_change.status = -1;
  342. if (write (fd, &rep, sizeof (rep)) != sizeof (rep)) {
  343. msg_err_pool ("cannot write reply to the control socket: %s",
  344. strerror (errno));
  345. }
  346. rspamd_session_destroy (session);
  347. }
  348. else {
  349. lua_settop (L, err_idx - 1);
  350. rspamd_session_pending (session);
  351. }
  352. return TRUE;
  353. }
  354. static gint
  355. lua_worker_add_control_handler (lua_State *L)
  356. {
  357. struct rspamd_worker *w = lua_check_worker (L, 1);
  358. struct rspamd_config *cfg = lua_check_config (L, 2);
  359. struct ev_loop *event_loop = lua_check_ev_base (L, 3);
  360. const gchar *cmd_name = luaL_checkstring (L, 4);
  361. enum rspamd_control_type cmd;
  362. struct rspamd_control_cbdata *cbd;
  363. if (w && cfg && event_loop && cmd_name && lua_isfunction (L, 5)) {
  364. cmd = rspamd_control_command_from_string (cmd_name);
  365. if (cmd == RSPAMD_CONTROL_MAX) {
  366. return luaL_error (L, "invalid command type: %s", cmd_name);
  367. }
  368. rspamd_mempool_t *pool = rspamd_mempool_new (
  369. rspamd_mempool_suggest_size (), "lua_control", 0);
  370. cbd = rspamd_mempool_alloc0 (pool, sizeof (*cbd));
  371. cbd->pool = pool;
  372. cbd->event_loop = event_loop;
  373. cbd->w = w;
  374. cbd->cfg = cfg;
  375. cbd->cmd = cmd;
  376. cbd->L = L;
  377. /* Refcount callback */
  378. lua_pushvalue (L, 5);
  379. cbd->cbref = luaL_ref (L, LUA_REGISTRYINDEX);
  380. rspamd_control_worker_add_cmd_handler (w, cmd, lua_worker_control_handler,
  381. cbd);
  382. }
  383. else {
  384. return luaL_error (L, "invalid arguments, need worker, cfg, "
  385. "ev_loop, cmd_name and callback function");
  386. }
  387. return 0;
  388. }
  389. #ifdef WITH_JEMALLOC
  390. static void
  391. lua_worker_jemalloc_stats_cb (void *ud, const char *msg)
  392. {
  393. lua_State *L = (lua_State *)ud;
  394. lua_pushstring (L, msg);
  395. }
  396. #endif
  397. static gint
  398. lua_worker_get_mem_stats (lua_State *L)
  399. {
  400. struct rspamd_worker *w = lua_check_worker (L, 1);
  401. if (w) {
  402. #ifdef WITH_JEMALLOC
  403. malloc_stats_print (lua_worker_jemalloc_stats_cb, (void *)L, NULL);
  404. #else
  405. lua_pushstring (L, "no stats, jemalloc support is required");
  406. #endif
  407. }
  408. else {
  409. return luaL_error (L, "invalid arguments");
  410. }
  411. return 1;
  412. }
  413. struct rspamd_lua_process_cbdata {
  414. gint sp[2];
  415. gint func_cbref;
  416. gint cb_cbref;
  417. gboolean replied;
  418. gboolean is_error;
  419. pid_t cpid;
  420. lua_State *L;
  421. guint64 sz;
  422. GString *io_buf;
  423. GString *out_buf;
  424. goffset out_pos;
  425. struct rspamd_worker *wrk;
  426. struct ev_loop *event_loop;
  427. ev_io ev;
  428. };
  429. static void
  430. rspamd_lua_execute_lua_subprocess (lua_State *L,
  431. struct rspamd_lua_process_cbdata *cbdata)
  432. {
  433. gint err_idx, r;
  434. guint64 wlen = 0;
  435. lua_pushcfunction (L, &rspamd_lua_traceback);
  436. err_idx = lua_gettop (L);
  437. lua_rawgeti (L, LUA_REGISTRYINDEX, cbdata->func_cbref);
  438. if (lua_pcall (L, 0, 1, err_idx) != 0) {
  439. const gchar *s = lua_tostring (L, -1);
  440. gsize slen = strlen (s);
  441. msg_err ("call to subprocess failed: %s", s);
  442. /* Indicate error */
  443. wlen = (1ULL << 63u) + slen;
  444. r = write (cbdata->sp[1], &wlen, sizeof (wlen));
  445. if (r == -1) {
  446. msg_err ("write failed: %s", strerror (errno));
  447. }
  448. r = write (cbdata->sp[1], s, slen);
  449. if (r == -1) {
  450. msg_err ("write failed: %s", strerror (errno));
  451. }
  452. }
  453. else {
  454. struct rspamd_lua_text *t = lua_check_text_or_string (L, -1);
  455. if (t) {
  456. wlen = t->len;
  457. r = write (cbdata->sp[1], &wlen, sizeof (wlen));
  458. if (r == -1) {
  459. msg_err ("write failed: %s", strerror (errno));
  460. }
  461. r = write (cbdata->sp[1], t->start, wlen);
  462. if (r == -1) {
  463. msg_err ("write failed: %s", strerror (errno));
  464. }
  465. }
  466. else {
  467. msg_err ("subprocess: invalid return value: %s",
  468. lua_typename (L, lua_type (L, -1)));
  469. }
  470. }
  471. lua_settop (L, err_idx - 1);
  472. }
  473. static void
  474. rspamd_lua_call_on_complete (lua_State *L,
  475. struct rspamd_lua_process_cbdata *cbdata,
  476. const gchar *err_msg,
  477. const gchar *data, gsize datalen)
  478. {
  479. gint err_idx;
  480. lua_pushcfunction (L, &rspamd_lua_traceback);
  481. err_idx = lua_gettop (L);
  482. lua_rawgeti (L, LUA_REGISTRYINDEX, cbdata->cb_cbref);
  483. if (err_msg) {
  484. lua_pushstring (L, err_msg);
  485. }
  486. else {
  487. lua_pushnil (L);
  488. }
  489. if (data) {
  490. lua_pushlstring (L, data, datalen);
  491. }
  492. else {
  493. lua_pushnil (L);
  494. }
  495. if (lua_pcall (L, 2, 0, err_idx) != 0) {
  496. msg_err ("call to on_complete script failed: %s",
  497. lua_tostring (L, -1));
  498. }
  499. lua_settop (L, err_idx - 1);
  500. }
  501. static gboolean
  502. rspamd_lua_cld_handler (struct rspamd_worker_signal_handler *sigh, void *ud)
  503. {
  504. struct rspamd_lua_process_cbdata *cbdata = ud;
  505. struct rspamd_srv_command srv_cmd;
  506. lua_State *L;
  507. pid_t died;
  508. gint res = 0;
  509. /* Are we called by a correct children ? */
  510. died = waitpid (cbdata->cpid, &res, WNOHANG);
  511. if (died <= 0) {
  512. /* Wait more */
  513. return TRUE;
  514. }
  515. L = cbdata->L;
  516. msg_info ("handled SIGCHLD from %P", cbdata->cpid);
  517. if (!cbdata->replied) {
  518. /* We still need to call on_complete callback */
  519. ev_io_stop (cbdata->event_loop, &cbdata->ev);
  520. rspamd_lua_call_on_complete (cbdata->L, cbdata,
  521. "Worker has died without reply", NULL, 0);
  522. }
  523. /* Free structures */
  524. close (cbdata->sp[0]);
  525. luaL_unref (L, LUA_REGISTRYINDEX, cbdata->func_cbref);
  526. luaL_unref (L, LUA_REGISTRYINDEX, cbdata->cb_cbref);
  527. g_string_free (cbdata->io_buf, TRUE);
  528. if (cbdata->out_buf) {
  529. g_string_free (cbdata->out_buf, TRUE);
  530. }
  531. /* Notify main */
  532. memset (&srv_cmd, 0, sizeof (srv_cmd));
  533. srv_cmd.type = RSPAMD_SRV_ON_FORK;
  534. srv_cmd.cmd.on_fork.state = child_dead;
  535. srv_cmd.cmd.on_fork.cpid = cbdata->cpid;
  536. srv_cmd.cmd.on_fork.ppid = getpid ();
  537. rspamd_srv_send_command (cbdata->wrk, cbdata->event_loop, &srv_cmd, -1,
  538. NULL, NULL);
  539. g_free (cbdata);
  540. /* We are done with this SIGCHLD */
  541. return FALSE;
  542. }
  543. static void
  544. rspamd_lua_subprocess_io (EV_P_ ev_io *w, int revents)
  545. {
  546. struct rspamd_lua_process_cbdata *cbdata =
  547. (struct rspamd_lua_process_cbdata *)w->data;
  548. gssize r;
  549. if (cbdata->sz == (guint64)-1) {
  550. guint64 sz;
  551. /* We read size of reply + flags first */
  552. r = read (cbdata->sp[0], cbdata->io_buf->str + cbdata->io_buf->len,
  553. sizeof (guint64) - cbdata->io_buf->len);
  554. if (r == 0) {
  555. ev_io_stop (cbdata->event_loop, &cbdata->ev);
  556. rspamd_lua_call_on_complete (cbdata->L, cbdata,
  557. "Unexpected EOF", NULL, 0);
  558. cbdata->replied = TRUE;
  559. kill (cbdata->cpid, SIGTERM);
  560. return;
  561. }
  562. else if (r == -1) {
  563. if (errno == EAGAIN || errno == EINTR) {
  564. return;
  565. }
  566. else {
  567. ev_io_stop (cbdata->event_loop, &cbdata->ev);
  568. rspamd_lua_call_on_complete (cbdata->L, cbdata,
  569. strerror (errno), NULL, 0);
  570. cbdata->replied = TRUE;
  571. kill (cbdata->cpid, SIGTERM);
  572. return;
  573. }
  574. }
  575. cbdata->io_buf->len += r;
  576. if (cbdata->io_buf->len == sizeof (guint64)) {
  577. memcpy ((guchar *)&sz, cbdata->io_buf->str, sizeof (sz));
  578. if (sz & (1ULL << 63)) {
  579. cbdata->is_error = TRUE;
  580. sz &= ~(1ULL << 63);
  581. }
  582. cbdata->io_buf->len = 0;
  583. cbdata->sz = sz;
  584. g_string_set_size (cbdata->io_buf, sz + 1);
  585. cbdata->io_buf->len = 0;
  586. }
  587. }
  588. else {
  589. /* Read data */
  590. r = read (cbdata->sp[0], cbdata->io_buf->str + cbdata->io_buf->len,
  591. cbdata->sz - cbdata->io_buf->len);
  592. if (r == 0) {
  593. ev_io_stop (cbdata->event_loop, &cbdata->ev);
  594. rspamd_lua_call_on_complete (cbdata->L, cbdata,
  595. "Unexpected EOF", NULL, 0);
  596. cbdata->replied = TRUE;
  597. kill (cbdata->cpid, SIGTERM);
  598. return;
  599. }
  600. else if (r == -1) {
  601. if (errno == EAGAIN || errno == EINTR) {
  602. return;
  603. }
  604. else {
  605. ev_io_stop (cbdata->event_loop, &cbdata->ev);
  606. rspamd_lua_call_on_complete (cbdata->L, cbdata,
  607. strerror (errno), NULL, 0);
  608. cbdata->replied = TRUE;
  609. kill (cbdata->cpid, SIGTERM);
  610. return;
  611. }
  612. }
  613. cbdata->io_buf->len += r;
  614. if (cbdata->io_buf->len == cbdata->sz) {
  615. gchar rep[4];
  616. ev_io_stop (cbdata->event_loop, &cbdata->ev);
  617. /* Finished reading data */
  618. if (cbdata->is_error) {
  619. cbdata->io_buf->str[cbdata->io_buf->len] = '\0';
  620. rspamd_lua_call_on_complete (cbdata->L, cbdata,
  621. cbdata->io_buf->str, NULL, 0);
  622. }
  623. else {
  624. rspamd_lua_call_on_complete (cbdata->L, cbdata,
  625. NULL, cbdata->io_buf->str, cbdata->io_buf->len);
  626. }
  627. cbdata->replied = TRUE;
  628. /* Write reply to the child */
  629. rspamd_socket_blocking (cbdata->sp[0]);
  630. memset (rep, 0, sizeof (rep));
  631. (void) !write (cbdata->sp[0], rep, sizeof (rep));
  632. }
  633. }
  634. }
  635. static gint
  636. lua_worker_spawn_process (lua_State *L)
  637. {
  638. struct rspamd_worker *w = lua_check_worker (L, 1);
  639. struct rspamd_lua_process_cbdata *cbdata;
  640. struct rspamd_abstract_worker_ctx *actx;
  641. struct rspamd_srv_command srv_cmd;
  642. const gchar *cmdline = NULL, *input = NULL, *proctitle = NULL;
  643. gsize inputlen = 0;
  644. pid_t pid;
  645. GError *err = NULL;
  646. gint func_cbref, cb_cbref;
  647. if (!rspamd_lua_parse_table_arguments (L, 2, &err,
  648. RSPAMD_LUA_PARSE_ARGUMENTS_DEFAULT,
  649. "func=F;exec=S;stdin=V;*on_complete=F;proctitle=S", &func_cbref,
  650. &cmdline, &inputlen, &input, &cb_cbref, &proctitle)) {
  651. msg_err ("cannot get parameters list: %e", err);
  652. if (err) {
  653. g_error_free (err);
  654. }
  655. return 0;
  656. }
  657. cbdata = g_malloc0 (sizeof (*cbdata));
  658. cbdata->cb_cbref = cb_cbref;
  659. cbdata->func_cbref = func_cbref;
  660. if (input) {
  661. cbdata->out_buf = g_string_new_len (input, inputlen);
  662. cbdata->out_pos = 0;
  663. }
  664. if (rspamd_socketpair (cbdata->sp, SOCK_STREAM) == -1) {
  665. msg_err ("cannot spawn socketpair: %s", strerror (errno));
  666. luaL_unref (L, LUA_REGISTRYINDEX, cbdata->func_cbref);
  667. luaL_unref (L, LUA_REGISTRYINDEX, cbdata->cb_cbref);
  668. g_free (cbdata);
  669. return 0;
  670. }
  671. actx = w->ctx;
  672. cbdata->wrk = w;
  673. cbdata->L = L;
  674. cbdata->event_loop = actx->event_loop;
  675. cbdata->sz = (guint64)-1;
  676. pid = fork ();
  677. if (pid == -1) {
  678. msg_err ("cannot spawn process: %s", strerror (errno));
  679. close (cbdata->sp[0]);
  680. close (cbdata->sp[1]);
  681. luaL_unref (L, LUA_REGISTRYINDEX, cbdata->func_cbref);
  682. luaL_unref (L, LUA_REGISTRYINDEX, cbdata->cb_cbref);
  683. g_free (cbdata);
  684. return 0;
  685. }
  686. else if (pid == 0) {
  687. /* Child */
  688. gint rc;
  689. gchar inbuf[4];
  690. rspamd_log_on_fork (w->cf->type, w->srv->cfg, w->srv->logger);
  691. rc = ottery_init (w->srv->cfg->libs_ctx->ottery_cfg);
  692. if (rc != OTTERY_ERR_NONE) {
  693. msg_err ("cannot initialize PRNG: %d", rc);
  694. abort ();
  695. }
  696. rspamd_random_seed_fast ();
  697. #ifdef HAVE_EVUTIL_RNG_INIT
  698. evutil_secure_rng_init ();
  699. #endif
  700. close (cbdata->sp[0]);
  701. /* Here we assume that we can block on writing results */
  702. rspamd_socket_blocking (cbdata->sp[1]);
  703. g_hash_table_remove_all (w->signal_events);
  704. ev_loop_destroy (cbdata->event_loop);
  705. if (proctitle) {
  706. setproctitle ("lua process: %s", proctitle);
  707. }
  708. else {
  709. setproctitle ("lua process: unnamed");
  710. }
  711. cbdata->event_loop = ev_loop_new (EVFLAG_SIGNALFD);
  712. rspamd_worker_unblock_signals ();
  713. rspamd_lua_execute_lua_subprocess (L, cbdata);
  714. /* Wait for parent to reply and exit */
  715. rc = read (cbdata->sp[1], inbuf, sizeof (inbuf));
  716. if (rc >= sizeof (inbuf) &&
  717. memcmp (inbuf, "\0\0\0\0", sizeof (inbuf)) == 0) {
  718. exit (EXIT_SUCCESS);
  719. }
  720. else {
  721. msg_err ("got invalid reply from parent");
  722. exit (EXIT_FAILURE);
  723. }
  724. }
  725. cbdata->cpid = pid;
  726. cbdata->io_buf = g_string_sized_new (8);
  727. /* Notify main */
  728. memset (&srv_cmd, 0, sizeof (srv_cmd));
  729. srv_cmd.type = RSPAMD_SRV_ON_FORK;
  730. srv_cmd.cmd.on_fork.state = child_create;
  731. srv_cmd.cmd.on_fork.cpid = pid;
  732. srv_cmd.cmd.on_fork.ppid = getpid ();
  733. rspamd_srv_send_command (w, cbdata->event_loop, &srv_cmd, -1, NULL, NULL);
  734. close (cbdata->sp[1]);
  735. rspamd_socket_nonblocking (cbdata->sp[0]);
  736. /* Parent */
  737. rspamd_worker_set_signal_handler (SIGCHLD, w, cbdata->event_loop,
  738. rspamd_lua_cld_handler,
  739. cbdata);
  740. /* Add result pipe waiting */
  741. ev_io_init (&cbdata->ev, rspamd_lua_subprocess_io, cbdata->sp[0], EV_READ);
  742. cbdata->ev.data = cbdata;
  743. ev_io_start (cbdata->event_loop, &cbdata->ev);
  744. return 0;
  745. }
  746. void
  747. luaopen_worker (lua_State * L)
  748. {
  749. rspamd_lua_new_class (L, "rspamd{worker}", worker_reg);
  750. }