xref: /trafficserver/plugins/lua/ts_lua_fetch.c (revision 7236e78f)
1 /*
2   Licensed to the Apache Software Foundation (ASF) under one
3   or more contributor license agreements.  See the NOTICE file
4   distributed with this work for additional information
5   regarding copyright ownership.  The ASF licenses this file
6   to you under the Apache License, Version 2.0 (the
7   "License"); you may not use this file except in compliance
8   with the License.  You may obtain a copy of the License at
9 
10   http://www.apache.org/licenses/LICENSE-2.0
11 
12   Unless required by applicable law or agreed to in writing, software
13   distributed under the License is distributed on an "AS IS" BASIS,
14   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15   See the License for the specific language governing permissions and
16   limitations under the License.
17 */
18 
19 #include "ts_lua_util.h"
20 #include "ts_lua_io.h"
21 #include "ts_lua_fetch.h"
22 
23 #define TS_LUA_EVENT_FETCH_OVER 20010
24 #define TS_LUA_FETCH_CLIENT_ADDRPORT "127.0.0.1:33333"
25 #define TS_LUA_FETCH_CLIENT_ADDRPORT_LEN 15
26 #define TS_LUA_FETCH_USER_AGENT "TS Fetcher/1.0"
27 
28 static int ts_lua_fetch(lua_State *L);
29 static int ts_lua_fetch_multi(lua_State *L);
30 static int ts_lua_fetch_handler(TSCont contp, TSEvent event, void *edata);
31 static int ts_lua_fetch_multi_cleanup(ts_lua_async_item *ai);
32 static int ts_lua_fetch_multi_handler(TSCont contp, TSEvent event, void *edata);
33 static int ts_lua_fetch_one_item(lua_State *L, const char *url, size_t url_len, ts_lua_fetch_info *fi);
34 static inline void ts_lua_destroy_fetch_multi_info(ts_lua_fetch_multi_info *fmi);
35 
36 void
ts_lua_inject_fetch_api(lua_State * L)37 ts_lua_inject_fetch_api(lua_State *L)
38 {
39   /* ts.fetch() */
40   lua_pushcfunction(L, ts_lua_fetch);
41   lua_setfield(L, -2, "fetch");
42 
43   /* ts.fetch_multi() */
44   lua_pushcfunction(L, ts_lua_fetch_multi);
45   lua_setfield(L, -2, "fetch_multi");
46 }
47 
48 static int
ts_lua_fetch(lua_State * L)49 ts_lua_fetch(lua_State *L)
50 {
51   int sz;
52   size_t n;
53   const char *url;
54   size_t url_len;
55   TSCont contp;
56   ts_lua_cont_info *ci;
57   ts_lua_async_item *ai;
58   ts_lua_fetch_info *fi;
59   ts_lua_fetch_multi_info *fmi;
60 
61   ci = ts_lua_get_cont_info(L);
62   if (ci == NULL) {
63     return 0;
64   }
65 
66   n = lua_gettop(L);
67   if (n < 1) {
68     return luaL_error(L, "'ts.fetch' requires parameter");
69   }
70 
71   /* url */
72   if (!lua_isstring(L, 1)) {
73     return luaL_error(L, "'ts.fetch' first param is not string");
74   }
75 
76   url = luaL_checklstring(L, 1, &url_len);
77 
78   /* replicate misc table */
79   if (n >= 2) {
80     lua_pushvalue(L, 2);
81 
82   } else {
83     lua_pushnil(L);
84   }
85 
86   contp = TSContCreate(ts_lua_fetch_multi_handler, ci->mutex);
87 
88   sz  = sizeof(ts_lua_fetch_multi_info) + 1 * sizeof(ts_lua_fetch_info);
89   fmi = (ts_lua_fetch_multi_info *)TSmalloc(sz);
90 
91   memset(fmi, 0, sz);
92   fmi->total = 1;
93   fmi->contp = contp;
94 
95   fi         = &fmi->fiv[0];
96   fi->fmi    = fmi;
97   fi->buffer = TSIOBufferCreate();
98   fi->reader = TSIOBufferReaderAlloc(fi->buffer);
99 
100   ts_lua_fetch_one_item(L, url, url_len, fi);
101 
102   // pop the replicated misc table
103   lua_pop(L, 1);
104 
105   ai = ts_lua_async_create_item(contp, ts_lua_fetch_multi_cleanup, fmi, ci);
106   TSContDataSet(contp, ai);
107 
108   return lua_yield(L, 0);
109   ;
110 }
111 
112 static int
ts_lua_fetch_multi(lua_State * L)113 ts_lua_fetch_multi(lua_State *L)
114 {
115   int type, sz;
116   size_t i, n;
117   const char *url;
118   size_t url_len;
119   TSCont contp;
120   ts_lua_cont_info *ci;
121   ts_lua_async_item *ai;
122   ts_lua_fetch_info *fi;
123   ts_lua_fetch_multi_info *fmi;
124 
125   ci = ts_lua_get_cont_info(L);
126   if (ci == NULL) {
127     return 0;
128   }
129 
130   if (lua_gettop(L) < 1) {
131     return luaL_error(L, "'ts.fetch_mutli' requires one parameter");
132   }
133 
134   type = lua_type(L, 1);
135   if (type != LUA_TTABLE) {
136     return luaL_error(L, "'ts.fetch_mutli' requires table as parameter");
137   }
138 
139   // main continuation handler
140   contp = TSContCreate(ts_lua_fetch_multi_handler, ci->mutex);
141 
142   // Iterate the table
143   n = lua_objlen(L, 1);
144 
145   sz  = sizeof(ts_lua_fetch_multi_info) + n * sizeof(ts_lua_fetch_info);
146   fmi = (ts_lua_fetch_multi_info *)TSmalloc(sz);
147 
148   memset(fmi, 0, sz);
149   fmi->total = n;
150   fmi->contp = contp;
151   fmi->multi = 1;
152 
153   for (i = 0; i < n; i++) {
154     /* push fetch item */
155     lua_pushinteger(L, i + 1);
156     lua_gettable(L, -2);
157 
158     if (lua_objlen(L, -1) < 1) {
159       ts_lua_destroy_fetch_multi_info(fmi);
160       TSContDestroy(contp);
161 
162       return luaL_error(L, "'ts.fetch_mutli' got empty table item");
163     }
164 
165     /* push url */
166     lua_pushnumber(L, 1);
167     lua_gettable(L, -2);
168 
169     if (!lua_isstring(L, -1)) {
170       ts_lua_destroy_fetch_multi_info(fmi);
171       TSContDestroy(contp);
172 
173       return luaL_error(L, "'ts.fetch_mutli' got invalid table item: url illegal");
174     }
175 
176     url = luaL_checklstring(L, -1, &url_len);
177 
178     /* push misc table */
179     lua_pushinteger(L, 2);
180     lua_gettable(L, -3);
181 
182     fi         = &fmi->fiv[i];
183     fi->fmi    = fmi;
184     fi->buffer = TSIOBufferCreate();
185     fi->reader = TSIOBufferReaderAlloc(fi->buffer);
186 
187     ts_lua_fetch_one_item(L, url, url_len, fi);
188     lua_pop(L, 3); // misc table, url, fetch item
189   }
190 
191   ai = ts_lua_async_create_item(contp, ts_lua_fetch_multi_cleanup, (void *)fmi, ci);
192   TSContDataSet(contp, ai);
193 
194   return lua_yield(L, 0);
195   ;
196 }
197 
198 static int
ts_lua_fetch_one_item(lua_State * L,const char * url,size_t url_len,ts_lua_fetch_info * fi)199 ts_lua_fetch_one_item(lua_State *L, const char *url, size_t url_len, ts_lua_fetch_info *fi)
200 {
201   TSCont contp;
202   int tb, flags, host_len, n;
203   int cl, ht, ua;
204   const char *method, *key, *value, *body, *opt;
205   const char *addr, *ptr, *host;
206   size_t method_len, key_len, value_len, body_len;
207   size_t addr_len, opt_len, i, left;
208   char c;
209   struct sockaddr clientaddr;
210   char buf[32];
211 
212   tb = lua_istable(L, -1);
213 
214   /* method */
215   if (tb) {
216     lua_pushlstring(L, "method", sizeof("method") - 1);
217     lua_gettable(L, -2);
218     if (lua_isstring(L, -1)) {
219       method = luaL_checklstring(L, -1, &method_len);
220 
221     } else {
222       method     = "GET";
223       method_len = sizeof("GET") - 1;
224     }
225 
226     lua_pop(L, 1);
227 
228   } else {
229     method     = "GET";
230     method_len = sizeof("GET") - 1;
231   }
232 
233   /* body */
234   body     = NULL;
235   body_len = 0;
236 
237   if (tb) {
238     lua_pushlstring(L, "body", sizeof("body") - 1);
239     lua_gettable(L, -2);
240 
241     if (lua_isstring(L, -1)) {
242       body = luaL_checklstring(L, -1, &body_len);
243     }
244 
245     lua_pop(L, 1);
246   }
247 
248   /* cliaddr */
249   memset(&clientaddr, 0, sizeof(clientaddr));
250 
251   if (tb) {
252     lua_pushlstring(L, "cliaddr", sizeof("cliaddr") - 1);
253     lua_gettable(L, -2);
254 
255     if (lua_isstring(L, -1)) {
256       addr = luaL_checklstring(L, -1, &addr_len);
257 
258       if (TS_ERROR == TSIpStringToAddr(addr, addr_len, &clientaddr)) {
259         TSError("[ts_lua][%s] Client ip parse failed! Using default.", TS_LUA_DEBUG_TAG);
260         if (TS_ERROR == TSIpStringToAddr(TS_LUA_FETCH_CLIENT_ADDRPORT, TS_LUA_FETCH_CLIENT_ADDRPORT_LEN, &clientaddr)) {
261           TSError("[ts_lua][%s] Default client ip parse failed!", TS_LUA_DEBUG_TAG);
262           return 0;
263         }
264       }
265     }
266 
267     lua_pop(L, 1);
268   }
269 
270   /* option */
271   flags = TS_FETCH_FLAGS_DECHUNK; // dechunk the body default
272 
273   if (tb) {
274     lua_pushlstring(L, "option", sizeof("option") - 1);
275     lua_gettable(L, -2);
276 
277     if (lua_isstring(L, -1)) {
278       opt = luaL_checklstring(L, -1, &opt_len);
279 
280       for (i = 0; i < opt_len; i++) {
281         c = opt[i];
282 
283         switch (c) {
284         case 'c':
285           flags &= (~TS_FETCH_FLAGS_DECHUNK);
286           break;
287 
288         default:
289           break;
290         }
291       }
292     }
293 
294     lua_pop(L, 1);
295   }
296 
297   contp = TSContCreate(ts_lua_fetch_handler, TSContMutexGet(fi->fmi->contp)); // reuse parent cont's mutex
298   TSContDataSet(contp, fi);
299 
300   fi->contp = contp;
301   fi->fch   = TSFetchCreate(contp, method, url, "HTTP/1.1", &clientaddr, flags);
302 
303   /* header */
304   cl = ht = ua = 0;
305 
306   if (tb) {
307     lua_pushlstring(L, "header", sizeof("header") - 1);
308     lua_gettable(L, -2);
309 
310     if (lua_istable(L, -1)) {
311       // iterate the header table
312       lua_pushnil(L);
313 
314       while (lua_next(L, -2)) {
315         lua_pushvalue(L, -2);
316 
317         key   = luaL_checklstring(L, -1, &key_len);
318         value = luaL_checklstring(L, -2, &value_len);
319 
320         if ((int)key_len == TS_MIME_LEN_CONTENT_LENGTH &&
321             !strncasecmp(TS_MIME_FIELD_CONTENT_LENGTH, key, key_len)) { // Content-Length
322           cl = 1;
323 
324         } else if ((int)key_len == TS_MIME_LEN_HOST && !strncasecmp(TS_MIME_FIELD_HOST, key, key_len)) { // Host
325           ht = 1;
326 
327         } else if ((int)key_len == TS_MIME_LEN_USER_AGENT && !strncasecmp(TS_MIME_FIELD_USER_AGENT, key, key_len)) { // User-Agent
328           ua = 1;
329         }
330 
331         TSFetchHeaderAdd(fi->fch, key, key_len, value, value_len);
332 
333         lua_pop(L, 2);
334       }
335     }
336 
337     lua_pop(L, 1);
338   }
339 
340   /* Host */
341   if (ht == 0) {
342     ptr = memchr(url, ':', url_len);
343 
344     if (ptr) {
345       host = ptr + 3;
346       left = url_len - (host - url);
347 
348       ptr = memchr(host, '/', left);
349 
350       if (ptr) {
351         host_len = ptr - host;
352 
353       } else {
354         host_len = left;
355       }
356 
357       TSFetchHeaderAdd(fi->fch, TS_MIME_FIELD_HOST, TS_MIME_LEN_HOST, host, host_len);
358     }
359   }
360 
361   /* User-Agent */
362   if (ua == 0) {
363     TSFetchHeaderAdd(fi->fch, TS_MIME_FIELD_USER_AGENT, TS_MIME_LEN_USER_AGENT, TS_LUA_FETCH_USER_AGENT,
364                      sizeof(TS_LUA_FETCH_USER_AGENT) - 1);
365   }
366 
367   if (body_len > 0 && cl == 0) { // add Content-Length header
368     n = sprintf(buf, "%zu", body_len);
369     TSFetchHeaderAdd(fi->fch, TS_MIME_FIELD_CONTENT_LENGTH, TS_MIME_LEN_CONTENT_LENGTH, buf, n);
370   }
371 
372   TSFetchLaunch(fi->fch);
373 
374   if (body_len > 0) {
375     TSFetchWriteData(fi->fch, body, body_len);
376   }
377 
378   return 0;
379 }
380 
381 static int
ts_lua_fetch_handler(TSCont contp,TSEvent ev,void * edata ATS_UNUSED)382 ts_lua_fetch_handler(TSCont contp, TSEvent ev, void *edata ATS_UNUSED)
383 {
384   int event;
385   char *from;
386   int64_t n, wavail;
387   TSIOBufferBlock blk;
388 
389   ts_lua_fetch_info *fi;
390   ts_lua_fetch_multi_info *fmi;
391 
392   event = (int)ev;
393   fi    = TSContDataGet(contp);
394   fmi   = fi->fmi;
395 
396   switch (event) {
397   case TS_FETCH_EVENT_EXT_HEAD_READY:
398   case TS_FETCH_EVENT_EXT_HEAD_DONE:
399     break;
400 
401   case TS_FETCH_EVENT_EXT_BODY_READY:
402   case TS_FETCH_EVENT_EXT_BODY_DONE:
403 
404     do {
405       blk  = TSIOBufferStart(fi->buffer);
406       from = TSIOBufferBlockWriteStart(blk, &wavail);
407       n    = TSFetchReadData(fi->fch, from, wavail);
408       TSIOBufferProduce(fi->buffer, n);
409     } while (n == wavail);
410 
411     if (event == TS_FETCH_EVENT_EXT_BODY_DONE) { // fetch over
412       fi->over = 1;
413     }
414 
415     break;
416 
417   default:
418     fi->failed = 1;
419     break;
420   }
421 
422   if (fmi && (fi->over || fi->failed)) {
423     TSContCall(fmi->contp, TS_LUA_EVENT_FETCH_OVER, fi); // error exist
424     ts_lua_destroy_fetch_multi_info(fmi);
425   }
426 
427   return 0;
428 }
429 
430 static int
ts_lua_fill_one_result(lua_State * L,ts_lua_fetch_info * fi)431 ts_lua_fill_one_result(lua_State *L, ts_lua_fetch_info *fi)
432 {
433   const char *name, *value;
434   int name_len, value_len;
435   char *dst;
436   int64_t ravail;
437   TSMBuffer bufp;
438   TSMLoc hdrp;
439   TSMLoc field_loc, next_field_loc;
440   TSHttpStatus status;
441 
442   bufp = TSFetchRespHdrMBufGet(fi->fch);
443   hdrp = TSFetchRespHdrMLocGet(fi->fch);
444 
445   // result table
446   lua_newtable(L);
447 
448   // status code
449   status = TSHttpHdrStatusGet(bufp, hdrp);
450   lua_pushlstring(L, "status", sizeof("status") - 1);
451   lua_pushnumber(L, status);
452   lua_rawset(L, -3);
453 
454   // header
455   lua_pushlstring(L, "header", sizeof("header") - 1);
456   lua_newtable(L);
457 
458   field_loc = TSMimeHdrFieldGet(bufp, hdrp, 0);
459   while (field_loc) {
460     name  = TSMimeHdrFieldNameGet(bufp, hdrp, field_loc, &name_len);
461     value = TSMimeHdrFieldValueStringGet(bufp, hdrp, field_loc, -1, &value_len);
462 
463     lua_pushlstring(L, name, name_len);
464     lua_pushlstring(L, value, value_len);
465     lua_rawset(L, -3);
466 
467     next_field_loc = TSMimeHdrFieldNext(bufp, hdrp, field_loc);
468     TSHandleMLocRelease(bufp, hdrp, field_loc);
469     field_loc = next_field_loc;
470   }
471   lua_rawset(L, -3);
472 
473   // body
474   ravail = TSIOBufferReaderAvail(fi->reader);
475   if (ravail > 0) {
476     lua_pushlstring(L, "body", sizeof("body") - 1);
477 
478     dst = (char *)TSmalloc(ravail);
479     IOBufferReaderCopy(fi->reader, dst, ravail);
480     lua_pushlstring(L, (char *)dst, ravail);
481 
482     lua_rawset(L, -3);
483     TSfree(dst);
484   }
485 
486   // truncated
487   lua_pushlstring(L, "truncated", sizeof("truncated") - 1);
488   if (fi->failed) {
489     lua_pushboolean(L, 1);
490 
491   } else {
492     lua_pushboolean(L, 0);
493   }
494 
495   lua_rawset(L, -3);
496 
497   return 0;
498 }
499 
500 static int
ts_lua_fetch_multi_handler(TSCont contp,TSEvent event ATS_UNUSED,void * edata)501 ts_lua_fetch_multi_handler(TSCont contp, TSEvent event ATS_UNUSED, void *edata)
502 {
503   int i;
504   lua_State *L;
505   TSMutex lmutex;
506 
507   ts_lua_async_item *ai;
508   ts_lua_cont_info *ci;
509   ts_lua_fetch_info *fi;
510   ts_lua_fetch_multi_info *fmi;
511 
512   ai = TSContDataGet(contp);
513   ci = ai->cinfo;
514 
515   fmi = (ts_lua_fetch_multi_info *)ai->data;
516   fi  = (ts_lua_fetch_info *)edata;
517 
518   L      = ai->cinfo->routine.lua;
519   lmutex = ai->cinfo->routine.mctx->mutexp;
520 
521   fmi->done++;
522 
523   if (fi->fmi != fmi && fmi->done != fmi->total) {
524     return 0;
525   }
526 
527   // all finish
528   TSMutexLock(lmutex);
529 
530   if (fmi->total == 1 && !fmi->multi) {
531     ts_lua_fill_one_result(L, fi);
532     TSContCall(ci->contp, TS_LUA_EVENT_COROUTINE_CONT, (void *)1);
533 
534   } else {
535     lua_newtable(L);
536 
537     for (i = 1; i <= fmi->total; i++) {
538       ts_lua_fill_one_result(L, &fmi->fiv[i - 1]);
539       lua_rawseti(L, -2, i);
540     }
541 
542     TSContCall(ci->contp, TS_LUA_EVENT_COROUTINE_CONT, (void *)1);
543   }
544 
545   TSMutexUnlock(lmutex);
546   return 0;
547 }
548 
549 static inline void
ts_lua_destroy_fetch_multi_info(ts_lua_fetch_multi_info * fmi)550 ts_lua_destroy_fetch_multi_info(ts_lua_fetch_multi_info *fmi)
551 {
552   int i;
553   ts_lua_fetch_info *fi;
554 
555   if (fmi == NULL) {
556     return;
557   }
558 
559   for (i = 0; i < fmi->total; i++) {
560     fi = &fmi->fiv[i];
561 
562     if (fi->reader) {
563       TSIOBufferReaderFree(fi->reader);
564     }
565 
566     if (fi->buffer) {
567       TSIOBufferDestroy(fi->buffer);
568     }
569 
570     if (fi->fch) {
571       TSFetchDestroy(fi->fch);
572     }
573 
574     if (fi->contp) {
575       TSContDestroy(fi->contp);
576     }
577   }
578 
579   TSfree(fmi);
580 }
581 
582 static int
ts_lua_fetch_multi_cleanup(ts_lua_async_item * ai)583 ts_lua_fetch_multi_cleanup(ts_lua_async_item *ai)
584 {
585   if (ai->deleted) {
586     return 0;
587   }
588 
589   if (ai->data) {
590     ai->data = NULL;
591     TSContDestroy(ai->contp);
592     ai->contp = NULL;
593   }
594 
595   ai->deleted = 1;
596 
597   return 0;
598 }
599