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_tcp.c 56KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441
  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. #include "lua_thread_pool.h"
  18. #include "libserver/ssl_util.h"
  19. #include "utlist.h"
  20. #include "unix-std.h"
  21. #include <math.h>
  22. static const gchar *M = "rspamd lua tcp";
  23. /***
  24. * @module rspamd_tcp
  25. * Rspamd TCP module represents generic TCP asynchronous client available from LUA code.
  26. * This module hides all complexity: DNS resolving, sessions management, zero-copy
  27. * text transfers and so on under the hood. It can work in partial or complete modes:
  28. *
  29. * - partial mode is used when you need to call a continuation routine each time data is available for read
  30. * - complete mode calls for continuation merely when all data is read from socket (e.g. when a server sends reply and closes a connection)
  31. * @example
  32. local logger = require "rspamd_logger"
  33. local tcp = require "rspamd_tcp"
  34. rspamd_config.SYM = function(task)
  35. local function cb(err, data)
  36. logger.infox('err: %1, data: %2', err, tostring(data))
  37. end
  38. tcp.request({
  39. task = task,
  40. host = "google.com",
  41. port = 80,
  42. data = {"GET / HTTP/1.0\r\n", "Host: google.com\r\n", "\r\n"},
  43. callback = cb})
  44. end
  45. -- New TCP syntax test
  46. rspamd_config:register_symbol({
  47. name = 'TCP_TEST',
  48. type = "normal",
  49. callback = function(task)
  50. local logger = require "rspamd_logger"
  51. local function rcpt_done_cb(err, data, conn)
  52. logger.errx(task, 'RCPT: got reply: %s, error: %s', data, err)
  53. conn:close()
  54. end
  55. local function rcpt_cb(err, conn)
  56. logger.errx(task, 'written rcpt, error: %s', err)
  57. conn:add_read(rcpt_done_cb, '\r\n')
  58. end
  59. local function from_done_cb(err, data, conn)
  60. logger.errx(task, 'FROM: got reply: %s, error: %s', data, err)
  61. conn:add_write(rcpt_cb, 'RCPT TO: <test@yandex.ru>\r\n')
  62. end
  63. local function from_cb(err, conn)
  64. logger.errx(task, 'written from, error: %s', err)
  65. conn:add_read(from_done_cb, '\r\n')
  66. end
  67. local function hello_done_cb(err, data, conn)
  68. logger.errx(task, 'HELO: got reply: %s, error: %s', data, err)
  69. conn:add_write(from_cb, 'MAIL FROM: <>\r\n')
  70. end
  71. local function hello_cb(err, conn)
  72. logger.errx(task, 'written hello, error: %s', err)
  73. conn:add_read(hello_done_cb, '\r\n')
  74. end
  75. local function init_cb(err, data, conn)
  76. logger.errx(task, 'got reply: %s, error: %s', data, err)
  77. conn:add_write(hello_cb, 'HELO example.com\r\n')
  78. end
  79. tcp.request{
  80. task = task,
  81. callback = init_cb,
  82. stop_pattern = '\r\n',
  83. host = 'mx.yandex.ru',
  84. port = 25
  85. }
  86. end,
  87. priority = 10,
  88. })
  89. */
  90. LUA_FUNCTION_DEF (tcp, request);
  91. /***
  92. * @function rspamd_tcp.connect_sync()
  93. *
  94. * Creates pseudo-synchronous TCP connection.
  95. * Each method of the connection requiring IO, becomes a yielding point,
  96. * i.e. current thread Lua thread is get suspended and resumes as soon as IO is done
  97. *
  98. * This class represents low-level API, using of "lua_tcp_sync" module is recommended.
  99. *
  100. * @example
  101. local rspamd_tcp = require "rspamd_tcp"
  102. local logger = require "rspamd_logger"
  103. local function http_simple_tcp_symbol(task)
  104. local err
  105. local is_ok, connection = rspamd_tcp.connect_sync {
  106. task = task,
  107. host = '127.0.0.1',
  108. timeout = 20,
  109. port = 18080,
  110. ssl = false, -- If SSL connection is needed
  111. ssl_verify = true, -- set to false if verify is not needed
  112. }
  113. is_ok, err = connection:write('GET /request_sync HTTP/1.1\r\nConnection: keep-alive\r\n\r\n')
  114. logger.errx(task, 'write %1, %2', is_ok, err)
  115. if not is_ok then
  116. logger.errx(task, 'write error: %1', err)
  117. end
  118. local data
  119. is_ok, data = connection:read_once();
  120. logger.errx(task, 'read_once: is_ok: %1, data: %2', is_ok, data)
  121. is_ok, err = connection:write("POST /request2 HTTP/1.1\r\n\r\n")
  122. logger.errx(task, 'write[2] %1, %2', is_ok, err)
  123. is_ok, data = connection:read_once();
  124. logger.errx(task, 'read_once[2]: is_ok %1, data: %2', is_ok, data)
  125. connection:close()
  126. end
  127. rspamd_config:register_symbol({
  128. name = 'SIMPLE_TCP_TEST',
  129. score = 1.0,
  130. callback = http_simple_tcp_symbol,
  131. no_squeeze = true
  132. })
  133. *
  134. */
  135. LUA_FUNCTION_DEF (tcp, connect_sync);
  136. /***
  137. * @method tcp:close()
  138. *
  139. * Closes TCP connection
  140. */
  141. LUA_FUNCTION_DEF (tcp, close);
  142. /***
  143. * @method tcp:add_read(callback, [pattern])
  144. *
  145. * Adds new read event to the tcp connection
  146. * @param {function} callback to be called when data is read
  147. * @param {string} pattern optional stop pattern
  148. */
  149. LUA_FUNCTION_DEF (tcp, add_read);
  150. /***
  151. * @method tcp:add_write(callback, data)
  152. *
  153. * Adds new write event to the tcp connection
  154. * @param {function} optional callback to be called when data is completely written
  155. * @param {table/string/text} data to send to a remote server
  156. */
  157. LUA_FUNCTION_DEF (tcp, add_write);
  158. /***
  159. * @method tcp:shift_callback()
  160. *
  161. * Shifts the current callback and go to the next one (if any)
  162. */
  163. LUA_FUNCTION_DEF (tcp, shift_callback);
  164. /***
  165. * @method tcp:starttls([no_verify])
  166. *
  167. * Starts tls connection
  168. * @param {boolean} no_verify used to skip ssl verification
  169. */
  170. LUA_FUNCTION_DEF (tcp, starttls);
  171. static const struct luaL_reg tcp_libf[] = {
  172. LUA_INTERFACE_DEF (tcp, request),
  173. {"new", lua_tcp_request},
  174. {"connect", lua_tcp_request},
  175. {"connect_sync", lua_tcp_connect_sync},
  176. {NULL, NULL}
  177. };
  178. static const struct luaL_reg tcp_libm[] = {
  179. LUA_INTERFACE_DEF (tcp, close),
  180. LUA_INTERFACE_DEF (tcp, add_read),
  181. LUA_INTERFACE_DEF (tcp, add_write),
  182. LUA_INTERFACE_DEF (tcp, shift_callback),
  183. LUA_INTERFACE_DEF (tcp, starttls),
  184. {"__tostring", rspamd_lua_class_tostring},
  185. {NULL, NULL}
  186. };
  187. /***
  188. * @method tcp:close()
  189. *
  190. * Closes TCP connection
  191. */
  192. LUA_FUNCTION_DEF (tcp_sync, close);
  193. /***
  194. * @method read_once()
  195. *
  196. * Performs one read operation. If syscall returned with EAGAIN/EINT,
  197. * restarts the operation, so it always returns either data or error.
  198. */
  199. LUA_FUNCTION_DEF (tcp_sync, read_once);
  200. /***
  201. * @method eof()
  202. *
  203. * True if last IO operation ended with EOF, i.e. endpoint closed connection
  204. */
  205. LUA_FUNCTION_DEF (tcp_sync, eof);
  206. /***
  207. * @method shutdown()
  208. *
  209. * Half-shutdown TCP connection
  210. */
  211. LUA_FUNCTION_DEF (tcp_sync, shutdown);
  212. /***
  213. * @method write()
  214. *
  215. * Writes data into the stream. If syscall returned with EAGAIN/EINT
  216. * restarts the operation. If performs write() until all the passed
  217. * data is written completely.
  218. */
  219. LUA_FUNCTION_DEF (tcp_sync, write);
  220. LUA_FUNCTION_DEF (tcp_sync, gc);
  221. static void lua_tcp_sync_session_dtor (gpointer ud);
  222. static const struct luaL_reg tcp_sync_libm[] = {
  223. LUA_INTERFACE_DEF (tcp_sync, close),
  224. LUA_INTERFACE_DEF (tcp_sync, read_once),
  225. LUA_INTERFACE_DEF (tcp_sync, write),
  226. LUA_INTERFACE_DEF (tcp_sync, eof),
  227. LUA_INTERFACE_DEF (tcp_sync, shutdown),
  228. {"__gc", lua_tcp_sync_gc},
  229. {"__tostring", rspamd_lua_class_tostring},
  230. {NULL, NULL}
  231. };
  232. struct lua_tcp_read_handler {
  233. gchar *stop_pattern;
  234. guint plen;
  235. gint cbref;
  236. };
  237. struct lua_tcp_write_handler {
  238. struct iovec *iov;
  239. guint iovlen;
  240. gint cbref;
  241. gsize pos;
  242. gsize total_bytes;
  243. };
  244. enum lua_tcp_handler_type {
  245. LUA_WANT_WRITE = 0,
  246. LUA_WANT_READ,
  247. LUA_WANT_CONNECT
  248. };
  249. struct lua_tcp_handler {
  250. union {
  251. struct lua_tcp_read_handler r;
  252. struct lua_tcp_write_handler w;
  253. } h;
  254. enum lua_tcp_handler_type type;
  255. };
  256. struct lua_tcp_dtor {
  257. rspamd_mempool_destruct_t dtor;
  258. void *data;
  259. struct lua_tcp_dtor *next;
  260. };
  261. #define LUA_TCP_FLAG_PARTIAL (1u << 0u)
  262. #define LUA_TCP_FLAG_SHUTDOWN (1u << 2u)
  263. #define LUA_TCP_FLAG_CONNECTED (1u << 3u)
  264. #define LUA_TCP_FLAG_FINISHED (1u << 4u)
  265. #define LUA_TCP_FLAG_SYNC (1u << 5u)
  266. #define LUA_TCP_FLAG_RESOLVED (1u << 6u)
  267. #define LUA_TCP_FLAG_SSL (1u << 7u)
  268. #define LUA_TCP_FLAG_SSL_NOVERIFY (1u << 8u)
  269. #undef TCP_DEBUG_REFS
  270. #ifdef TCP_DEBUG_REFS
  271. #define TCP_RETAIN(x) do { \
  272. msg_info ("retain ref %p, refcount: %d", (x), (x)->ref.refcount); \
  273. REF_RETAIN(x); \
  274. } while (0)
  275. #define TCP_RELEASE(x) do { \
  276. msg_info ("release ref %p, refcount: %d", (x), (x)->ref.refcount); \
  277. REF_RELEASE(x); \
  278. } while (0)
  279. #else
  280. #define TCP_RETAIN(x) REF_RETAIN(x)
  281. #define TCP_RELEASE(x) REF_RELEASE(x)
  282. #endif
  283. struct lua_tcp_cbdata {
  284. struct rspamd_async_session *session;
  285. struct rspamd_async_event *async_ev;
  286. struct ev_loop *event_loop;
  287. rspamd_inet_addr_t *addr;
  288. GByteArray *in;
  289. GQueue *handlers;
  290. gint fd;
  291. gint connect_cb;
  292. guint port;
  293. guint flags;
  294. gchar tag[7];
  295. struct rspamd_io_ev ev;
  296. struct lua_tcp_dtor *dtors;
  297. ref_entry_t ref;
  298. struct rspamd_task *task;
  299. struct rspamd_symcache_item *item;
  300. struct thread_entry *thread;
  301. struct rspamd_config *cfg;
  302. struct rspamd_ssl_connection *ssl_conn;
  303. gchar *hostname;
  304. gboolean eof;
  305. };
  306. #define IS_SYNC(c) (((c)->flags & LUA_TCP_FLAG_SYNC) != 0)
  307. #define msg_debug_tcp(...) rspamd_conditional_debug_fast (NULL, cbd->addr, \
  308. rspamd_lua_tcp_log_id, "lua_tcp", cbd->tag, \
  309. G_STRFUNC, \
  310. __VA_ARGS__)
  311. INIT_LOG_MODULE(lua_tcp)
  312. static void lua_tcp_handler (int fd, short what, gpointer ud);
  313. static void lua_tcp_plan_handler_event (struct lua_tcp_cbdata *cbd,
  314. gboolean can_read, gboolean can_write);
  315. static void lua_tcp_unregister_event (struct lua_tcp_cbdata *cbd);
  316. static void
  317. lua_tcp_void_finalyser (gpointer arg) {}
  318. static const gdouble default_tcp_timeout = 5.0;
  319. static struct rspamd_dns_resolver *
  320. lua_tcp_global_resolver (struct ev_loop *ev_base,
  321. struct rspamd_config *cfg)
  322. {
  323. static struct rspamd_dns_resolver *global_resolver;
  324. if (cfg && cfg->dns_resolver) {
  325. return cfg->dns_resolver;
  326. }
  327. if (global_resolver == NULL) {
  328. global_resolver = rspamd_dns_resolver_init (NULL, ev_base, cfg);
  329. }
  330. return global_resolver;
  331. }
  332. static gboolean
  333. lua_tcp_shift_handler (struct lua_tcp_cbdata *cbd)
  334. {
  335. struct lua_tcp_handler *hdl;
  336. hdl = g_queue_pop_head (cbd->handlers);
  337. if (hdl == NULL) {
  338. /* We are done */
  339. return FALSE;
  340. }
  341. if (hdl->type == LUA_WANT_READ) {
  342. msg_debug_tcp ("switch from read handler %d", hdl->h.r.cbref);
  343. if (hdl->h.r.cbref && hdl->h.r.cbref != -1) {
  344. luaL_unref (cbd->cfg->lua_state, LUA_REGISTRYINDEX, hdl->h.r.cbref);
  345. }
  346. if (hdl->h.r.stop_pattern) {
  347. g_free (hdl->h.r.stop_pattern);
  348. }
  349. }
  350. else if (hdl->type == LUA_WANT_WRITE) {
  351. msg_debug_tcp ("switch from write handler %d", hdl->h.r.cbref);
  352. if (hdl->h.w.cbref && hdl->h.w.cbref != -1) {
  353. luaL_unref (cbd->cfg->lua_state, LUA_REGISTRYINDEX, hdl->h.w.cbref);
  354. }
  355. if (hdl->h.w.iov) {
  356. g_free (hdl->h.w.iov);
  357. }
  358. }
  359. else {
  360. msg_debug_tcp ("removing connect handler");
  361. /* LUA_WANT_CONNECT: it doesn't allocate anything, nothing to do here */
  362. }
  363. g_free (hdl);
  364. return TRUE;
  365. }
  366. static void
  367. lua_tcp_fin (gpointer arg)
  368. {
  369. struct lua_tcp_cbdata *cbd = (struct lua_tcp_cbdata *)arg;
  370. struct lua_tcp_dtor *dtor, *dttmp;
  371. if (IS_SYNC (cbd) && cbd->task) {
  372. /*
  373. pointer is now becoming invalid, we should remove registered destructor,
  374. all the necessary steps are done here
  375. */
  376. rspamd_mempool_replace_destructor (cbd->task->task_pool,
  377. lua_tcp_sync_session_dtor, cbd, NULL);
  378. }
  379. msg_debug_tcp ("finishing TCP %s connection", IS_SYNC (cbd) ? "sync" : "async");
  380. if (cbd->connect_cb != -1) {
  381. luaL_unref (cbd->cfg->lua_state, LUA_REGISTRYINDEX, cbd->connect_cb);
  382. }
  383. if (cbd->ssl_conn) {
  384. /* TODO: postpone close in case ssl is used ! */
  385. rspamd_ssl_connection_free (cbd->ssl_conn);
  386. }
  387. if (cbd->fd != -1) {
  388. rspamd_ev_watcher_stop (cbd->event_loop, &cbd->ev);
  389. close (cbd->fd);
  390. cbd->fd = -1;
  391. }
  392. if (cbd->addr) {
  393. rspamd_inet_address_free (cbd->addr);
  394. }
  395. while (lua_tcp_shift_handler (cbd)) {}
  396. g_queue_free (cbd->handlers);
  397. LL_FOREACH_SAFE (cbd->dtors, dtor, dttmp) {
  398. dtor->dtor (dtor->data);
  399. g_free (dtor);
  400. }
  401. g_byte_array_unref (cbd->in);
  402. g_free (cbd->hostname);
  403. g_free (cbd);
  404. }
  405. static struct lua_tcp_cbdata *
  406. lua_check_tcp (lua_State *L, gint pos)
  407. {
  408. void *ud = rspamd_lua_check_udata (L, pos, "rspamd{tcp}");
  409. luaL_argcheck (L, ud != NULL, pos, "'tcp' expected");
  410. return ud ? *((struct lua_tcp_cbdata **)ud) : NULL;
  411. }
  412. static void
  413. lua_tcp_maybe_free (struct lua_tcp_cbdata *cbd)
  414. {
  415. if (IS_SYNC (cbd)) {
  416. /*
  417. * in this mode, we don't remove object, we only remove the event
  418. * Object is owned by lua and will be destroyed on __gc()
  419. */
  420. if (cbd->item) {
  421. rspamd_symcache_item_async_dec_check (cbd->task, cbd->item, M);
  422. cbd->item = NULL;
  423. }
  424. if (cbd->async_ev) {
  425. rspamd_session_remove_event (cbd->session, lua_tcp_void_finalyser, cbd);
  426. }
  427. cbd->async_ev = NULL;
  428. }
  429. else {
  430. if (cbd->item) {
  431. rspamd_symcache_item_async_dec_check (cbd->task, cbd->item, M);
  432. cbd->item = NULL;
  433. }
  434. if (cbd->async_ev) {
  435. rspamd_session_remove_event (cbd->session, lua_tcp_fin, cbd);
  436. }
  437. else {
  438. lua_tcp_fin (cbd);
  439. }
  440. }
  441. }
  442. #ifdef __GNUC__
  443. static void
  444. lua_tcp_push_error (struct lua_tcp_cbdata *cbd, gboolean is_fatal,
  445. const char *err, ...) __attribute__ ((format(printf, 3, 4)));
  446. #endif
  447. static void lua_tcp_resume_thread_error_argp (struct lua_tcp_cbdata *cbd, const gchar *error, va_list argp);
  448. static void
  449. lua_tcp_push_error (struct lua_tcp_cbdata *cbd, gboolean is_fatal,
  450. const char *err, ...)
  451. {
  452. va_list ap, ap_copy;
  453. struct lua_tcp_cbdata **pcbd;
  454. struct lua_tcp_handler *hdl;
  455. gint cbref, top;
  456. struct lua_callback_state cbs;
  457. lua_State *L;
  458. gboolean callback_called = FALSE;
  459. if (cbd->thread) {
  460. va_start (ap, err);
  461. lua_tcp_resume_thread_error_argp (cbd, err, ap);
  462. va_end (ap);
  463. return;
  464. }
  465. lua_thread_pool_prepare_callback (cbd->cfg->lua_thread_pool, &cbs);
  466. L = cbs.L;
  467. va_start (ap, err);
  468. for (;;) {
  469. hdl = g_queue_peek_head (cbd->handlers);
  470. if (hdl == NULL) {
  471. break;
  472. }
  473. if (hdl->type == LUA_WANT_READ) {
  474. cbref = hdl->h.r.cbref;
  475. }
  476. else {
  477. cbref = hdl->h.w.cbref;
  478. }
  479. if (cbref != -1) {
  480. top = lua_gettop (L);
  481. lua_rawgeti (L, LUA_REGISTRYINDEX, cbref);
  482. /* Error message */
  483. va_copy (ap_copy, ap);
  484. lua_pushvfstring (L, err, ap_copy);
  485. va_end (ap_copy);
  486. /* Body */
  487. lua_pushnil (L);
  488. /* Connection */
  489. pcbd = lua_newuserdata (L, sizeof (*pcbd));
  490. *pcbd = cbd;
  491. rspamd_lua_setclass (L, "rspamd{tcp}", -1);
  492. TCP_RETAIN (cbd);
  493. if (cbd->item) {
  494. rspamd_symcache_set_cur_item (cbd->task, cbd->item);
  495. }
  496. if (lua_pcall (L, 3, 0, 0) != 0) {
  497. msg_info ("callback call failed: %s", lua_tostring (L, -1));
  498. }
  499. lua_settop (L, top);
  500. TCP_RELEASE (cbd);
  501. callback_called = TRUE;
  502. }
  503. if (!is_fatal) {
  504. if (callback_called) {
  505. /* Stop on the first callback found */
  506. break;
  507. }
  508. else {
  509. /* Shift to another callback to inform about non fatal error */
  510. msg_debug_tcp ("non fatal error find matching callback");
  511. lua_tcp_shift_handler (cbd);
  512. continue;
  513. }
  514. }
  515. else {
  516. msg_debug_tcp ("fatal error rollback all handlers");
  517. lua_tcp_shift_handler (cbd);
  518. }
  519. }
  520. va_end (ap);
  521. lua_thread_pool_restore_callback (&cbs);
  522. }
  523. static void lua_tcp_resume_thread (struct lua_tcp_cbdata *cbd, const guint8 *str, gsize len);
  524. static void
  525. lua_tcp_push_data (struct lua_tcp_cbdata *cbd, const guint8 *str, gsize len)
  526. {
  527. struct rspamd_lua_text *t;
  528. struct lua_tcp_cbdata **pcbd;
  529. struct lua_tcp_handler *hdl;
  530. gint cbref, arg_cnt, top;
  531. struct lua_callback_state cbs;
  532. lua_State *L;
  533. if (cbd->thread) {
  534. lua_tcp_resume_thread (cbd, str, len);
  535. return;
  536. }
  537. lua_thread_pool_prepare_callback (cbd->cfg->lua_thread_pool, &cbs);
  538. L = cbs.L;
  539. hdl = g_queue_peek_head (cbd->handlers);
  540. g_assert (hdl != NULL);
  541. if (hdl->type == LUA_WANT_READ) {
  542. cbref = hdl->h.r.cbref;
  543. }
  544. else {
  545. cbref = hdl->h.w.cbref;
  546. }
  547. if (cbref != -1) {
  548. top = lua_gettop (L);
  549. lua_rawgeti (L, LUA_REGISTRYINDEX, cbref);
  550. /* Error */
  551. lua_pushnil (L);
  552. /* Body */
  553. if (hdl->type == LUA_WANT_READ) {
  554. t = lua_newuserdata (L, sizeof (*t));
  555. rspamd_lua_setclass (L, "rspamd{text}", -1);
  556. t->start = (const gchar *)str;
  557. t->len = len;
  558. t->flags = 0;
  559. arg_cnt = 3;
  560. }
  561. else {
  562. arg_cnt = 2;
  563. }
  564. /* Connection */
  565. pcbd = lua_newuserdata (L, sizeof (*pcbd));
  566. *pcbd = cbd;
  567. rspamd_lua_setclass (L, "rspamd{tcp}", -1);
  568. TCP_RETAIN (cbd);
  569. if (cbd->item) {
  570. rspamd_symcache_set_cur_item (cbd->task, cbd->item);
  571. }
  572. if (lua_pcall (L, arg_cnt, 0, 0) != 0) {
  573. msg_info ("callback call failed: %s", lua_tostring (L, -1));
  574. }
  575. lua_settop (L, top);
  576. TCP_RELEASE (cbd);
  577. }
  578. lua_thread_pool_restore_callback (&cbs);
  579. }
  580. static void
  581. lua_tcp_resume_thread_error_argp (struct lua_tcp_cbdata *cbd, const gchar *error, va_list argp)
  582. {
  583. struct thread_entry *thread = cbd->thread;
  584. lua_State *L = thread->lua_state;
  585. lua_pushboolean (L, FALSE);
  586. lua_pushvfstring (L, error, argp);
  587. lua_tcp_shift_handler (cbd);
  588. // lua_tcp_unregister_event (cbd);
  589. lua_thread_pool_set_running_entry (cbd->cfg->lua_thread_pool, cbd->thread);
  590. lua_thread_resume (thread, 2);
  591. TCP_RELEASE (cbd);
  592. }
  593. static void
  594. lua_tcp_resume_thread (struct lua_tcp_cbdata *cbd, const guint8 *str, gsize len)
  595. {
  596. /*
  597. * typical call returns:
  598. *
  599. * read:
  600. * error:
  601. * (nil, error message)
  602. * got data:
  603. * (true, data)
  604. * write/connect:
  605. * error:
  606. * (nil, error message)
  607. * wrote
  608. * (true)
  609. */
  610. lua_State *L = cbd->thread->lua_state;
  611. struct lua_tcp_handler *hdl;
  612. hdl = g_queue_peek_head (cbd->handlers);
  613. lua_pushboolean (L, TRUE);
  614. if (hdl->type == LUA_WANT_READ) {
  615. lua_pushlstring (L, str, len);
  616. }
  617. else {
  618. lua_pushnil (L);
  619. }
  620. lua_tcp_shift_handler (cbd);
  621. lua_thread_pool_set_running_entry (cbd->cfg->lua_thread_pool,
  622. cbd->thread);
  623. if (cbd->item) {
  624. rspamd_symcache_set_cur_item (cbd->task, cbd->item);
  625. }
  626. lua_thread_resume (cbd->thread, 2);
  627. TCP_RELEASE (cbd);
  628. }
  629. static void
  630. lua_tcp_plan_read (struct lua_tcp_cbdata *cbd)
  631. {
  632. rspamd_ev_watcher_reschedule (cbd->event_loop, &cbd->ev, EV_READ);
  633. }
  634. static void
  635. lua_tcp_connect_helper (struct lua_tcp_cbdata *cbd)
  636. {
  637. /* This is used for sync mode only */
  638. lua_State *L = cbd->thread->lua_state;
  639. struct lua_tcp_cbdata **pcbd;
  640. lua_pushboolean (L, TRUE);
  641. lua_thread_pool_set_running_entry (cbd->cfg->lua_thread_pool, cbd->thread);
  642. pcbd = lua_newuserdata (L, sizeof (*pcbd));
  643. *pcbd = cbd;
  644. rspamd_lua_setclass (L, "rspamd{tcp_sync}", -1);
  645. msg_debug_tcp ("tcp connected");
  646. lua_tcp_shift_handler (cbd);
  647. // lua_tcp_unregister_event (cbd);
  648. lua_thread_resume (cbd->thread, 2);
  649. TCP_RELEASE (cbd);
  650. }
  651. static void
  652. lua_tcp_write_helper (struct lua_tcp_cbdata *cbd)
  653. {
  654. struct iovec *start;
  655. guint niov, i;
  656. gint flags = 0;
  657. bool allocated_iov = false;
  658. gsize remain;
  659. gssize r;
  660. struct iovec *cur_iov;
  661. struct lua_tcp_handler *hdl;
  662. struct lua_tcp_write_handler *wh;
  663. struct msghdr msg;
  664. hdl = g_queue_peek_head (cbd->handlers);
  665. g_assert (hdl != NULL && hdl->type == LUA_WANT_WRITE);
  666. wh = &hdl->h.w;
  667. if (wh->pos == wh->total_bytes) {
  668. goto call_finish_handler;
  669. }
  670. start = &wh->iov[0];
  671. niov = wh->iovlen;
  672. remain = wh->pos;
  673. /* We know that niov is small enough for that */
  674. if (niov < 1024) {
  675. cur_iov = g_alloca (niov * sizeof (struct iovec));
  676. }
  677. else {
  678. cur_iov = g_malloc0 (niov * sizeof (struct iovec));
  679. allocated_iov = true;
  680. }
  681. memcpy (cur_iov, wh->iov, niov * sizeof (struct iovec));
  682. for (i = 0; i < wh->iovlen && remain > 0; i++) {
  683. /* Find out the first iov required */
  684. start = &cur_iov[i];
  685. if (start->iov_len <= remain) {
  686. remain -= start->iov_len;
  687. start = &cur_iov[i + 1];
  688. niov--;
  689. }
  690. else {
  691. start->iov_base = (void *)((char *)start->iov_base + remain);
  692. start->iov_len -= remain;
  693. remain = 0;
  694. }
  695. }
  696. memset (&msg, 0, sizeof (msg));
  697. msg.msg_iov = start;
  698. msg.msg_iovlen = MIN (IOV_MAX, niov);
  699. g_assert (niov > 0);
  700. #ifdef MSG_NOSIGNAL
  701. flags = MSG_NOSIGNAL;
  702. #endif
  703. msg_debug_tcp ("want write %d io vectors of %d", (int)msg.msg_iovlen,
  704. (int)niov);
  705. if (cbd->ssl_conn) {
  706. r = rspamd_ssl_writev (cbd->ssl_conn, msg.msg_iov, msg.msg_iovlen);
  707. }
  708. else {
  709. r = sendmsg (cbd->fd, &msg, flags);
  710. }
  711. if (allocated_iov) {
  712. g_free (cur_iov);
  713. }
  714. if (r == -1) {
  715. if (!(cbd->ssl_conn)) {
  716. if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
  717. msg_debug_tcp ("got temporary failure, retry write");
  718. lua_tcp_plan_handler_event (cbd, TRUE, TRUE);
  719. return;
  720. }
  721. else {
  722. lua_tcp_push_error (cbd, TRUE,
  723. "IO write error while trying to write %d bytes: %s",
  724. (gint) remain, strerror (errno));
  725. msg_debug_tcp ("write error, terminate connection");
  726. TCP_RELEASE (cbd);
  727. }
  728. }
  729. return;
  730. }
  731. else {
  732. wh->pos += r;
  733. }
  734. msg_debug_tcp ("written %z bytes: %z/%z", r,
  735. wh->pos, wh->total_bytes);
  736. if (wh->pos >= wh->total_bytes) {
  737. goto call_finish_handler;
  738. }
  739. else {
  740. /* Want to write more */
  741. if (r > 0) {
  742. /* XXX: special case: we know that we want to write more data
  743. * than it is available in iov function.
  744. *
  745. * Hence, we need to check if we can write more at some point...
  746. */
  747. lua_tcp_write_helper (cbd);
  748. }
  749. }
  750. return;
  751. call_finish_handler:
  752. msg_debug_tcp ("finishing TCP write, calling TCP handler");
  753. if ((cbd->flags & LUA_TCP_FLAG_SHUTDOWN)) {
  754. /* Half close the connection */
  755. shutdown (cbd->fd, SHUT_WR);
  756. cbd->flags &= ~LUA_TCP_FLAG_SHUTDOWN;
  757. }
  758. lua_tcp_push_data (cbd, NULL, 0);
  759. if (!IS_SYNC (cbd)) {
  760. lua_tcp_shift_handler (cbd);
  761. lua_tcp_plan_handler_event (cbd, TRUE, TRUE);
  762. }
  763. }
  764. static gboolean
  765. lua_tcp_process_read_handler (struct lua_tcp_cbdata *cbd,
  766. struct lua_tcp_read_handler *rh, gboolean eof)
  767. {
  768. guint slen;
  769. goffset pos;
  770. if (rh->stop_pattern) {
  771. slen = rh->plen;
  772. if (cbd->in->len >= slen) {
  773. if ((pos = rspamd_substring_search (cbd->in->data, cbd->in->len,
  774. rh->stop_pattern, slen)) != -1) {
  775. msg_debug_tcp ("found TCP stop pattern");
  776. lua_tcp_push_data (cbd, cbd->in->data, pos);
  777. if (!IS_SYNC (cbd)) {
  778. lua_tcp_shift_handler (cbd);
  779. }
  780. if (pos + slen < cbd->in->len) {
  781. /* We have a leftover */
  782. memmove (cbd->in->data, cbd->in->data + pos + slen,
  783. cbd->in->len - (pos + slen));
  784. cbd->in->len = cbd->in->len - (pos + slen);
  785. }
  786. else {
  787. cbd->in->len = 0;
  788. }
  789. return TRUE;
  790. }
  791. else {
  792. /* Plan new read */
  793. msg_debug_tcp ("NOT found TCP stop pattern");
  794. if (!cbd->eof) {
  795. lua_tcp_plan_read (cbd);
  796. }
  797. else {
  798. /* Got session finished but no stop pattern */
  799. lua_tcp_push_error (cbd, TRUE,
  800. "IO read error: connection terminated");
  801. }
  802. }
  803. }
  804. }
  805. else {
  806. msg_debug_tcp ("read TCP partial data %d bytes", cbd->in->len);
  807. slen = cbd->in->len;
  808. /* we have eaten all the data, handler should not know that there is something */
  809. cbd->in->len = 0;
  810. lua_tcp_push_data (cbd, cbd->in->data, slen);
  811. if (!IS_SYNC (cbd)) {
  812. lua_tcp_shift_handler (cbd);
  813. }
  814. return TRUE;
  815. }
  816. return FALSE;
  817. }
  818. static void
  819. lua_tcp_process_read (struct lua_tcp_cbdata *cbd,
  820. guchar *in, gssize r)
  821. {
  822. struct lua_tcp_handler *hdl;
  823. struct lua_tcp_read_handler *rh;
  824. hdl = g_queue_peek_head (cbd->handlers);
  825. g_assert (hdl != NULL && hdl->type == LUA_WANT_READ);
  826. rh = &hdl->h.r;
  827. if (r > 0) {
  828. if (cbd->flags & LUA_TCP_FLAG_PARTIAL) {
  829. lua_tcp_push_data (cbd, in, r);
  830. /* Plan next event */
  831. lua_tcp_plan_read (cbd);
  832. }
  833. else {
  834. g_byte_array_append (cbd->in, in, r);
  835. if (!lua_tcp_process_read_handler (cbd, rh, FALSE)) {
  836. /* Plan more read */
  837. lua_tcp_plan_read (cbd);
  838. }
  839. else {
  840. /* Go towards the next handler */
  841. if (!IS_SYNC (cbd)) {
  842. lua_tcp_plan_handler_event (cbd, TRUE, TRUE);
  843. }
  844. }
  845. }
  846. }
  847. else if (r == 0) {
  848. /* EOF */
  849. cbd->eof = TRUE;
  850. if (cbd->in->len > 0) {
  851. /* We have some data to process */
  852. lua_tcp_process_read_handler (cbd, rh, TRUE);
  853. }
  854. else {
  855. lua_tcp_push_error (cbd, TRUE, "IO read error: connection terminated");
  856. }
  857. lua_tcp_plan_handler_event (cbd, FALSE, FALSE);
  858. }
  859. else {
  860. /* An error occurred */
  861. if (errno == EAGAIN || errno == EINTR) {
  862. /* Restart call */
  863. lua_tcp_plan_read (cbd);
  864. return;
  865. }
  866. /* Fatal error */
  867. cbd->eof = TRUE;
  868. if (cbd->in->len > 0) {
  869. /* We have some data to process */
  870. lua_tcp_process_read_handler (cbd, rh, TRUE);
  871. }
  872. else {
  873. lua_tcp_push_error (cbd, TRUE,
  874. "IO read error while trying to read data: %s",
  875. strerror (errno));
  876. }
  877. lua_tcp_plan_handler_event (cbd, FALSE, FALSE);
  878. }
  879. }
  880. static void
  881. lua_tcp_handler (int fd, short what, gpointer ud)
  882. {
  883. struct lua_tcp_cbdata *cbd = ud;
  884. guchar inbuf[8192];
  885. gssize r;
  886. gint so_error = 0;
  887. socklen_t so_len = sizeof (so_error);
  888. struct lua_callback_state cbs;
  889. lua_State *L;
  890. enum lua_tcp_handler_type event_type;
  891. TCP_RETAIN (cbd);
  892. msg_debug_tcp ("processed TCP event: %d", what);
  893. struct lua_tcp_handler *rh = g_queue_peek_head (cbd->handlers);
  894. event_type = rh->type;
  895. rspamd_ev_watcher_stop (cbd->event_loop, &cbd->ev);
  896. if (what == EV_READ) {
  897. if (cbd->ssl_conn) {
  898. r = rspamd_ssl_read (cbd->ssl_conn, inbuf, sizeof (inbuf));
  899. }
  900. else {
  901. r = read (cbd->fd, inbuf, sizeof (inbuf));
  902. }
  903. lua_tcp_process_read (cbd, inbuf, r);
  904. }
  905. else if (what == EV_WRITE) {
  906. if (!(cbd->flags & LUA_TCP_FLAG_CONNECTED)) {
  907. if (getsockopt (fd, SOL_SOCKET, SO_ERROR, &so_error, &so_len) == -1) {
  908. lua_tcp_push_error (cbd, TRUE, "Cannot get socket error: %s",
  909. strerror (errno));
  910. TCP_RELEASE (cbd);
  911. goto out;
  912. }
  913. else if (so_error != 0) {
  914. lua_tcp_push_error (cbd, TRUE, "Socket error detected: %s",
  915. strerror (so_error));
  916. TCP_RELEASE (cbd);
  917. goto out;
  918. }
  919. else {
  920. cbd->flags |= LUA_TCP_FLAG_CONNECTED;
  921. if (cbd->connect_cb != -1) {
  922. struct lua_tcp_cbdata **pcbd;
  923. gint top;
  924. lua_thread_pool_prepare_callback (cbd->cfg->lua_thread_pool, &cbs);
  925. L = cbs.L;
  926. top = lua_gettop (L);
  927. lua_rawgeti (L, LUA_REGISTRYINDEX, cbd->connect_cb);
  928. pcbd = lua_newuserdata (L, sizeof (*pcbd));
  929. *pcbd = cbd;
  930. TCP_RETAIN (cbd);
  931. rspamd_lua_setclass (L, "rspamd{tcp}", -1);
  932. if (cbd->item) {
  933. rspamd_symcache_set_cur_item (cbd->task, cbd->item);
  934. }
  935. if (lua_pcall (L, 1, 0, 0) != 0) {
  936. msg_info ("callback call failed: %s", lua_tostring (L, -1));
  937. }
  938. lua_settop (L, top);
  939. TCP_RELEASE (cbd);
  940. lua_thread_pool_restore_callback (&cbs);
  941. }
  942. }
  943. }
  944. if (event_type == LUA_WANT_WRITE) {
  945. lua_tcp_write_helper (cbd);
  946. }
  947. else if (event_type == LUA_WANT_CONNECT) {
  948. lua_tcp_connect_helper (cbd);
  949. }
  950. else {
  951. g_assert_not_reached ();
  952. }
  953. }
  954. else {
  955. lua_tcp_push_error (cbd, TRUE, "IO timeout");
  956. TCP_RELEASE (cbd);
  957. }
  958. out:
  959. TCP_RELEASE (cbd);
  960. }
  961. static void
  962. lua_tcp_plan_handler_event (struct lua_tcp_cbdata *cbd, gboolean can_read,
  963. gboolean can_write)
  964. {
  965. struct lua_tcp_handler *hdl;
  966. hdl = g_queue_peek_head (cbd->handlers);
  967. if (hdl == NULL) {
  968. if (!(cbd->flags & LUA_TCP_FLAG_FINISHED)) {
  969. /* We are finished with a connection */
  970. msg_debug_tcp ("no handlers left, finish session");
  971. TCP_RELEASE (cbd);
  972. cbd->flags |= LUA_TCP_FLAG_FINISHED;
  973. }
  974. }
  975. else {
  976. if (hdl->type == LUA_WANT_READ) {
  977. /* We need to check if we have some leftover in the buffer */
  978. if (cbd->in->len > 0) {
  979. msg_debug_tcp ("process read buffer leftover");
  980. if (lua_tcp_process_read_handler (cbd, &hdl->h.r, FALSE)) {
  981. if (!IS_SYNC(cbd)) {
  982. /* We can go to the next handler */
  983. lua_tcp_plan_handler_event (cbd, can_read, can_write);
  984. }
  985. }
  986. }
  987. else {
  988. if (can_read) {
  989. /* We need to plan a new event */
  990. msg_debug_tcp ("plan new read");
  991. rspamd_ev_watcher_reschedule (cbd->event_loop, &cbd->ev,
  992. EV_READ);
  993. }
  994. else {
  995. /* Cannot read more */
  996. msg_debug_tcp ("cannot read more");
  997. lua_tcp_push_error (cbd, FALSE, "EOF, cannot read more data");
  998. if (!IS_SYNC (cbd)) {
  999. lua_tcp_shift_handler (cbd);
  1000. lua_tcp_plan_handler_event (cbd, can_read, can_write);
  1001. }
  1002. }
  1003. }
  1004. }
  1005. else if (hdl->type == LUA_WANT_WRITE) {
  1006. /*
  1007. * We need to plan write event if there is something in the
  1008. * write request
  1009. */
  1010. if (hdl->h.w.pos < hdl->h.w.total_bytes) {
  1011. msg_debug_tcp ("plan new write");
  1012. if (can_write) {
  1013. rspamd_ev_watcher_reschedule (cbd->event_loop, &cbd->ev,
  1014. EV_WRITE);
  1015. }
  1016. else {
  1017. /* Cannot write more */
  1018. lua_tcp_push_error (cbd, FALSE, "EOF, cannot write more data");
  1019. if (!IS_SYNC(cbd)) {
  1020. lua_tcp_shift_handler (cbd);
  1021. lua_tcp_plan_handler_event (cbd, can_read, can_write);
  1022. }
  1023. }
  1024. }
  1025. else {
  1026. /* We shouldn't have empty write handlers */
  1027. g_assert_not_reached ();
  1028. }
  1029. }
  1030. else { /* LUA_WANT_CONNECT */
  1031. msg_debug_tcp ("plan new connect");
  1032. rspamd_ev_watcher_reschedule (cbd->event_loop, &cbd->ev,
  1033. EV_WRITE);
  1034. }
  1035. }
  1036. }
  1037. static gboolean
  1038. lua_tcp_register_event (struct lua_tcp_cbdata *cbd)
  1039. {
  1040. if (cbd->session) {
  1041. event_finalizer_t fin = IS_SYNC (cbd) ? lua_tcp_void_finalyser : lua_tcp_fin;
  1042. if (cbd->item) {
  1043. cbd->async_ev = rspamd_session_add_event_full (cbd->session, fin, cbd, M,
  1044. rspamd_symcache_item_name (cbd->item));
  1045. }
  1046. else {
  1047. cbd->async_ev = rspamd_session_add_event (cbd->session, fin, cbd, M);
  1048. }
  1049. if (!cbd->async_ev) {
  1050. return FALSE;
  1051. }
  1052. }
  1053. return TRUE;
  1054. }
  1055. static void
  1056. lua_tcp_register_watcher (struct lua_tcp_cbdata *cbd)
  1057. {
  1058. if (cbd->item && cbd->task) {
  1059. rspamd_symcache_item_async_inc (cbd->task, cbd->item, M);
  1060. }
  1061. }
  1062. static void
  1063. lua_tcp_ssl_on_error (gpointer ud, GError *err)
  1064. {
  1065. struct lua_tcp_cbdata *cbd = (struct lua_tcp_cbdata *)ud;
  1066. if (err) {
  1067. lua_tcp_push_error (cbd, TRUE, "ssl error: %s", err->message);
  1068. }
  1069. else {
  1070. lua_tcp_push_error (cbd, TRUE, "ssl error: unknown error");
  1071. }
  1072. TCP_RELEASE (cbd);
  1073. }
  1074. static gboolean
  1075. lua_tcp_make_connection (struct lua_tcp_cbdata *cbd)
  1076. {
  1077. int fd;
  1078. rspamd_inet_address_set_port (cbd->addr, cbd->port);
  1079. fd = rspamd_inet_address_connect (cbd->addr, SOCK_STREAM, TRUE);
  1080. if (fd == -1) {
  1081. if (cbd->session) {
  1082. rspamd_mempool_t *pool = rspamd_session_mempool (cbd->session);
  1083. msg_info_pool ("cannot connect to %s (%s): %s",
  1084. rspamd_inet_address_to_string (cbd->addr),
  1085. cbd->hostname,
  1086. strerror (errno));
  1087. }
  1088. else {
  1089. msg_info ("cannot connect to %s (%s): %s",
  1090. rspamd_inet_address_to_string (cbd->addr),
  1091. cbd->hostname,
  1092. strerror (errno));
  1093. }
  1094. return FALSE;
  1095. }
  1096. cbd->fd = fd;
  1097. #if 0
  1098. if (!(cbd->flags & LUA_TCP_FLAG_RESOLVED)) {
  1099. /* We come here without resolving, so we need to add a watcher */
  1100. lua_tcp_register_watcher (cbd);
  1101. }
  1102. else {
  1103. cbd->flags |= LUA_TCP_FLAG_RESOLVED;
  1104. }
  1105. #endif
  1106. if (cbd->flags & LUA_TCP_FLAG_SSL) {
  1107. gpointer ssl_ctx;
  1108. gboolean verify_peer;
  1109. if (cbd->flags & LUA_TCP_FLAG_SSL_NOVERIFY) {
  1110. ssl_ctx = cbd->cfg->libs_ctx->ssl_ctx_noverify;
  1111. verify_peer = FALSE;
  1112. }
  1113. else {
  1114. ssl_ctx = cbd->cfg->libs_ctx->ssl_ctx;
  1115. verify_peer = TRUE;
  1116. }
  1117. cbd->ssl_conn = rspamd_ssl_connection_new (ssl_ctx,
  1118. cbd->event_loop,
  1119. verify_peer,
  1120. cbd->tag);
  1121. if (!rspamd_ssl_connect_fd (cbd->ssl_conn, fd, cbd->hostname, &cbd->ev,
  1122. cbd->ev.timeout, lua_tcp_handler, lua_tcp_ssl_on_error, cbd)) {
  1123. lua_tcp_push_error (cbd, TRUE, "ssl connection failed: %s",
  1124. strerror (errno));
  1125. return FALSE;
  1126. }
  1127. else {
  1128. lua_tcp_register_event (cbd);
  1129. }
  1130. }
  1131. else {
  1132. rspamd_ev_watcher_init (&cbd->ev, cbd->fd, EV_WRITE,
  1133. lua_tcp_handler, cbd);
  1134. lua_tcp_register_event (cbd);
  1135. lua_tcp_plan_handler_event (cbd, TRUE, TRUE);
  1136. }
  1137. return TRUE;
  1138. }
  1139. static void
  1140. lua_tcp_dns_handler (struct rdns_reply *reply, gpointer ud)
  1141. {
  1142. struct lua_tcp_cbdata *cbd = (struct lua_tcp_cbdata *)ud;
  1143. const struct rdns_request_name *rn;
  1144. if (reply->code != RDNS_RC_NOERROR) {
  1145. rn = rdns_request_get_name (reply->request, NULL);
  1146. lua_tcp_push_error (cbd, TRUE, "unable to resolve host: %s",
  1147. rn->name);
  1148. TCP_RELEASE (cbd);
  1149. }
  1150. else {
  1151. /*
  1152. * We set this flag as it means that we have already registered the watcher
  1153. * when started DNS query
  1154. */
  1155. cbd->flags |= LUA_TCP_FLAG_RESOLVED;
  1156. if (reply->entries->type == RDNS_REQUEST_A) {
  1157. cbd->addr = rspamd_inet_address_new (AF_INET,
  1158. &reply->entries->content.a.addr);
  1159. }
  1160. else if (reply->entries->type == RDNS_REQUEST_AAAA) {
  1161. cbd->addr = rspamd_inet_address_new (AF_INET6,
  1162. &reply->entries->content.aaa.addr);
  1163. }
  1164. rspamd_inet_address_set_port (cbd->addr, cbd->port);
  1165. if (!lua_tcp_make_connection (cbd)) {
  1166. lua_tcp_push_error (cbd, TRUE, "unable to make connection to the host %s",
  1167. rspamd_inet_address_to_string (cbd->addr));
  1168. TCP_RELEASE (cbd);
  1169. }
  1170. }
  1171. }
  1172. static gboolean
  1173. lua_tcp_arg_toiovec (lua_State *L, gint pos, struct lua_tcp_cbdata *cbd,
  1174. struct iovec *vec)
  1175. {
  1176. struct rspamd_lua_text *t;
  1177. gsize len;
  1178. const gchar *str;
  1179. struct lua_tcp_dtor *dtor;
  1180. if (lua_type (L, pos) == LUA_TUSERDATA) {
  1181. t = lua_check_text (L, pos);
  1182. if (t) {
  1183. vec->iov_base = (void *)t->start;
  1184. vec->iov_len = t->len;
  1185. if (t->flags & RSPAMD_TEXT_FLAG_OWN) {
  1186. /* Steal ownership */
  1187. t->flags = 0;
  1188. dtor = g_malloc0 (sizeof (*dtor));
  1189. dtor->dtor = g_free;
  1190. dtor->data = (void *)t->start;
  1191. LL_PREPEND (cbd->dtors, dtor);
  1192. }
  1193. }
  1194. else {
  1195. msg_err ("bad userdata argument at position %d", pos);
  1196. return FALSE;
  1197. }
  1198. }
  1199. else if (lua_type (L, pos) == LUA_TSTRING) {
  1200. str = luaL_checklstring (L, pos, &len);
  1201. vec->iov_base = g_malloc (len);
  1202. dtor = g_malloc0 (sizeof (*dtor));
  1203. dtor->dtor = g_free;
  1204. dtor->data = vec->iov_base;
  1205. LL_PREPEND (cbd->dtors, dtor);
  1206. memcpy (vec->iov_base, str, len);
  1207. vec->iov_len = len;
  1208. }
  1209. else {
  1210. msg_err ("bad argument at position %d", pos);
  1211. return FALSE;
  1212. }
  1213. return TRUE;
  1214. }
  1215. /***
  1216. * @function rspamd_tcp.request({params})
  1217. * This function creates and sends TCP request to the specified host and port,
  1218. * resolves hostname (if needed) and invokes continuation callback upon data received
  1219. * from the remote peer. This function accepts table of arguments with the following
  1220. * attributes
  1221. *
  1222. * - `task`: rspamd task objects (implies `pool`, `session`, `ev_base` and `resolver` arguments)
  1223. * - `ev_base`: event base (if no task specified)
  1224. * - `resolver`: DNS resolver (no task)
  1225. * - `session`: events session (no task)
  1226. * - `host`: IP or name of the peer (required)
  1227. * - `port`: remote port to use
  1228. * - `data`: a table of strings or `rspamd_text` objects that contains data pieces
  1229. * - `callback`: continuation function (required)
  1230. * - `on_connect`: callback called on connection success
  1231. * - `timeout`: floating point value that specifies timeout for IO operations in **seconds**
  1232. * - `partial`: boolean flag that specifies that callback should be called on any data portion received
  1233. * - `stop_pattern`: stop reading on finding a certain pattern (e.g. \r\n.\r\n for smtp)
  1234. * - `shutdown`: half-close socket after writing (boolean: default false)
  1235. * - `read`: read response after sending request (boolean: default true)
  1236. * @return {boolean} true if request has been sent
  1237. */
  1238. static gint
  1239. lua_tcp_request (lua_State *L)
  1240. {
  1241. LUA_TRACE_POINT;
  1242. const gchar *host;
  1243. gchar *stop_pattern = NULL;
  1244. guint port;
  1245. gint cbref, tp, conn_cbref = -1;
  1246. gsize plen = 0;
  1247. struct ev_loop *event_loop = NULL;
  1248. struct lua_tcp_cbdata *cbd;
  1249. struct rspamd_dns_resolver *resolver = NULL;
  1250. struct rspamd_async_session *session = NULL;
  1251. struct rspamd_task *task = NULL;
  1252. struct rspamd_config *cfg = NULL;
  1253. struct iovec *iov = NULL;
  1254. guint niov = 0, total_out;
  1255. guint64 h;
  1256. gdouble timeout = default_tcp_timeout;
  1257. gboolean partial = FALSE, do_shutdown = FALSE, do_read = TRUE,
  1258. ssl = FALSE, ssl_noverify = FALSE;
  1259. if (lua_type (L, 1) == LUA_TTABLE) {
  1260. lua_pushstring (L, "host");
  1261. lua_gettable (L, -2);
  1262. host = luaL_checkstring (L, -1);
  1263. lua_pop (L, 1);
  1264. lua_pushstring (L, "port");
  1265. lua_gettable (L, -2);
  1266. if (lua_type (L, -1) == LUA_TNUMBER) {
  1267. port = lua_tointeger (L, -1);
  1268. }
  1269. else {
  1270. /* We assume that it is a unix socket */
  1271. port = 0;
  1272. }
  1273. lua_pop (L, 1);
  1274. lua_pushstring (L, "callback");
  1275. lua_gettable (L, -2);
  1276. if (host == NULL || lua_type (L, -1) != LUA_TFUNCTION) {
  1277. lua_pop (L, 1);
  1278. msg_err ("tcp request has bad params");
  1279. lua_pushboolean (L, FALSE);
  1280. return 1;
  1281. }
  1282. cbref = luaL_ref (L, LUA_REGISTRYINDEX);
  1283. cbd = g_malloc0 (sizeof (*cbd));
  1284. lua_pushstring (L, "task");
  1285. lua_gettable (L, -2);
  1286. if (lua_type (L, -1) == LUA_TUSERDATA) {
  1287. task = lua_check_task (L, -1);
  1288. event_loop = task->event_loop;
  1289. resolver = task->resolver;
  1290. session = task->s;
  1291. cfg = task->cfg;
  1292. }
  1293. lua_pop (L, 1);
  1294. if (task == NULL) {
  1295. lua_pushstring (L, "ev_base");
  1296. lua_gettable (L, -2);
  1297. if (rspamd_lua_check_udata_maybe (L, -1, "rspamd{ev_base}")) {
  1298. event_loop = *(struct ev_loop **)lua_touserdata (L, -1);
  1299. }
  1300. else {
  1301. g_free (cbd);
  1302. return luaL_error (L, "event loop is required");
  1303. }
  1304. lua_pop (L, 1);
  1305. lua_pushstring (L, "session");
  1306. lua_gettable (L, -2);
  1307. if (rspamd_lua_check_udata_maybe (L, -1, "rspamd{session}")) {
  1308. session = *(struct rspamd_async_session **)lua_touserdata (L, -1);
  1309. }
  1310. else {
  1311. session = NULL;
  1312. }
  1313. lua_pop (L, 1);
  1314. lua_pushstring (L, "config");
  1315. lua_gettable (L, -2);
  1316. if (rspamd_lua_check_udata_maybe (L, -1, "rspamd{config}")) {
  1317. cfg = *(struct rspamd_config **)lua_touserdata (L, -1);
  1318. }
  1319. else {
  1320. cfg = NULL;
  1321. }
  1322. lua_pop (L, 1);
  1323. lua_pushstring (L, "resolver");
  1324. lua_gettable (L, -2);
  1325. if (rspamd_lua_check_udata_maybe (L, -1, "rspamd{resolver}")) {
  1326. resolver = *(struct rspamd_dns_resolver **)lua_touserdata (L, -1);
  1327. }
  1328. else {
  1329. resolver = lua_tcp_global_resolver (event_loop, cfg);
  1330. }
  1331. lua_pop (L, 1);
  1332. }
  1333. lua_pushstring (L, "timeout");
  1334. lua_gettable (L, -2);
  1335. if (lua_type (L, -1) == LUA_TNUMBER) {
  1336. timeout = lua_tonumber (L, -1);
  1337. }
  1338. lua_pop (L, 1);
  1339. lua_pushstring (L, "stop_pattern");
  1340. lua_gettable (L, -2);
  1341. if (lua_type (L, -1) == LUA_TSTRING) {
  1342. const gchar *p;
  1343. p = lua_tolstring (L, -1, &plen);
  1344. if (p && plen > 0) {
  1345. stop_pattern = g_malloc (plen);
  1346. memcpy (stop_pattern, p, plen);
  1347. }
  1348. }
  1349. lua_pop (L, 1);
  1350. lua_pushstring (L, "partial");
  1351. lua_gettable (L, -2);
  1352. if (lua_type (L, -1) == LUA_TBOOLEAN) {
  1353. partial = lua_toboolean (L, -1);
  1354. }
  1355. lua_pop (L, 1);
  1356. lua_pushstring (L, "shutdown");
  1357. lua_gettable (L, -2);
  1358. if (lua_type (L, -1) == LUA_TBOOLEAN) {
  1359. do_shutdown = lua_toboolean (L, -1);
  1360. }
  1361. lua_pop (L, 1);
  1362. lua_pushstring (L, "read");
  1363. lua_gettable (L, -2);
  1364. if (lua_type (L, -1) == LUA_TBOOLEAN) {
  1365. do_read = lua_toboolean (L, -1);
  1366. }
  1367. lua_pop (L, 1);
  1368. lua_pushstring (L, "ssl");
  1369. lua_gettable (L, -2);
  1370. if (lua_type (L, -1) == LUA_TBOOLEAN) {
  1371. ssl = lua_toboolean (L, -1);
  1372. }
  1373. lua_pop (L, 1);
  1374. lua_pushstring (L, "ssl_noverify");
  1375. lua_gettable (L, -2);
  1376. if (lua_type (L, -1) == LUA_TBOOLEAN) {
  1377. ssl_noverify = lua_toboolean (L, -1);
  1378. lua_pop (L, 1);
  1379. }
  1380. else {
  1381. lua_pop (L, 1); /* Previous nil... */
  1382. /* Similar to lua http, meh... */
  1383. lua_pushstring (L, "no_ssl_verify");
  1384. lua_gettable (L, -2);
  1385. if (lua_type (L, -1) == LUA_TBOOLEAN) {
  1386. ssl_noverify = lua_toboolean (L, -1);
  1387. }
  1388. lua_pop (L, 1);
  1389. }
  1390. lua_pushstring (L, "on_connect");
  1391. lua_gettable (L, -2);
  1392. if (lua_type (L, -1) == LUA_TFUNCTION) {
  1393. conn_cbref = luaL_ref (L, LUA_REGISTRYINDEX);
  1394. }
  1395. else {
  1396. lua_pop (L, 1);
  1397. }
  1398. lua_pushstring (L, "data");
  1399. lua_gettable (L, -2);
  1400. total_out = 0;
  1401. tp = lua_type (L, -1);
  1402. if (tp == LUA_TSTRING || tp == LUA_TUSERDATA) {
  1403. iov = g_malloc (sizeof (*iov));
  1404. niov = 1;
  1405. if (!lua_tcp_arg_toiovec (L, -1, cbd, iov)) {
  1406. lua_pop (L, 1);
  1407. msg_err ("tcp request has bad data argument");
  1408. lua_pushboolean (L, FALSE);
  1409. g_free (iov);
  1410. g_free (cbd);
  1411. return 1;
  1412. }
  1413. total_out = iov[0].iov_len;
  1414. }
  1415. else if (tp == LUA_TTABLE) {
  1416. /* Count parts */
  1417. lua_pushnil (L);
  1418. while (lua_next (L, -2) != 0) {
  1419. niov ++;
  1420. lua_pop (L, 1);
  1421. }
  1422. iov = g_malloc (sizeof (*iov) * niov);
  1423. lua_pushnil (L);
  1424. niov = 0;
  1425. while (lua_next (L, -2) != 0) {
  1426. if (!lua_tcp_arg_toiovec (L, -1, cbd, &iov[niov])) {
  1427. lua_pop (L, 2);
  1428. msg_err ("tcp request has bad data argument at pos %d", niov);
  1429. lua_pushboolean (L, FALSE);
  1430. g_free (iov);
  1431. g_free (cbd);
  1432. return 1;
  1433. }
  1434. total_out += iov[niov].iov_len;
  1435. niov ++;
  1436. lua_pop (L, 1);
  1437. }
  1438. }
  1439. lua_pop (L, 1);
  1440. }
  1441. else {
  1442. return luaL_error (L, "tcp request has bad params");
  1443. }
  1444. if (resolver == NULL && cfg == NULL && task == NULL) {
  1445. g_free (cbd);
  1446. g_free (iov);
  1447. return luaL_error (L, "tcp request has bad params: one of "
  1448. "{resolver,task,config} should be set");
  1449. }
  1450. cbd->task = task;
  1451. if (task) {
  1452. cbd->item = rspamd_symcache_get_cur_item (task);
  1453. }
  1454. cbd->cfg = cfg;
  1455. h = rspamd_random_uint64_fast ();
  1456. rspamd_snprintf (cbd->tag, sizeof (cbd->tag), "%uxL", h);
  1457. cbd->handlers = g_queue_new ();
  1458. cbd->hostname = g_strdup (host);
  1459. if (total_out > 0) {
  1460. struct lua_tcp_handler *wh;
  1461. wh = g_malloc0 (sizeof (*wh));
  1462. wh->type = LUA_WANT_WRITE;
  1463. wh->h.w.iov = iov;
  1464. wh->h.w.iovlen = niov;
  1465. wh->h.w.total_bytes = total_out;
  1466. wh->h.w.pos = 0;
  1467. /* Cannot set write handler here */
  1468. wh->h.w.cbref = -1;
  1469. if (cbref != -1 && !do_read) {
  1470. /* We have write only callback */
  1471. wh->h.w.cbref = cbref;
  1472. }
  1473. else {
  1474. /* We have simple client callback */
  1475. wh->h.w.cbref = -1;
  1476. }
  1477. g_queue_push_tail (cbd->handlers, wh);
  1478. }
  1479. cbd->event_loop = event_loop;
  1480. cbd->fd = -1;
  1481. cbd->port = port;
  1482. cbd->ev.timeout = timeout;
  1483. if (ssl) {
  1484. cbd->flags |= LUA_TCP_FLAG_SSL;
  1485. if (ssl_noverify) {
  1486. cbd->flags |= LUA_TCP_FLAG_SSL_NOVERIFY;
  1487. }
  1488. }
  1489. if (do_read) {
  1490. cbd->in = g_byte_array_sized_new (8192);
  1491. }
  1492. else {
  1493. /* Save some space... */
  1494. cbd->in = g_byte_array_new ();
  1495. }
  1496. if (partial) {
  1497. cbd->flags |= LUA_TCP_FLAG_PARTIAL;
  1498. }
  1499. if (do_shutdown) {
  1500. cbd->flags |= LUA_TCP_FLAG_SHUTDOWN;
  1501. }
  1502. if (do_read) {
  1503. struct lua_tcp_handler *rh;
  1504. rh = g_malloc0 (sizeof (*rh));
  1505. rh->type = LUA_WANT_READ;
  1506. rh->h.r.cbref = cbref;
  1507. rh->h.r.stop_pattern = stop_pattern;
  1508. rh->h.r.plen = plen;
  1509. g_queue_push_tail (cbd->handlers, rh);
  1510. }
  1511. cbd->connect_cb = conn_cbref;
  1512. REF_INIT_RETAIN (cbd, lua_tcp_maybe_free);
  1513. if (session) {
  1514. cbd->session = session;
  1515. if (rspamd_session_blocked (session)) {
  1516. lua_tcp_push_error (cbd, TRUE, "async session is the blocked state");
  1517. TCP_RELEASE (cbd);
  1518. cbd->item = NULL; /* To avoid decrease with no watcher */
  1519. lua_pushboolean (L, FALSE);
  1520. return 1;
  1521. }
  1522. }
  1523. if (rspamd_parse_inet_address (&cbd->addr,
  1524. host, strlen (host), RSPAMD_INET_ADDRESS_PARSE_DEFAULT)) {
  1525. rspamd_inet_address_set_port (cbd->addr, port);
  1526. /* Host is numeric IP, no need to resolve */
  1527. lua_tcp_register_watcher (cbd);
  1528. if (!lua_tcp_make_connection (cbd)) {
  1529. lua_tcp_push_error (cbd, TRUE, "cannot connect to the host: %s", host);
  1530. lua_pushboolean (L, FALSE);
  1531. /* No reset of the item as watcher has been registered */
  1532. TCP_RELEASE (cbd);
  1533. return 1;
  1534. }
  1535. }
  1536. else {
  1537. if (task == NULL) {
  1538. if (!rspamd_dns_resolver_request (resolver, session, NULL, lua_tcp_dns_handler, cbd,
  1539. RDNS_REQUEST_A, host)) {
  1540. lua_tcp_push_error (cbd, TRUE, "cannot resolve host: %s", host);
  1541. lua_pushboolean (L, FALSE);
  1542. cbd->item = NULL; /* To avoid decrease with no watcher */
  1543. TCP_RELEASE (cbd);
  1544. return 1;
  1545. }
  1546. else {
  1547. lua_tcp_register_watcher (cbd);
  1548. }
  1549. }
  1550. else {
  1551. if (!rspamd_dns_resolver_request_task (task, lua_tcp_dns_handler, cbd,
  1552. RDNS_REQUEST_A, host)) {
  1553. lua_tcp_push_error (cbd, TRUE, "cannot resolve host: %s", host);
  1554. lua_pushboolean (L, FALSE);
  1555. cbd->item = NULL; /* To avoid decrease with no watcher */
  1556. TCP_RELEASE (cbd);
  1557. return 1;
  1558. }
  1559. else {
  1560. lua_tcp_register_watcher (cbd);
  1561. }
  1562. }
  1563. }
  1564. lua_pushboolean (L, TRUE);
  1565. return 1;
  1566. }
  1567. /***
  1568. * @function rspamd_tcp.connect_sync({params})
  1569. * Creates new pseudo-synchronous connection to the specific address:port
  1570. *
  1571. * - `task`: rspamd task objects (implies `pool`, `session`, `ev_base` and `resolver` arguments)
  1572. * - `ev_base`: event base (if no task specified)
  1573. * - `resolver`: DNS resolver (no task)
  1574. * - `session`: events session (no task)
  1575. * - `config`: config (no task)
  1576. * - `host`: IP or name of the peer (required)
  1577. * - `port`: remote port to use
  1578. * - `timeout`: floating point value that specifies timeout for IO operations in **seconds**
  1579. * @return {boolean} true if request has been sent
  1580. */
  1581. static gint
  1582. lua_tcp_connect_sync (lua_State *L)
  1583. {
  1584. LUA_TRACE_POINT;
  1585. GError *err = NULL;
  1586. gint64 port = -1;
  1587. gdouble timeout = default_tcp_timeout;
  1588. const gchar *host = NULL;
  1589. gint ret;
  1590. guint64 h;
  1591. struct rspamd_task *task = NULL;
  1592. struct rspamd_async_session *session = NULL;
  1593. struct rspamd_dns_resolver *resolver = NULL;
  1594. struct rspamd_config *cfg = NULL;
  1595. struct ev_loop *ev_base = NULL;
  1596. struct lua_tcp_cbdata *cbd;
  1597. int arguments_validated = rspamd_lua_parse_table_arguments (L, 1, &err,
  1598. RSPAMD_LUA_PARSE_ARGUMENTS_DEFAULT,
  1599. "task=U{task};session=U{session};resolver=U{resolver};ev_base=U{ev_base};"
  1600. "*host=S;*port=I;timeout=D;config=U{config}",
  1601. &task, &session, &resolver, &ev_base,
  1602. &host, &port, &timeout, &cfg);
  1603. if (!arguments_validated) {
  1604. if (err) {
  1605. ret = luaL_error (L, "invalid arguments: %s", err->message);
  1606. g_error_free (err);
  1607. return ret;
  1608. }
  1609. return luaL_error (L, "invalid arguments");
  1610. }
  1611. if (0 > port || port > 65535) {
  1612. return luaL_error (L, "invalid port given (correct values: 1..65535)");
  1613. }
  1614. if (task == NULL && (cfg == NULL || ev_base == NULL || session == NULL)) {
  1615. return luaL_error (L, "invalid arguments: either task or config+ev_base+session should be set");
  1616. }
  1617. if (isnan (timeout)) {
  1618. /* rspamd_lua_parse_table_arguments() sets missing N field to zero */
  1619. timeout = default_tcp_timeout;
  1620. }
  1621. cbd = g_new0 (struct lua_tcp_cbdata, 1);
  1622. if (task) {
  1623. static const gchar hexdigests[16] = "0123456789abcdef";
  1624. cfg = task->cfg;
  1625. ev_base = task->event_loop;
  1626. session = task->s;
  1627. /* Make a readable tag */
  1628. memcpy (cbd->tag, task->task_pool->tag.uid, sizeof (cbd->tag) - 2);
  1629. cbd->tag[sizeof (cbd->tag) - 2] = hexdigests[GPOINTER_TO_INT (cbd) & 0xf];
  1630. cbd->tag[sizeof (cbd->tag) - 1] = 0;
  1631. }
  1632. else {
  1633. h = rspamd_random_uint64_fast ();
  1634. rspamd_snprintf (cbd->tag, sizeof (cbd->tag), "%uxL", h);
  1635. }
  1636. if (resolver == NULL) {
  1637. if (task) {
  1638. resolver = task->resolver;
  1639. }
  1640. else {
  1641. resolver = lua_tcp_global_resolver (ev_base, cfg);
  1642. }
  1643. }
  1644. cbd->task = task;
  1645. cbd->cfg = cfg;
  1646. cbd->thread = lua_thread_pool_get_running_entry (cfg->lua_thread_pool);
  1647. cbd->handlers = g_queue_new ();
  1648. cbd->event_loop = ev_base;
  1649. cbd->flags |= LUA_TCP_FLAG_SYNC;
  1650. cbd->fd = -1;
  1651. cbd->port = (guint16)port;
  1652. cbd->in = g_byte_array_new ();
  1653. cbd->connect_cb = -1;
  1654. REF_INIT_RETAIN (cbd, lua_tcp_maybe_free);
  1655. if (task) {
  1656. rspamd_mempool_add_destructor (task->task_pool, lua_tcp_sync_session_dtor, cbd);
  1657. }
  1658. struct lua_tcp_handler *wh;
  1659. wh = g_malloc0 (sizeof (*wh));
  1660. wh->type = LUA_WANT_CONNECT;
  1661. g_queue_push_tail (cbd->handlers, wh);
  1662. if (session) {
  1663. cbd->session = session;
  1664. if (rspamd_session_blocked (session)) {
  1665. TCP_RELEASE (cbd);
  1666. lua_pushboolean (L, FALSE);
  1667. lua_pushliteral (L, "Session is being destroyed, requests are not allowed");
  1668. return 2;
  1669. }
  1670. }
  1671. if (rspamd_parse_inet_address (&cbd->addr,
  1672. host, strlen (host), RSPAMD_INET_ADDRESS_PARSE_DEFAULT)) {
  1673. rspamd_inet_address_set_port (cbd->addr, (guint16)port);
  1674. /* Host is numeric IP, no need to resolve */
  1675. if (!lua_tcp_make_connection (cbd)) {
  1676. lua_pushboolean (L, FALSE);
  1677. lua_pushliteral (L, "Failed to initiate connection");
  1678. TCP_RELEASE (cbd);
  1679. return 2;
  1680. }
  1681. }
  1682. else {
  1683. if (task == NULL) {
  1684. if (!rspamd_dns_resolver_request (resolver, session, NULL, lua_tcp_dns_handler, cbd,
  1685. RDNS_REQUEST_A, host)) {
  1686. lua_pushboolean (L, FALSE);
  1687. lua_pushliteral (L, "Failed to initiate dns request");
  1688. TCP_RELEASE (cbd);
  1689. return 2;
  1690. }
  1691. else {
  1692. lua_tcp_register_watcher (cbd);
  1693. }
  1694. }
  1695. else {
  1696. cbd->item = rspamd_symcache_get_cur_item (task);
  1697. if (!rspamd_dns_resolver_request_task (task, lua_tcp_dns_handler, cbd,
  1698. RDNS_REQUEST_A, host)) {
  1699. lua_pushboolean (L, FALSE);
  1700. lua_pushliteral (L, "Failed to initiate dns request");
  1701. TCP_RELEASE (cbd);
  1702. return 2;
  1703. }
  1704. else {
  1705. lua_tcp_register_watcher (cbd);
  1706. }
  1707. }
  1708. }
  1709. return lua_thread_yield (cbd->thread, 0);
  1710. }
  1711. static gint
  1712. lua_tcp_close (lua_State *L)
  1713. {
  1714. LUA_TRACE_POINT;
  1715. struct lua_tcp_cbdata *cbd = lua_check_tcp (L, 1);
  1716. if (cbd == NULL) {
  1717. return luaL_error (L, "invalid arguments");
  1718. }
  1719. cbd->flags |= LUA_TCP_FLAG_FINISHED;
  1720. TCP_RELEASE (cbd);
  1721. return 0;
  1722. }
  1723. static gint
  1724. lua_tcp_add_read (lua_State *L)
  1725. {
  1726. LUA_TRACE_POINT;
  1727. struct lua_tcp_cbdata *cbd = lua_check_tcp (L, 1);
  1728. struct lua_tcp_handler *rh;
  1729. gchar *stop_pattern = NULL;
  1730. const gchar *p;
  1731. gsize plen = 0;
  1732. gint cbref = -1;
  1733. if (cbd == NULL) {
  1734. return luaL_error (L, "invalid arguments");
  1735. }
  1736. if (lua_type (L, 2) == LUA_TFUNCTION) {
  1737. lua_pushvalue (L, 2);
  1738. cbref = luaL_ref (L, LUA_REGISTRYINDEX);
  1739. }
  1740. if (lua_type (L, 3) == LUA_TSTRING) {
  1741. p = lua_tolstring (L, 3, &plen);
  1742. if (p && plen > 0) {
  1743. stop_pattern = g_malloc (plen);
  1744. memcpy (stop_pattern, p, plen);
  1745. }
  1746. }
  1747. rh = g_malloc0 (sizeof (*rh));
  1748. rh->type = LUA_WANT_READ;
  1749. rh->h.r.cbref = cbref;
  1750. rh->h.r.stop_pattern = stop_pattern;
  1751. rh->h.r.plen = plen;
  1752. msg_debug_tcp ("added read event, cbref: %d", cbref);
  1753. g_queue_push_tail (cbd->handlers, rh);
  1754. return 0;
  1755. }
  1756. static gint
  1757. lua_tcp_add_write (lua_State *L)
  1758. {
  1759. LUA_TRACE_POINT;
  1760. struct lua_tcp_cbdata *cbd = lua_check_tcp (L, 1);
  1761. struct lua_tcp_handler *wh;
  1762. gint cbref = -1, tp;
  1763. struct iovec *iov = NULL;
  1764. guint niov = 0, total_out = 0;
  1765. if (cbd == NULL) {
  1766. return luaL_error (L, "invalid arguments");
  1767. }
  1768. if (lua_type (L, 2) == LUA_TFUNCTION) {
  1769. lua_pushvalue (L, 2);
  1770. cbref = luaL_ref (L, LUA_REGISTRYINDEX);
  1771. }
  1772. tp = lua_type (L, 3);
  1773. if (tp == LUA_TSTRING || tp == LUA_TUSERDATA) {
  1774. iov = g_malloc (sizeof (*iov));
  1775. niov = 1;
  1776. if (!lua_tcp_arg_toiovec (L, 3, cbd, iov)) {
  1777. msg_err ("tcp request has bad data argument");
  1778. lua_pushboolean (L, FALSE);
  1779. g_free (iov);
  1780. return 1;
  1781. }
  1782. total_out = iov[0].iov_len;
  1783. }
  1784. else if (tp == LUA_TTABLE) {
  1785. /* Count parts */
  1786. lua_pushvalue (L, 3);
  1787. lua_pushnil (L);
  1788. while (lua_next (L, -2) != 0) {
  1789. niov ++;
  1790. lua_pop (L, 1);
  1791. }
  1792. iov = g_malloc (sizeof (*iov) * niov);
  1793. lua_pushnil (L);
  1794. niov = 0;
  1795. while (lua_next (L, -2) != 0) {
  1796. if (!lua_tcp_arg_toiovec (L, -1, cbd, &iov[niov])) {
  1797. lua_pop (L, 2);
  1798. msg_err ("tcp request has bad data argument at pos %d", niov);
  1799. lua_pushboolean (L, FALSE);
  1800. g_free (iov);
  1801. g_free (cbd);
  1802. return 1;
  1803. }
  1804. total_out += iov[niov].iov_len;
  1805. niov ++;
  1806. lua_pop (L, 1);
  1807. }
  1808. lua_pop (L, 1);
  1809. }
  1810. wh = g_malloc0 (sizeof (*wh));
  1811. wh->type = LUA_WANT_WRITE;
  1812. wh->h.w.iov = iov;
  1813. wh->h.w.iovlen = niov;
  1814. wh->h.w.total_bytes = total_out;
  1815. wh->h.w.pos = 0;
  1816. /* Cannot set write handler here */
  1817. wh->h.w.cbref = cbref;
  1818. msg_debug_tcp ("added write event, cbref: %d", cbref);
  1819. g_queue_push_tail (cbd->handlers, wh);
  1820. lua_pushboolean (L, TRUE);
  1821. return 1;
  1822. }
  1823. static gint
  1824. lua_tcp_shift_callback (lua_State *L)
  1825. {
  1826. LUA_TRACE_POINT;
  1827. struct lua_tcp_cbdata *cbd = lua_check_tcp (L, 1);
  1828. if (cbd == NULL) {
  1829. return luaL_error (L, "invalid arguments");
  1830. }
  1831. lua_tcp_shift_handler (cbd);
  1832. lua_tcp_plan_handler_event (cbd, TRUE, TRUE);
  1833. return 0;
  1834. }
  1835. static struct lua_tcp_cbdata *
  1836. lua_check_sync_tcp (lua_State *L, gint pos)
  1837. {
  1838. void *ud = rspamd_lua_check_udata (L, pos, "rspamd{tcp_sync}");
  1839. luaL_argcheck (L, ud != NULL, pos, "'tcp' expected");
  1840. return ud ? *((struct lua_tcp_cbdata **)ud) : NULL;
  1841. }
  1842. static int
  1843. lua_tcp_sync_close (lua_State *L)
  1844. {
  1845. LUA_TRACE_POINT;
  1846. struct lua_tcp_cbdata *cbd = lua_check_sync_tcp (L, 1);
  1847. if (cbd == NULL) {
  1848. return luaL_error (L, "invalid arguments [self is not rspamd{tcp_sync}]");
  1849. }
  1850. cbd->flags |= LUA_TCP_FLAG_FINISHED;
  1851. if (cbd->fd != -1) {
  1852. rspamd_ev_watcher_stop (cbd->event_loop, &cbd->ev);
  1853. close (cbd->fd);
  1854. cbd->fd = -1;
  1855. }
  1856. return 0;
  1857. }
  1858. static void
  1859. lua_tcp_sync_session_dtor (gpointer ud)
  1860. {
  1861. struct lua_tcp_cbdata *cbd = ud;
  1862. cbd->flags |= LUA_TCP_FLAG_FINISHED;
  1863. if (cbd->fd != -1) {
  1864. msg_debug ("closing sync TCP connection");
  1865. rspamd_ev_watcher_stop (cbd->event_loop, &cbd->ev);
  1866. close (cbd->fd);
  1867. cbd->fd = -1;
  1868. }
  1869. /* Task is gone, we should not try use it anymore */
  1870. cbd->task = NULL;
  1871. /* All events are removed when task is done, we should not refer them */
  1872. cbd->async_ev = NULL;
  1873. }
  1874. static int
  1875. lua_tcp_sync_read_once (lua_State *L)
  1876. {
  1877. LUA_TRACE_POINT;
  1878. struct lua_tcp_cbdata *cbd = lua_check_sync_tcp (L, 1);
  1879. struct lua_tcp_handler *rh;
  1880. if (cbd == NULL) {
  1881. return luaL_error (L, "invalid arguments [self is not rspamd{tcp_sync}]");
  1882. }
  1883. struct thread_entry *thread = lua_thread_pool_get_running_entry (cbd->cfg->lua_thread_pool);
  1884. rh = g_malloc0 (sizeof (*rh));
  1885. rh->type = LUA_WANT_READ;
  1886. rh->h.r.cbref = -1;
  1887. msg_debug_tcp ("added read sync event, thread: %p", thread);
  1888. g_queue_push_tail (cbd->handlers, rh);
  1889. lua_tcp_plan_handler_event (cbd, TRUE, TRUE);
  1890. TCP_RETAIN (cbd);
  1891. return lua_thread_yield (thread, 0);
  1892. }
  1893. static int
  1894. lua_tcp_sync_write (lua_State *L)
  1895. {
  1896. LUA_TRACE_POINT;
  1897. struct lua_tcp_cbdata *cbd = lua_check_sync_tcp (L, 1);
  1898. struct lua_tcp_handler *wh;
  1899. gint tp;
  1900. struct iovec *iov = NULL;
  1901. guint niov = 0;
  1902. gsize total_out = 0;
  1903. if (cbd == NULL) {
  1904. return luaL_error (L, "invalid arguments [self is not rspamd{tcp_sync}]");
  1905. }
  1906. struct thread_entry *thread = lua_thread_pool_get_running_entry (cbd->cfg->lua_thread_pool);
  1907. tp = lua_type (L, 2);
  1908. if (tp == LUA_TSTRING || tp == LUA_TUSERDATA) {
  1909. iov = g_malloc (sizeof (*iov));
  1910. niov = 1;
  1911. if (!lua_tcp_arg_toiovec (L, 2, cbd, iov)) {
  1912. msg_err ("tcp request has bad data argument");
  1913. g_free (iov);
  1914. g_free (cbd);
  1915. return luaL_error (L, "invalid arguments second parameter (data) is expected to be either string or rspamd{text}");
  1916. }
  1917. total_out = iov[0].iov_len;
  1918. }
  1919. else if (tp == LUA_TTABLE) {
  1920. /* Count parts */
  1921. lua_pushvalue (L, 3);
  1922. lua_pushnil (L);
  1923. while (lua_next (L, -2) != 0) {
  1924. niov ++;
  1925. lua_pop (L, 1);
  1926. }
  1927. iov = g_malloc (sizeof (*iov) * niov);
  1928. lua_pushnil (L);
  1929. niov = 0;
  1930. while (lua_next (L, -2) != 0) {
  1931. if (!lua_tcp_arg_toiovec (L, -1, cbd, &iov[niov])) {
  1932. msg_err ("tcp request has bad data argument at pos %d", niov);
  1933. g_free (iov);
  1934. g_free (cbd);
  1935. return luaL_error (L, "invalid arguments second parameter (data) is expected to be either string or rspamd{text}");
  1936. }
  1937. total_out += iov[niov].iov_len;
  1938. niov ++;
  1939. lua_pop (L, 1);
  1940. }
  1941. lua_pop (L, 1);
  1942. }
  1943. wh = g_malloc0 (sizeof (*wh));
  1944. wh->type = LUA_WANT_WRITE;
  1945. wh->h.w.iov = iov;
  1946. wh->h.w.iovlen = niov;
  1947. wh->h.w.total_bytes = total_out;
  1948. wh->h.w.pos = 0;
  1949. wh->h.w.cbref = -1;
  1950. msg_debug_tcp ("added sync write event, thread: %p", thread);
  1951. g_queue_push_tail (cbd->handlers, wh);
  1952. lua_tcp_plan_handler_event (cbd, TRUE, TRUE);
  1953. TCP_RETAIN (cbd);
  1954. return lua_thread_yield (thread, 0);
  1955. }
  1956. static gint
  1957. lua_tcp_sync_eof(lua_State *L)
  1958. {
  1959. LUA_TRACE_POINT;
  1960. struct lua_tcp_cbdata *cbd = lua_check_sync_tcp (L, 1);
  1961. if (cbd == NULL) {
  1962. return luaL_error (L, "invalid arguments [self is not rspamd{tcp_sync}]");
  1963. }
  1964. lua_pushboolean(L, cbd->eof);
  1965. return 1;
  1966. }
  1967. static gint
  1968. lua_tcp_sync_shutdown (lua_State *L)
  1969. {
  1970. LUA_TRACE_POINT;
  1971. struct lua_tcp_cbdata *cbd = lua_check_sync_tcp (L, 1);
  1972. if (cbd == NULL) {
  1973. return luaL_error (L, "invalid arguments [self is not rspamd{tcp_sync}]");
  1974. }
  1975. shutdown (cbd->fd, SHUT_WR);
  1976. return 0;
  1977. }
  1978. static gint
  1979. lua_tcp_starttls (lua_State * L)
  1980. {
  1981. LUA_TRACE_POINT;
  1982. struct lua_tcp_cbdata *cbd = lua_check_tcp (L, 1);
  1983. gpointer ssl_ctx;
  1984. gboolean verify_peer;
  1985. if (cbd == NULL || cbd->ssl_conn != NULL) {
  1986. return luaL_error (L, "invalid arguments");
  1987. }
  1988. if (cbd->flags & LUA_TCP_FLAG_SSL_NOVERIFY) {
  1989. ssl_ctx = cbd->cfg->libs_ctx->ssl_ctx_noverify;
  1990. verify_peer = FALSE;
  1991. }
  1992. else {
  1993. ssl_ctx = cbd->cfg->libs_ctx->ssl_ctx;
  1994. verify_peer = TRUE;
  1995. }
  1996. cbd->ssl_conn = rspamd_ssl_connection_new (ssl_ctx,
  1997. cbd->event_loop,
  1998. verify_peer,
  1999. cbd->tag);
  2000. if (!rspamd_ssl_connect_fd (cbd->ssl_conn, cbd->fd, cbd->hostname, &cbd->ev,
  2001. cbd->ev.timeout, lua_tcp_handler, lua_tcp_ssl_on_error, cbd)) {
  2002. lua_tcp_push_error (cbd, TRUE, "ssl connection failed: %s",
  2003. strerror (errno));
  2004. }
  2005. return 0;
  2006. }
  2007. static gint
  2008. lua_tcp_sync_gc (lua_State * L)
  2009. {
  2010. struct lua_tcp_cbdata *cbd = lua_check_sync_tcp (L, 1);
  2011. if (!cbd) {
  2012. return luaL_error (L, "invalid arguments [self is not rspamd{tcp_sync}]");
  2013. }
  2014. lua_tcp_maybe_free(cbd);
  2015. lua_tcp_fin (cbd);
  2016. return 0;
  2017. }
  2018. static gint
  2019. lua_load_tcp (lua_State * L)
  2020. {
  2021. lua_newtable (L);
  2022. luaL_register (L, NULL, tcp_libf);
  2023. return 1;
  2024. }
  2025. void
  2026. luaopen_tcp (lua_State * L)
  2027. {
  2028. rspamd_lua_add_preload (L, "rspamd_tcp", lua_load_tcp);
  2029. rspamd_lua_new_class (L, "rspamd{tcp}", tcp_libm);
  2030. rspamd_lua_new_class (L, "rspamd{tcp_sync}", tcp_sync_libm);
  2031. lua_pop (L, 1);
  2032. }