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.

fuzzy_check.c 23KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726
  1. /*
  2. * Copyright (c) 2009, Rambler media
  3. * All rights reserved.
  4. *
  5. * Redistribution and use in source and binary forms, with or without
  6. * modification, are permitted provided that the following conditions are met:
  7. * * Redistributions of source code must retain the above copyright
  8. * notice, this list of conditions and the following disclaimer.
  9. * * Redistributions in binary form must reproduce the above copyright
  10. * notice, this list of conditions and the following disclaimer in the
  11. * documentation and/or other materials provided with the distribution.
  12. *
  13. * THIS SOFTWARE IS PROVIDED BY Rambler media ''AS IS'' AND ANY
  14. * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  15. * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  16. * DISCLAIMED. IN NO EVENT SHALL Rambler BE LIABLE FOR ANY
  17. * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  18. * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  19. * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
  20. * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  21. * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  22. * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  23. */
  24. /***MODULE:fuzzy
  25. * rspamd module that checks fuzzy checksums for messages
  26. *
  27. * Allowed options:
  28. * - symbol (string): symbol to insert (default: 'R_FUZZY')
  29. * - max_score (double): maximum score to that weights of hashes would be normalized (default: 0 - no normalization)
  30. *
  31. * - fuzzy_map (string): a string that contains map in format { fuzzy_key => [ symbol, weight ] } where fuzzy_key is number of
  32. * fuzzy list. This string itself should be in format 1:R_FUZZY_SAMPLE1:10,2:R_FUZZY_SAMPLE2:1 etc, where first number is fuzzy
  33. * key, second is symbol to insert and third - weight for normalization
  34. *
  35. * - min_length (integer): minimum length (in characters) for text part to be checked for fuzzy hash (default: 0 - no limit)
  36. * - whitelist (map string): map of ip addresses that should not be checked with this module
  37. * - servers (string): list of fuzzy servers in format "server1:port,server2:port" - these servers would be used for checking and storing
  38. * fuzzy hashes
  39. */
  40. #include "../config.h"
  41. #include "../main.h"
  42. #include "../message.h"
  43. #include "../modules.h"
  44. #include "../cfg_file.h"
  45. #include "../expressions.h"
  46. #include "../util.h"
  47. #include "../view.h"
  48. #include "../map.h"
  49. #include "../fuzzy_storage.h"
  50. #define DEFAULT_SYMBOL "R_FUZZY_HASH"
  51. #define DEFAULT_UPSTREAM_ERROR_TIME 10
  52. #define DEFAULT_UPSTREAM_DEAD_TIME 300
  53. #define DEFAULT_UPSTREAM_MAXERRORS 10
  54. #define IO_TIMEOUT 5
  55. #define DEFAULT_PORT 11335
  56. struct storage_server {
  57. struct upstream up;
  58. char *name;
  59. struct in_addr addr;
  60. uint16_t port;
  61. };
  62. struct fuzzy_mapping {
  63. uint32_t fuzzy_flag;
  64. char *symbol;
  65. double weight;
  66. };
  67. struct fuzzy_ctx {
  68. int (*filter) (struct worker_task * task);
  69. char *symbol;
  70. struct storage_server *servers;
  71. int servers_num;
  72. memory_pool_t *fuzzy_pool;
  73. double max_score;
  74. uint32_t min_hash_len;
  75. radix_tree_t *whitelist;
  76. GHashTable *mappings;
  77. };
  78. struct fuzzy_client_session {
  79. int state;
  80. fuzzy_hash_t *h;
  81. struct event ev;
  82. struct timeval tv;
  83. struct worker_task *task;
  84. struct storage_server *server;
  85. int fd;
  86. };
  87. struct fuzzy_learn_session {
  88. struct event ev;
  89. fuzzy_hash_t *h;
  90. int cmd;
  91. int value;
  92. int flag;
  93. int *saved;
  94. struct timeval tv;
  95. struct controller_session *session;
  96. struct storage_server *server;
  97. struct worker_task *task;
  98. int fd;
  99. };
  100. static struct fuzzy_ctx *fuzzy_module_ctx = NULL;
  101. static int fuzzy_mime_filter (struct worker_task *task);
  102. static void fuzzy_symbol_callback (struct worker_task *task, void *unused);
  103. static void fuzzy_add_handler (char **args, struct controller_session *session);
  104. static void fuzzy_delete_handler (char **args, struct controller_session *session);
  105. /* Flags string is in format <numeric_flag>:<SYMBOL>:weight[, <numeric_flag>:<SYMBOL>:weight...] */
  106. static void
  107. parse_flags_string (char *str)
  108. {
  109. char **strvec, *item, *err_str, **map_str;
  110. int num, i, t;
  111. struct fuzzy_mapping *map;
  112. strvec = g_strsplit (str, ", ;", 0);
  113. num = g_strv_length (strvec);
  114. for (i = 0; i < num; i ++) {
  115. item = strvec[i];
  116. map_str = g_strsplit (item, ":", 3);
  117. t = g_strv_length (map_str);
  118. if (t != 3 && t != 2) {
  119. msg_err ("invalid fuzzy mapping: %s", item);
  120. }
  121. else {
  122. map = memory_pool_alloc (fuzzy_module_ctx->fuzzy_pool, sizeof (struct fuzzy_mapping));
  123. map->symbol = memory_pool_strdup (fuzzy_module_ctx->fuzzy_pool, map_str[1]);
  124. errno = 0;
  125. map->fuzzy_flag = strtol (map_str[0], &err_str, 10);
  126. if (errno != 0 || (err_str && *err_str != '\0')) {
  127. msg_info ("cannot parse flag %s: %s", map_str[0], strerror (errno));
  128. }
  129. else if (t == 2) {
  130. /* Weight is skipped in definition */
  131. map->weight = fuzzy_module_ctx->max_score;
  132. }
  133. else {
  134. map->weight = strtol (map_str[2], &err_str, 10);
  135. /* Add flag to hash table */
  136. g_hash_table_insert (fuzzy_module_ctx->mappings, GINT_TO_POINTER(map->fuzzy_flag), map);
  137. }
  138. }
  139. g_strfreev (map_str);
  140. }
  141. g_strfreev (strvec);
  142. }
  143. static void
  144. parse_servers_string (char *str)
  145. {
  146. char **strvec, *p, portbuf[6], *name;
  147. int num, i, j, port;
  148. struct hostent *hent;
  149. struct in_addr addr;
  150. strvec = g_strsplit (str, ",", 0);
  151. num = g_strv_length (strvec);
  152. fuzzy_module_ctx->servers = memory_pool_alloc0 (fuzzy_module_ctx->fuzzy_pool, sizeof (struct storage_server) * num);
  153. for (i = 0; i < num; i++) {
  154. g_strstrip (strvec[i]);
  155. if ((p = strchr (strvec[i], ':')) != NULL) {
  156. j = 0;
  157. p++;
  158. while (g_ascii_isdigit (*(p + j)) && j < sizeof (portbuf) - 1) {
  159. portbuf[j] = *(p + j);
  160. j++;
  161. }
  162. portbuf[j] = '\0';
  163. port = atoi (portbuf);
  164. }
  165. else {
  166. /* Default http port */
  167. port = DEFAULT_PORT;
  168. }
  169. name = memory_pool_alloc (fuzzy_module_ctx->fuzzy_pool, p - strvec[i]);
  170. g_strlcpy (name, strvec[i], p - strvec[i]);
  171. if (!inet_aton (name, &addr)) {
  172. /* Resolve using dns */
  173. hent = gethostbyname (name);
  174. if (hent == NULL) {
  175. msg_info ("cannot resolve: %s", name);
  176. continue;
  177. }
  178. else {
  179. fuzzy_module_ctx->servers[fuzzy_module_ctx->servers_num].port = port;
  180. fuzzy_module_ctx->servers[fuzzy_module_ctx->servers_num].name = name;
  181. memcpy (&fuzzy_module_ctx->servers[fuzzy_module_ctx->servers_num].addr, hent->h_addr, sizeof (struct in_addr));
  182. fuzzy_module_ctx->servers_num++;
  183. }
  184. }
  185. else {
  186. fuzzy_module_ctx->servers[fuzzy_module_ctx->servers_num].port = port;
  187. fuzzy_module_ctx->servers[fuzzy_module_ctx->servers_num].name = name;
  188. memcpy (&fuzzy_module_ctx->servers[fuzzy_module_ctx->servers_num].addr, hent->h_addr, sizeof (struct in_addr));
  189. fuzzy_module_ctx->servers_num++;
  190. }
  191. }
  192. g_strfreev (strvec);
  193. }
  194. static double
  195. fuzzy_normalize (int32_t in, double weight)
  196. {
  197. double ms = weight, ams = fabs (ms), ain = fabs (in);
  198. if (ams > 0.001) {
  199. if (ain < ams / 2.) {
  200. return in;
  201. }
  202. else if (ain < ams * 2.) {
  203. ain = ain / 3. + ams / 3.;
  204. return in > 0 ? ain : -(ain);
  205. }
  206. else {
  207. return in > 0 ? ms : -(ms);
  208. }
  209. }
  210. return (double)in;
  211. }
  212. int
  213. fuzzy_check_module_init (struct config_file *cfg, struct module_ctx **ctx)
  214. {
  215. fuzzy_module_ctx = g_malloc (sizeof (struct fuzzy_ctx));
  216. fuzzy_module_ctx->filter = fuzzy_mime_filter;
  217. fuzzy_module_ctx->fuzzy_pool = memory_pool_new (memory_pool_get_size ());
  218. fuzzy_module_ctx->servers = NULL;
  219. fuzzy_module_ctx->servers_num = 0;
  220. fuzzy_module_ctx->mappings = g_hash_table_new (g_direct_hash, g_direct_equal);
  221. *ctx = (struct module_ctx *)fuzzy_module_ctx;
  222. return 0;
  223. }
  224. int
  225. fuzzy_check_module_config (struct config_file *cfg)
  226. {
  227. char *value;
  228. int res = TRUE;
  229. if ((value = get_module_opt (cfg, "fuzzy_check", "symbol")) != NULL) {
  230. fuzzy_module_ctx->symbol = memory_pool_strdup (fuzzy_module_ctx->fuzzy_pool, value);
  231. }
  232. else {
  233. fuzzy_module_ctx->symbol = DEFAULT_SYMBOL;
  234. }
  235. if ((value = get_module_opt (cfg, "fuzzy_check", "max_score")) != NULL) {
  236. fuzzy_module_ctx->max_score = strtod (value, NULL);
  237. }
  238. else {
  239. fuzzy_module_ctx->max_score = 0.;
  240. }
  241. if ((value = get_module_opt (cfg, "fuzzy_check", "min_length")) != NULL) {
  242. fuzzy_module_ctx->min_hash_len = strtoul (value, NULL, 10);
  243. }
  244. else {
  245. fuzzy_module_ctx->min_hash_len = 0.;
  246. }
  247. if ((value = get_module_opt (cfg, "fuzzy_check", "whitelist")) != NULL) {
  248. fuzzy_module_ctx->whitelist = radix_tree_create ();
  249. if (!add_map (value, read_radix_list, fin_radix_list, (void **)&fuzzy_module_ctx->whitelist)) {
  250. msg_err ("cannot add whitelist '%s'", value);
  251. }
  252. }
  253. else {
  254. fuzzy_module_ctx->whitelist = NULL;
  255. }
  256. if ((value = get_module_opt (cfg, "fuzzy_check", "servers")) != NULL) {
  257. parse_servers_string (value);
  258. }
  259. if ((value = get_module_opt (cfg, "fuzzy_check", "fuzzy_map")) != NULL) {
  260. parse_flags_string (value);
  261. }
  262. register_symbol (&cfg->cache, fuzzy_module_ctx->symbol, fuzzy_module_ctx->max_score, fuzzy_symbol_callback, NULL);
  263. register_custom_controller_command ("fuzzy_add", fuzzy_add_handler, TRUE, TRUE);
  264. register_custom_controller_command ("fuzzy_del", fuzzy_delete_handler, TRUE, TRUE);
  265. return res;
  266. }
  267. int
  268. fuzzy_check_module_reconfig (struct config_file *cfg)
  269. {
  270. memory_pool_delete (fuzzy_module_ctx->fuzzy_pool);
  271. fuzzy_module_ctx->servers = NULL;
  272. fuzzy_module_ctx->servers_num = 0;
  273. fuzzy_module_ctx->fuzzy_pool = memory_pool_new (memory_pool_get_size ());
  274. g_hash_table_remove_all (fuzzy_module_ctx->mappings);
  275. return fuzzy_check_module_config (cfg);
  276. }
  277. /* Finalize IO */
  278. static void
  279. fuzzy_io_fin (void *ud)
  280. {
  281. struct fuzzy_client_session *session = ud;
  282. event_del (&session->ev);
  283. close (session->fd);
  284. session->task->save.saved--;
  285. if (session->task->save.saved == 0) {
  286. /* Call other filters */
  287. session->task->save.saved = 1;
  288. process_filters (session->task);
  289. }
  290. }
  291. /* Call this whenever we got data from fuzzy storage */
  292. static void
  293. fuzzy_io_callback (int fd, short what, void *arg)
  294. {
  295. struct fuzzy_client_session *session = arg;
  296. struct fuzzy_cmd cmd;
  297. struct fuzzy_mapping *map;
  298. char buf[62], *err_str, *symbol;
  299. int value = 0, flag = 0, r;
  300. double nval;
  301. if (what == EV_WRITE) {
  302. /* Send command to storage */
  303. cmd.blocksize = session->h->block_size;
  304. cmd.value = 0;
  305. memcpy (cmd.hash, session->h->hash_pipe, sizeof (cmd.hash));
  306. cmd.cmd = FUZZY_CHECK;
  307. if (write (fd, &cmd, sizeof (struct fuzzy_cmd)) == -1) {
  308. goto err;
  309. }
  310. else {
  311. event_set (&session->ev, fd, EV_READ, fuzzy_io_callback, session);
  312. event_add (&session->ev, &session->tv);
  313. }
  314. }
  315. else if (what == EV_READ) {
  316. /* Got reply */
  317. if ((r = read (fd, buf, sizeof (buf) - 1)) == -1) {
  318. goto err;
  319. }
  320. else if (buf[0] == 'O' && buf[1] == 'K') {
  321. buf[r] = 0;
  322. /* Now try to get value */
  323. value = strtol (buf + 3, &err_str, 10);
  324. if (*err_str == ' ') {
  325. /* Now read flag */
  326. flag = strtol (err_str + 1, &err_str, 10);
  327. }
  328. *err_str = '\0';
  329. /* Get mapping by flag */
  330. if ((map = g_hash_table_lookup (fuzzy_module_ctx->mappings, GINT_TO_POINTER (flag))) == NULL) {
  331. /* Default symbol and default weight */
  332. symbol = fuzzy_module_ctx->symbol;
  333. nval = fuzzy_normalize (value, fuzzy_module_ctx->max_score);
  334. }
  335. else {
  336. /* Get symbol and weight from map */
  337. symbol = map->symbol;
  338. nval = fuzzy_normalize (value, map->weight);
  339. }
  340. snprintf (buf, sizeof (buf), "%d: %d / %.2f", flag, value, nval);
  341. insert_result (session->task, symbol, nval, g_list_prepend (NULL,
  342. memory_pool_strdup (session->task->task_pool, buf)));
  343. }
  344. goto ok;
  345. }
  346. else {
  347. errno = ETIMEDOUT;
  348. goto err;
  349. }
  350. return;
  351. err:
  352. msg_err ("got error on IO with server %s:%d, %d, %s", session->server->name, session->server->port, errno, strerror (errno));
  353. ok:
  354. remove_normal_event (session->task->s, fuzzy_io_fin, session);
  355. }
  356. static void
  357. fuzzy_learn_fin (void *arg)
  358. {
  359. struct fuzzy_learn_session *session = arg;
  360. event_del (&session->ev);
  361. close (session->fd);
  362. (*session->saved)--;
  363. if (*session->saved == 0) {
  364. session->session->state = STATE_REPLY;
  365. session->session->dispatcher->write_callback (session);
  366. }
  367. }
  368. static void
  369. fuzzy_learn_callback (int fd, short what, void *arg)
  370. {
  371. struct fuzzy_learn_session *session = arg;
  372. struct fuzzy_cmd cmd;
  373. char buf[64];
  374. int r;
  375. if (what == EV_WRITE) {
  376. /* Send command to storage */
  377. cmd.blocksize = session->h->block_size;
  378. memcpy (cmd.hash, session->h->hash_pipe, sizeof (cmd.hash));
  379. cmd.cmd = session->cmd;
  380. cmd.value = session->value;
  381. cmd.flag = session->flag;
  382. if (write (fd, &cmd, sizeof (struct fuzzy_cmd)) == -1) {
  383. goto err;
  384. }
  385. else {
  386. event_set (&session->ev, fd, EV_READ, fuzzy_learn_callback, session);
  387. event_add (&session->ev, &session->tv);
  388. }
  389. }
  390. else if (what == EV_READ) {
  391. if (read (fd, buf, sizeof (buf)) == -1) {
  392. goto err;
  393. }
  394. else if (buf[0] == 'O' && buf[1] == 'K') {
  395. r = snprintf (buf, sizeof (buf), "OK" CRLF);
  396. rspamd_dispatcher_write (session->session->dispatcher, buf, r, FALSE, FALSE);
  397. goto ok;
  398. }
  399. else {
  400. r = snprintf (buf, sizeof (buf), "ERR" CRLF);
  401. rspamd_dispatcher_write (session->session->dispatcher, buf, r, FALSE, FALSE);
  402. goto ok;
  403. }
  404. }
  405. else {
  406. errno = ETIMEDOUT;
  407. goto err;
  408. }
  409. return;
  410. err:
  411. msg_err ("got error in IO with server %s:%d, %d, %s", session->server->name, session->server->port, errno, strerror (errno));
  412. r = snprintf (buf, sizeof (buf), "Error" CRLF);
  413. rspamd_dispatcher_write (session->session->dispatcher, buf, r, FALSE, FALSE);
  414. ok:
  415. remove_normal_event (session->session->s, fuzzy_learn_fin, session);
  416. }
  417. /* This callback is called when we check message via fuzzy hashes storage */
  418. static void
  419. fuzzy_symbol_callback (struct worker_task *task, void *unused)
  420. {
  421. struct mime_text_part *part;
  422. struct fuzzy_client_session *session;
  423. struct storage_server *selected;
  424. GList *cur;
  425. int sock;
  426. /* Check whitelist */
  427. if (fuzzy_module_ctx->whitelist && task->from_addr.s_addr != 0) {
  428. if (radix32tree_find (fuzzy_module_ctx->whitelist, ntohl ((uint32_t) task->from_addr.s_addr)) != RADIX_NO_VALUE) {
  429. msg_info ("address %s is whitelisted, skip fuzzy check", inet_ntoa (task->from_addr));
  430. return;
  431. }
  432. }
  433. cur = task->text_parts;
  434. while (cur) {
  435. part = cur->data;
  436. if (part->is_empty) {
  437. cur = g_list_next (cur);
  438. continue;
  439. }
  440. /* Check length of hash */
  441. if (fuzzy_module_ctx->min_hash_len != 0 &&
  442. strlen (part->fuzzy->hash_pipe) * part->fuzzy->block_size < fuzzy_module_ctx->min_hash_len) {
  443. msg_info ("part hash is shorter than %d symbols, skip fuzzy check", fuzzy_module_ctx->min_hash_len);
  444. cur = g_list_next (cur);
  445. continue;
  446. }
  447. /* Get upstream */
  448. #ifdef HAVE_CLOCK_GETTIME
  449. selected = (struct storage_server *)get_upstream_by_hash (fuzzy_module_ctx->servers, fuzzy_module_ctx->servers_num,
  450. sizeof (struct storage_server), task->ts.tv_sec,
  451. DEFAULT_UPSTREAM_ERROR_TIME, DEFAULT_UPSTREAM_DEAD_TIME, DEFAULT_UPSTREAM_MAXERRORS, part->fuzzy->hash_pipe, sizeof (part->fuzzy->hash_pipe));
  452. #else
  453. selected = (struct storage_server *)get_upstream_by_hash (fuzzy_module_ctx->servers, fuzzy_module_ctx->servers_num,
  454. sizeof (struct storage_server), task->tv.tv_sec,
  455. DEFAULT_UPSTREAM_ERROR_TIME, DEFAULT_UPSTREAM_DEAD_TIME, DEFAULT_UPSTREAM_MAXERRORS, part->fuzzy->hash_pipe, sizeof (part->fuzzy->hash_pipe));
  456. #endif
  457. if (selected) {
  458. if ((sock = make_udp_socket (&selected->addr, selected->port, FALSE, TRUE)) == -1) {
  459. msg_warn ("cannot connect to %s, %d, %s", selected->name, errno, strerror (errno));
  460. }
  461. else {
  462. /* Create session for a socket */
  463. session = memory_pool_alloc (task->task_pool, sizeof (struct fuzzy_client_session));
  464. event_set (&session->ev, sock, EV_WRITE, fuzzy_io_callback, session);
  465. session->tv.tv_sec = IO_TIMEOUT;
  466. session->tv.tv_usec = 0;
  467. session->state = 0;
  468. session->h = part->fuzzy;
  469. session->task = task;
  470. session->fd = sock;
  471. session->server = selected;
  472. event_add (&session->ev, &session->tv);
  473. register_async_event (task->s, fuzzy_io_fin, session, FALSE);
  474. task->save.saved++;
  475. }
  476. }
  477. cur = g_list_next (cur);
  478. }
  479. }
  480. static void
  481. fuzzy_process_handler (struct controller_session *session, f_str_t * in)
  482. {
  483. struct worker_task *task;
  484. struct fuzzy_learn_session *s;
  485. struct mime_text_part *part;
  486. struct storage_server *selected;
  487. GList *cur;
  488. int sock, r, cmd = 0, value = 0, flag = 0, *saved, *sargs;
  489. char out_buf[BUFSIZ];
  490. /* Extract arguments */
  491. if (session->other_data) {
  492. sargs = session->other_data;
  493. cmd = sargs[0];
  494. value = sargs[1];
  495. flag = sargs[2];
  496. }
  497. /* Prepare task */
  498. task = construct_task (session->worker);
  499. session->other_data = task;
  500. session->state = STATE_WAIT;
  501. /* Allocate message from string */
  502. task->msg = memory_pool_alloc (task->task_pool, sizeof (f_str_t));
  503. task->msg->begin = in->begin;
  504. task->msg->len = in->len;
  505. saved = memory_pool_alloc0 (session->session_pool, sizeof (int));
  506. r = process_message (task);
  507. if (r == -1) {
  508. msg_warn ("processing of message failed");
  509. free_task (task, FALSE);
  510. session->state = STATE_REPLY;
  511. r = snprintf (out_buf, sizeof (out_buf), "cannot process message" CRLF);
  512. rspamd_dispatcher_write (session->dispatcher, out_buf, r, FALSE, FALSE);
  513. return;
  514. }
  515. else {
  516. /* Plan new event for writing */
  517. cur = task->text_parts;
  518. while (cur) {
  519. part = cur->data;
  520. if (part->is_empty) {
  521. cur = g_list_next (cur);
  522. continue;
  523. }
  524. /* Get upstream */
  525. #ifdef HAVE_CLOCK_GETTIME
  526. selected = (struct storage_server *)get_upstream_by_hash (fuzzy_module_ctx->servers, fuzzy_module_ctx->servers_num,
  527. sizeof (struct storage_server), task->ts.tv_sec,
  528. DEFAULT_UPSTREAM_ERROR_TIME, DEFAULT_UPSTREAM_DEAD_TIME, DEFAULT_UPSTREAM_MAXERRORS, part->fuzzy->hash_pipe, sizeof (part->fuzzy->hash_pipe));
  529. #else
  530. selected = (struct storage_server *)get_upstream_by_hash (fuzzy_module_ctx->servers, fuzzy_module_ctx->servers_num,
  531. sizeof (struct storage_server), task->tv.tv_sec,
  532. DEFAULT_UPSTREAM_ERROR_TIME, DEFAULT_UPSTREAM_DEAD_TIME, DEFAULT_UPSTREAM_MAXERRORS, part->fuzzy->hash_pipe, sizeof (part->fuzzy->hash_pipe));
  533. #endif
  534. if (selected) {
  535. /* Create UDP socket */
  536. if ((sock = make_udp_socket (&selected->addr, selected->port, FALSE, TRUE)) == -1) {
  537. msg_warn ("cannot connect to %s, %d, %s", selected->name, errno, strerror (errno));
  538. session->state = STATE_REPLY;
  539. r = snprintf (out_buf, sizeof (out_buf), "no hashes written" CRLF);
  540. rspamd_dispatcher_write (session->dispatcher, out_buf, r, FALSE, FALSE);
  541. free_task (task, FALSE);
  542. return;
  543. }
  544. else {
  545. /* Socket is made, create session */
  546. s = memory_pool_alloc (session->session_pool, sizeof (struct fuzzy_learn_session));
  547. event_set (&s->ev, sock, EV_WRITE, fuzzy_learn_callback, s);
  548. s->tv.tv_sec = IO_TIMEOUT;
  549. s->tv.tv_usec = 0;
  550. s->task = task;
  551. s->h = memory_pool_alloc (session->session_pool, sizeof (fuzzy_hash_t));
  552. memcpy (s->h, part->fuzzy, sizeof (fuzzy_hash_t));
  553. s->session = session;
  554. s->server = selected;
  555. s->cmd = cmd;
  556. s->value = value;
  557. s->flag = flag;
  558. s->saved = saved;
  559. s->fd = sock;
  560. event_add (&s->ev, &s->tv);
  561. (*saved)++;
  562. register_async_event (session->s, fuzzy_learn_fin, s, FALSE);
  563. }
  564. }
  565. else {
  566. /* Cannot write hash */
  567. session->state = STATE_REPLY;
  568. r = snprintf (out_buf, sizeof (out_buf), "cannot write fuzzy hash" CRLF);
  569. rspamd_dispatcher_write (session->dispatcher, out_buf, r, FALSE, FALSE);
  570. free_task (task, FALSE);
  571. return;
  572. }
  573. cur = g_list_next (cur);
  574. }
  575. }
  576. free_task (task, FALSE);
  577. if (*saved == 0) {
  578. session->state = STATE_REPLY;
  579. r = snprintf (out_buf, sizeof (out_buf), "no hashes written" CRLF);
  580. rspamd_dispatcher_write (session->dispatcher, out_buf, r, FALSE, FALSE);
  581. }
  582. }
  583. static void
  584. fuzzy_controller_handler (char **args, struct controller_session *session, int cmd)
  585. {
  586. char *arg, out_buf[BUFSIZ], *err_str;
  587. uint32_t size;
  588. int r, value = 1, flag = 0, *sargs;
  589. /* Process size */
  590. arg = args[0];
  591. if (!arg || *arg == '\0') {
  592. msg_info ("empty content length");
  593. r = snprintf (out_buf, sizeof (out_buf), "fuzzy command requires length as argument" CRLF);
  594. rspamd_dispatcher_write (session->dispatcher, out_buf, r, FALSE, FALSE);
  595. session->state = STATE_REPLY;
  596. return;
  597. }
  598. errno = 0;
  599. size = strtoul (arg, &err_str, 10);
  600. if (errno != 0 || (err_str && *err_str != '\0')) {
  601. r = snprintf (out_buf, sizeof (out_buf), "learn size is invalid" CRLF);
  602. rspamd_dispatcher_write (session->dispatcher, out_buf, r, FALSE, FALSE);
  603. session->state = STATE_REPLY;
  604. return;
  605. }
  606. /* Process value */
  607. arg = args[1];
  608. if (arg && *arg != '\0') {
  609. errno = 0;
  610. value = strtol (arg, &err_str, 10);
  611. if (errno != 0 || *err_str != '\0') {
  612. msg_info ("error converting numeric argument %s", arg);
  613. value = 0;
  614. }
  615. }
  616. /* Process flag */
  617. arg = args[2];
  618. if (arg && *arg != '\0') {
  619. errno = 0;
  620. flag = strtol (arg, &err_str, 10);
  621. if (errno != 0 || *err_str != '\0') {
  622. msg_info ("error converting numeric argument %s", arg);
  623. flag = 0;
  624. }
  625. }
  626. session->state = STATE_OTHER;
  627. rspamd_set_dispatcher_policy (session->dispatcher, BUFFER_CHARACTER, size);
  628. session->other_handler = fuzzy_process_handler;
  629. /* Prepare args */
  630. sargs = memory_pool_alloc (session->session_pool, sizeof (int) * 3);
  631. sargs[0] = cmd;
  632. sargs[1] = value;
  633. sargs[2] = flag;
  634. session->other_data = sargs;
  635. }
  636. static void
  637. fuzzy_add_handler (char **args, struct controller_session *session)
  638. {
  639. fuzzy_controller_handler (args, session, FUZZY_WRITE);
  640. }
  641. static void
  642. fuzzy_delete_handler (char **args, struct controller_session *session)
  643. {
  644. fuzzy_controller_handler (args, session, FUZZY_DEL);
  645. }
  646. static int
  647. fuzzy_mime_filter (struct worker_task *task)
  648. {
  649. /* XXX: remove this */
  650. return 0;
  651. }