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 <string.h>
20 #include <errno.h>
21 #include <pthread.h>
22 #include <getopt.h>
23 #include <inttypes.h>
24 #include <pthread.h>
25 
26 #include "ts_lua_util.h"
27 
28 #define TS_LUA_MAX_STATE_COUNT 256
29 
30 #define TS_LUA_STATS_TIMEOUT 5000   // 5s -- convert to configurable
31 #define TS_LUA_STATS_BUFFER_SIZE 10 // stats buffer
32 
33 #define TS_LUA_IND_STATE 0
34 #define TS_LUA_IND_GC_BYTES 1
35 #define TS_LUA_IND_THREADS 2
36 #define TS_LUA_IND_SIZE 3
37 
38 static uint64_t ts_lua_http_next_id   = 0;
39 static uint64_t ts_lua_g_http_next_id = 0;
40 
41 static ts_lua_main_ctx *ts_lua_main_ctx_array   = NULL;
42 static ts_lua_main_ctx *ts_lua_g_main_ctx_array = NULL;
43 
44 static pthread_key_t lua_g_state_key;
45 static pthread_key_t lua_state_key;
46 
47 // records.config entry injected by plugin
48 static char const *const ts_lua_mgmt_state_str   = "proxy.config.plugin.lua.max_states";
49 static char const *const ts_lua_mgmt_state_regex = "^[1-9][0-9]*$";
50 
51 // this is set the first time global configuration is probed.
52 static int ts_lua_max_state_count = 0;
53 
54 // lifecycle message tag
55 static char const *const print_tag = "stats_print";
56 static char const *const reset_tag = "stats_reset";
57 
58 // stat record strings
59 static char const *const ts_lua_stat_strs[] = {
60   "plugin.lua.remap.states",
61   "plugin.lua.remap.gc_bytes",
62   "plugin.lua.remap.threads",
63   NULL,
64 };
65 static char const *const ts_lua_g_stat_strs[] = {
66   "plugin.lua.global.states",
67   "plugin.lua.global.gc_bytes",
68   "plugin.lua.global.threads",
69   NULL,
70 };
71 
72 typedef struct {
73   ts_lua_main_ctx *main_ctx_array;
74 
75   int gc_kb;   // last collected gc in kb
76   int threads; // last collected number active threads
77 
78   int stat_inds[TS_LUA_IND_SIZE]; // stats indices
79 
80 } ts_lua_plugin_stats;
81 
82 ts_lua_plugin_stats *
create_plugin_stats(ts_lua_main_ctx * const main_ctx_array,char const * const * stat_strs)83 create_plugin_stats(ts_lua_main_ctx *const main_ctx_array, char const *const *stat_strs)
84 {
85   ts_lua_plugin_stats *const stats = TSmalloc(sizeof(ts_lua_plugin_stats));
86   memset(stats, 0, sizeof(ts_lua_plugin_stats));
87 
88   stats->main_ctx_array = main_ctx_array;
89 
90   // sample buffers
91   stats->gc_kb   = 0;
92   stats->threads = 0;
93 
94   int const max_state_count = ts_lua_max_state_count;
95 
96   for (int ind = 0; ind < TS_LUA_IND_SIZE; ++ind) {
97     stats->stat_inds[ind] = TSStatCreate(stat_strs[ind], TS_RECORDDATATYPE_INT, TS_STAT_NON_PERSISTENT, TS_STAT_SYNC_SUM);
98   }
99 
100   // initialize the number of states stat
101   int const sid = stats->stat_inds[TS_LUA_IND_STATE];
102   if (TS_ERROR != sid) {
103     TSStatIntSet(sid, max_state_count);
104   }
105 
106   return stats;
107 }
108 
109 ts_lua_main_ctx *
create_lua_vms()110 create_lua_vms()
111 {
112   ts_lua_main_ctx *ctx_array = NULL;
113 
114   // Inject the setting into records.config
115   static bool ts_mgt_int_inserted = false;
116   if (!ts_mgt_int_inserted) {
117     if (TS_SUCCESS == TSMgmtIntCreate(TS_RECORDTYPE_CONFIG, ts_lua_mgmt_state_str, TS_LUA_MAX_STATE_COUNT,
118                                       TS_RECORDUPDATE_RESTART_TS, TS_RECORDCHECK_INT, ts_lua_mgmt_state_regex,
119                                       TS_RECORDACCESS_READ_ONLY)) {
120       TSDebug(TS_LUA_DEBUG_TAG, "[%s] registered config string %s: with default [%d]", __FUNCTION__, ts_lua_mgmt_state_str,
121               TS_LUA_MAX_STATE_COUNT);
122     } else {
123       TSError("[%s][%s] failed to register %s", TS_LUA_DEBUG_TAG, __FUNCTION__, ts_lua_mgmt_state_str);
124     }
125     ts_mgt_int_inserted = true;
126   }
127 
128   if (0 == ts_lua_max_state_count) {
129     TSMgmtInt mgmt_state = 0;
130 
131     if (TS_SUCCESS != TSMgmtIntGet(ts_lua_mgmt_state_str, &mgmt_state)) {
132       TSDebug(TS_LUA_DEBUG_TAG, "[%s] setting max state to default: %d", __FUNCTION__, TS_LUA_MAX_STATE_COUNT);
133       ts_lua_max_state_count = TS_LUA_MAX_STATE_COUNT;
134     } else {
135       ts_lua_max_state_count = (int)mgmt_state;
136       TSDebug(TS_LUA_DEBUG_TAG, "[%s] found %s: [%d]", __FUNCTION__, ts_lua_mgmt_state_str, ts_lua_max_state_count);
137     }
138 
139     if (ts_lua_max_state_count < 1) {
140       TSError("[ts_lua][%s] invalid %s: %d", __FUNCTION__, ts_lua_mgmt_state_str, ts_lua_max_state_count);
141       ts_lua_max_state_count = 0;
142       return NULL;
143     }
144   }
145 
146   ctx_array = TSmalloc(sizeof(ts_lua_main_ctx) * ts_lua_max_state_count);
147   memset(ctx_array, 0, sizeof(ts_lua_main_ctx) * ts_lua_max_state_count);
148 
149   int const ret = ts_lua_create_vm(ctx_array, ts_lua_max_state_count);
150 
151   if (ret) {
152     ts_lua_destroy_vm(ctx_array, ts_lua_max_state_count);
153     TSfree(ctx_array);
154     ctx_array = NULL;
155     return NULL;
156   }
157 
158   // Initialize the GC numbers, no need to lock here
159   for (int index = 0; index < ts_lua_max_state_count; ++index) {
160     ts_lua_main_ctx *const main_ctx = (ctx_array + index);
161     lua_State *const lstate         = main_ctx->lua;
162     ts_lua_ctx_stats *const stats   = main_ctx->stats;
163 
164     stats->gc_kb = stats->gc_kb_max = lua_getgccount(lstate);
165   }
166 
167   return ctx_array;
168 }
169 
170 // dump exhaustive per state summary stats
171 static void
collectStats(ts_lua_plugin_stats * const plugin_stats)172 collectStats(ts_lua_plugin_stats *const plugin_stats)
173 {
174   TSMgmtInt gc_kb_total   = 0;
175   TSMgmtInt threads_total = 0;
176 
177   ts_lua_main_ctx *const main_ctx_array = plugin_stats->main_ctx_array;
178 
179   // aggregate stats on the states
180   for (int index = 0; index < ts_lua_max_state_count; ++index) {
181     ts_lua_main_ctx *const main_ctx = (main_ctx_array + index);
182     if (NULL != main_ctx) {
183       ts_lua_ctx_stats *const stats = main_ctx->stats;
184 
185       TSMutexLock(stats->mutexp);
186       gc_kb_total += (TSMgmtInt)stats->gc_kb;
187       threads_total += (TSMgmtInt)stats->threads;
188       TSMutexUnlock(stats->mutexp);
189     }
190   }
191 
192   // set the stats sample slot
193   plugin_stats->gc_kb   = gc_kb_total;
194   plugin_stats->threads = threads_total;
195 }
196 
197 static void
publishStats(ts_lua_plugin_stats * const plugin_stats)198 publishStats(ts_lua_plugin_stats *const plugin_stats)
199 {
200   TSMgmtInt const gc_bytes = plugin_stats->gc_kb * 1024;
201   TSStatIntSet(plugin_stats->stat_inds[TS_LUA_IND_GC_BYTES], gc_bytes);
202   TSStatIntSet(plugin_stats->stat_inds[TS_LUA_IND_THREADS], plugin_stats->threads);
203 }
204 
205 // dump exhaustive per state summary stats
206 static int
statsHandler(TSCont contp,TSEvent event,void * edata)207 statsHandler(TSCont contp, TSEvent event, void *edata)
208 {
209   ts_lua_plugin_stats *const plugin_stats = (ts_lua_plugin_stats *)TSContDataGet(contp);
210 
211   collectStats(plugin_stats);
212   publishStats(plugin_stats);
213 
214   TSContScheduleOnPool(contp, TS_LUA_STATS_TIMEOUT, TS_THREAD_POOL_TASK);
215 
216   return TS_EVENT_NONE;
217 }
218 
219 static void
get_time_now_str(char * const buf,size_t const buflen)220 get_time_now_str(char *const buf, size_t const buflen)
221 {
222   TSHRTime const timenowusec = TShrtime();
223   int64_t const timemsec     = (int64_t)(timenowusec / 1000000);
224   time_t const timesec       = (time_t)(timemsec / 1000);
225   int const ms               = (int)(timemsec % 1000);
226 
227   struct tm tm;
228   gmtime_r(&timesec, &tm);
229   size_t const dtlen = strftime(buf, buflen, "%b %e %H:%M:%S", &tm);
230 
231   // tack on the ms
232   snprintf(buf + dtlen, buflen - dtlen, ".%03d", ms);
233 }
234 
235 // dump exhaustive per state summary stats
236 static int
lifecycleHandler(TSCont contp,TSEvent event,void * edata)237 lifecycleHandler(TSCont contp, TSEvent event, void *edata)
238 {
239   // ensure the message is for ts_lua
240   TSPluginMsg *const msgp = (TSPluginMsg *)edata;
241   if (0 != strncasecmp(msgp->tag, TS_LUA_DEBUG_TAG, strlen(msgp->tag))) {
242     return TS_EVENT_NONE;
243   }
244 
245   ts_lua_main_ctx *const main_ctx_array = (ts_lua_main_ctx *)TSContDataGet(contp);
246 
247   static char const *const remapstr  = "remap";
248   static char const *const globalstr = "global";
249 
250   char const *labelstr = NULL;
251 
252   if (main_ctx_array == ts_lua_main_ctx_array) {
253     labelstr = remapstr;
254   } else {
255     labelstr = globalstr;
256   }
257 
258   char timebuf[128];
259   get_time_now_str(timebuf, 128);
260 
261   char const *const msgstr = (char *)msgp->data;
262   enum State { Print, Reset } state;
263   state                      = Print;
264   size_t const reset_tag_len = strlen(reset_tag);
265 
266   if (reset_tag_len <= msgp->data_size && 0 == strncasecmp(reset_tag, msgstr, reset_tag_len)) {
267     TSDebug(TS_LUA_DEBUG_TAG, "[%s] LIFECYCLE_MSG: %s", __FUNCTION__, reset_tag);
268     state = Reset;
269     fprintf(stderr, "[%s] %s (%s) resetting per state gc_kb_max and threads_max\n", timebuf, TS_LUA_DEBUG_TAG, labelstr);
270   } else {
271     TSDebug(TS_LUA_DEBUG_TAG, "[%s] LIFECYCLE_MSG: %s", __FUNCTION__, print_tag);
272   }
273 
274   for (int index = 0; index < ts_lua_max_state_count; ++index) {
275     ts_lua_main_ctx *const main_ctx = (main_ctx_array + index);
276     if (NULL != main_ctx) {
277       ts_lua_ctx_stats *const stats = main_ctx->stats;
278       if (NULL != main_ctx) {
279         TSMutexLock(stats->mutexp);
280 
281         switch (state) {
282         case Reset:
283           stats->threads_max = stats->threads;
284           stats->gc_kb_max   = stats->gc_kb;
285           break;
286 
287         case Print:
288         default:
289           fprintf(stderr, "[%s] %s (%s) id: %3d gc_kb: %6d gc_kb_max: %6d threads: %4d threads_max: %4d\n", timebuf,
290                   TS_LUA_DEBUG_TAG, labelstr, index, stats->gc_kb, stats->gc_kb_max, stats->threads, stats->threads_max);
291           break;
292         }
293 
294         TSMutexUnlock(stats->mutexp);
295       }
296     }
297   }
298 
299   return TS_EVENT_NONE;
300 }
301 
302 TSReturnCode
TSRemapInit(TSRemapInterface * api_info,char * errbuf,int errbuf_size)303 TSRemapInit(TSRemapInterface *api_info, char *errbuf, int errbuf_size)
304 {
305   if (!api_info || api_info->size < sizeof(TSRemapInterface)) {
306     strncpy(errbuf, "[TSRemapInit] - Incorrect size of TSRemapInterface structure", errbuf_size - 1);
307     errbuf[errbuf_size - 1] = '\0';
308     return TS_ERROR;
309   }
310 
311   if (NULL == ts_lua_main_ctx_array) {
312     ts_lua_main_ctx_array = create_lua_vms();
313     if (NULL != ts_lua_main_ctx_array) {
314       pthread_key_create(&lua_state_key, NULL);
315 
316       TSCont const lcontp = TSContCreate(lifecycleHandler, TSMutexCreate());
317       TSContDataSet(lcontp, ts_lua_main_ctx_array);
318       TSLifecycleHookAdd(TS_LIFECYCLE_MSG_HOOK, lcontp);
319 
320       ts_lua_plugin_stats *const plugin_stats = create_plugin_stats(ts_lua_main_ctx_array, ts_lua_stat_strs);
321 
322       // start the stats management
323       if (NULL != plugin_stats) {
324         TSDebug(TS_LUA_DEBUG_TAG, "Starting up stats management continuation");
325         TSCont const scontp = TSContCreate(statsHandler, TSMutexCreate());
326         TSContDataSet(scontp, plugin_stats);
327         TSContScheduleOnPool(scontp, TS_LUA_STATS_TIMEOUT, TS_THREAD_POOL_TASK);
328       }
329     } else {
330       return TS_ERROR;
331     }
332   }
333 
334   return TS_SUCCESS;
335 }
336 
337 TSReturnCode
TSRemapNewInstance(int argc,char * argv[],void ** ih,char * errbuf,int errbuf_size)338 TSRemapNewInstance(int argc, char *argv[], void **ih, char *errbuf, int errbuf_size)
339 {
340   int ret;
341   char script[TS_LUA_MAX_SCRIPT_FNAME_LENGTH];
342   char *inline_script                  = "";
343   int fn                               = 0;
344   int states                           = ts_lua_max_state_count;
345   static const struct option longopt[] = {
346     {"states", required_argument, 0, 's'},
347     {"inline", required_argument, 0, 'i'},
348     {0, 0, 0, 0},
349   };
350 
351   argc--;
352   argv++;
353 
354   for (;;) {
355     int opt;
356 
357     opt = getopt_long(argc, (char *const *)argv, "", longopt, NULL);
358     switch (opt) {
359     case 's':
360       states = atoi(optarg);
361       TSDebug(TS_LUA_DEBUG_TAG, "[%s] setting number of lua VMs [%d]", __FUNCTION__, states);
362       // set state
363       break;
364     case 'i':
365       inline_script = optarg;
366     }
367 
368     if (opt == -1) {
369       break;
370     }
371   }
372 
373   if (states < 1 || ts_lua_max_state_count < states) {
374     snprintf(errbuf, errbuf_size, "[TSRemapNewInstance] - invalid state in option input. Must be between 1 and %d",
375              ts_lua_max_state_count);
376     return TS_ERROR;
377   }
378 
379   if (argc - optind > 0) {
380     fn = 1;
381     if (argv[optind][0] == '/') {
382       snprintf(script, sizeof(script), "%s", argv[optind]);
383     } else {
384       snprintf(script, sizeof(script), "%s/%s", TSConfigDirGet(), argv[optind]);
385     }
386   }
387 
388   if (strlen(inline_script) == 0 && argc - optind < 1) {
389     strncpy(errbuf, "[TSRemapNewInstance] - lua script file or string is required !!", errbuf_size - 1);
390     errbuf[errbuf_size - 1] = '\0';
391     return TS_ERROR;
392   }
393 
394   if (strlen(script) >= TS_LUA_MAX_SCRIPT_FNAME_LENGTH - 16) {
395     strncpy(errbuf, "[TSRemapNewInstance] - lua script file name too long !!", errbuf_size - 1);
396     errbuf[errbuf_size - 1] = '\0';
397     return TS_ERROR;
398   }
399 
400   ts_lua_instance_conf *conf = NULL;
401 
402   // check to make sure it is a lua file and there is no parameter for the lua file
403   if (fn && (argc - optind < 2)) {
404     TSDebug(TS_LUA_DEBUG_TAG, "[%s] checking if script has been registered", __FUNCTION__);
405 
406     // we only need to check the first lua VM for script registration
407     TSMutexLock(ts_lua_main_ctx_array[0].mutexp);
408     conf = ts_lua_script_registered(ts_lua_main_ctx_array[0].lua, script);
409     TSMutexUnlock(ts_lua_main_ctx_array[0].mutexp);
410   }
411 
412   if (!conf) {
413     TSDebug(TS_LUA_DEBUG_TAG, "[%s] creating new conf instance", __FUNCTION__);
414 
415     conf = TSmalloc(sizeof(ts_lua_instance_conf));
416     if (!conf) {
417       strncpy(errbuf, "[TSRemapNewInstance] TSmalloc failed!!", errbuf_size - 1);
418       errbuf[errbuf_size - 1] = '\0';
419       return TS_ERROR;
420     }
421 
422     memset(conf, 0, sizeof(ts_lua_instance_conf));
423     conf->states    = states;
424     conf->remap     = 1;
425     conf->init_func = 0;
426 
427     if (fn) {
428       snprintf(conf->script, TS_LUA_MAX_SCRIPT_FNAME_LENGTH, "%s", script);
429     } else {
430       conf->content = inline_script;
431     }
432 
433     ts_lua_init_instance(conf);
434 
435     ret = ts_lua_add_module(conf, ts_lua_main_ctx_array, conf->states, argc - optind, &argv[optind], errbuf, errbuf_size);
436 
437     if (ret != 0) {
438       return TS_ERROR;
439     }
440 
441     // register the script only if it is from a file and has no __init__ function
442     if (fn && !conf->init_func) {
443       // we only need to register the script for the first lua VM
444       TSMutexLock(ts_lua_main_ctx_array[0].mutexp);
445       ts_lua_script_register(ts_lua_main_ctx_array[0].lua, conf->script, conf);
446       TSMutexUnlock(ts_lua_main_ctx_array[0].mutexp);
447     }
448   }
449 
450   *ih = conf;
451 
452   return TS_SUCCESS;
453 }
454 
455 void
TSRemapDeleteInstance(void * ih)456 TSRemapDeleteInstance(void *ih)
457 {
458   int states = ((ts_lua_instance_conf *)ih)->states;
459   ts_lua_del_module((ts_lua_instance_conf *)ih, ts_lua_main_ctx_array, states);
460   ts_lua_del_instance(ih);
461   // because we now reuse ts_lua_instance_conf / ih for remap rules sharing the same lua script
462   // we cannot safely free it in this function during the configuration reloads
463   // we therefore are leaking memory on configuration reloads
464   return;
465 }
466 
467 static TSRemapStatus
ts_lua_remap_plugin_init(void * ih,TSHttpTxn rh,TSRemapRequestInfo * rri)468 ts_lua_remap_plugin_init(void *ih, TSHttpTxn rh, TSRemapRequestInfo *rri)
469 {
470   int ret;
471   uint64_t req_id;
472 
473   TSCont contp;
474   lua_State *L;
475 
476   ts_lua_main_ctx *main_ctx;
477   ts_lua_http_ctx *http_ctx;
478   ts_lua_cont_info *ci;
479 
480   ts_lua_instance_conf *instance_conf;
481 
482   int remap     = (rri == NULL ? 0 : 1);
483   instance_conf = (ts_lua_instance_conf *)ih;
484 
485   main_ctx = pthread_getspecific(lua_state_key);
486   if (main_ctx == NULL) {
487     req_id   = __sync_fetch_and_add(&ts_lua_http_next_id, 1);
488     main_ctx = &ts_lua_main_ctx_array[req_id % instance_conf->states];
489     pthread_setspecific(lua_state_key, main_ctx);
490   }
491 
492   TSMutexLock(main_ctx->mutexp);
493 
494   http_ctx = ts_lua_create_http_ctx(main_ctx, instance_conf);
495 
496   http_ctx->txnp     = rh;
497   http_ctx->has_hook = 0;
498   http_ctx->rri      = rri;
499   if (rri != NULL) {
500     http_ctx->client_request_bufp = rri->requestBufp;
501     http_ctx->client_request_hdrp = rri->requestHdrp;
502     http_ctx->client_request_url  = rri->requestUrl;
503   }
504 
505   ci = &http_ctx->cinfo;
506   L  = ci->routine.lua;
507 
508   contp = TSContCreate(ts_lua_http_cont_handler, NULL);
509   TSContDataSet(contp, http_ctx);
510 
511   ci->contp = contp;
512   ci->mutex = TSContMutexGet((TSCont)rh);
513 
514   lua_getglobal(L, (remap ? TS_LUA_FUNCTION_REMAP : TS_LUA_FUNCTION_OS_RESPONSE));
515   if (lua_type(L, -1) != LUA_TFUNCTION) {
516     lua_pop(L, 1);
517     ts_lua_destroy_http_ctx(http_ctx);
518     TSMutexUnlock(main_ctx->mutexp);
519     return TSREMAP_NO_REMAP;
520   }
521 
522   ts_lua_set_cont_info(L, NULL);
523   if (lua_pcall(L, 0, 1, 0) != 0) {
524     TSError("[ts_lua][%s] lua_pcall failed: %s", __FUNCTION__, lua_tostring(L, -1));
525     ret = TSREMAP_NO_REMAP;
526 
527   } else {
528     ret = lua_tointeger(L, -1);
529   }
530 
531   lua_pop(L, 1);
532 
533   if (http_ctx->has_hook) {
534     TSDebug(TS_LUA_DEBUG_TAG, "[%s] has txn hook -> adding txn close hook handler to release resources", __FUNCTION__);
535     TSHttpTxnHookAdd(rh, TS_HTTP_TXN_CLOSE_HOOK, contp);
536   } else {
537     TSDebug(TS_LUA_DEBUG_TAG, "[%s] no txn hook -> release resources now", __FUNCTION__);
538     ts_lua_destroy_http_ctx(http_ctx);
539   }
540 
541   TSMutexUnlock(main_ctx->mutexp);
542 
543   return ret;
544 }
545 
546 void
TSRemapOSResponse(void * ih,TSHttpTxn rh,int os_response_type)547 TSRemapOSResponse(void *ih, TSHttpTxn rh, int os_response_type)
548 {
549   TSDebug(TS_LUA_DEBUG_TAG, "[%s] os response function and type - %d", __FUNCTION__, os_response_type);
550   ts_lua_remap_plugin_init(ih, rh, NULL);
551 }
552 
553 TSRemapStatus
TSRemapDoRemap(void * ih,TSHttpTxn rh,TSRemapRequestInfo * rri)554 TSRemapDoRemap(void *ih, TSHttpTxn rh, TSRemapRequestInfo *rri)
555 {
556   TSDebug(TS_LUA_DEBUG_TAG, "[%s] remap function", __FUNCTION__);
557   return ts_lua_remap_plugin_init(ih, rh, rri);
558 }
559 
560 static int
configHandler(TSCont contp,TSEvent event ATS_UNUSED,void * edata ATS_UNUSED)561 configHandler(TSCont contp, TSEvent event ATS_UNUSED, void *edata ATS_UNUSED)
562 {
563   TSDebug(TS_LUA_DEBUG_TAG, "[%s] calling configuration handler", __FUNCTION__);
564   ts_lua_instance_conf *conf = (ts_lua_instance_conf *)TSContDataGet(contp);
565   ts_lua_reload_module(conf, ts_lua_g_main_ctx_array, conf->states);
566   return 0;
567 }
568 
569 static int
globalHookHandler(TSCont contp,TSEvent event ATS_UNUSED,void * edata)570 globalHookHandler(TSCont contp, TSEvent event ATS_UNUSED, void *edata)
571 {
572   TSHttpTxn txnp = (TSHttpTxn)edata;
573 
574   TSMBuffer bufp;
575   TSMLoc hdr_loc;
576   TSMLoc url_loc;
577 
578   int ret;
579   uint64_t req_id;
580   TSCont txn_contp;
581 
582   lua_State *l;
583 
584   ts_lua_main_ctx *main_ctx;
585   ts_lua_http_ctx *http_ctx;
586   ts_lua_cont_info *ci;
587 
588   ts_lua_instance_conf *conf = (ts_lua_instance_conf *)TSContDataGet(contp);
589 
590   main_ctx = pthread_getspecific(lua_g_state_key);
591   if (main_ctx == NULL) {
592     req_id = __sync_fetch_and_add(&ts_lua_g_http_next_id, 1);
593     TSDebug(TS_LUA_DEBUG_TAG, "[%s] req_id: %" PRId64, __FUNCTION__, req_id);
594     main_ctx = &ts_lua_g_main_ctx_array[req_id % conf->states];
595     pthread_setspecific(lua_g_state_key, main_ctx);
596   }
597 
598   TSMutexLock(main_ctx->mutexp);
599 
600   http_ctx           = ts_lua_create_http_ctx(main_ctx, conf);
601   http_ctx->txnp     = txnp;
602   http_ctx->rri      = NULL;
603   http_ctx->has_hook = 0;
604 
605   if (!http_ctx->client_request_bufp) {
606     if (TSHttpTxnClientReqGet(txnp, &bufp, &hdr_loc) == TS_SUCCESS) {
607       http_ctx->client_request_bufp = bufp;
608       http_ctx->client_request_hdrp = hdr_loc;
609 
610       if (TSHttpHdrUrlGet(bufp, hdr_loc, &url_loc) == TS_SUCCESS) {
611         http_ctx->client_request_url = url_loc;
612       }
613     }
614   }
615 
616   if (!http_ctx->client_request_hdrp) {
617     ts_lua_destroy_http_ctx(http_ctx);
618     TSMutexUnlock(main_ctx->mutexp);
619 
620     TSHttpTxnReenable(txnp, TS_EVENT_HTTP_CONTINUE);
621     return 0;
622   }
623 
624   txn_contp = TSContCreate(ts_lua_http_cont_handler, NULL);
625   TSContDataSet(txn_contp, http_ctx);
626 
627   ci        = &http_ctx->cinfo;
628   ci->contp = txn_contp;
629   ci->mutex = TSContMutexGet((TSCont)txnp);
630 
631   l = ci->routine.lua;
632 
633   switch (event) {
634   case TS_EVENT_HTTP_READ_REQUEST_HDR:
635     lua_getglobal(l, TS_LUA_FUNCTION_G_READ_REQUEST);
636     break;
637 
638   case TS_EVENT_HTTP_SEND_REQUEST_HDR:
639     lua_getglobal(l, TS_LUA_FUNCTION_G_SEND_REQUEST);
640     break;
641 
642   case TS_EVENT_HTTP_READ_RESPONSE_HDR:
643     lua_getglobal(l, TS_LUA_FUNCTION_G_READ_RESPONSE);
644     break;
645 
646   case TS_EVENT_HTTP_SEND_RESPONSE_HDR:
647     // client response can be changed within a transaction
648     // (e.g. due to the follow redirect feature). So, clearing the pointers
649     // to allow API(s) to fetch the pointers again when it re-enters the hook
650     if (http_ctx->client_response_hdrp != NULL) {
651       TSHandleMLocRelease(http_ctx->client_response_bufp, TS_NULL_MLOC, http_ctx->client_response_hdrp);
652       http_ctx->client_response_bufp = NULL;
653       http_ctx->client_response_hdrp = NULL;
654     }
655     lua_getglobal(l, TS_LUA_FUNCTION_G_SEND_RESPONSE);
656     break;
657 
658   case TS_EVENT_HTTP_CACHE_LOOKUP_COMPLETE:
659     lua_getglobal(l, TS_LUA_FUNCTION_G_CACHE_LOOKUP_COMPLETE);
660     break;
661 
662   case TS_EVENT_HTTP_TXN_START:
663     lua_getglobal(l, TS_LUA_FUNCTION_G_TXN_START);
664     break;
665 
666   case TS_EVENT_HTTP_PRE_REMAP:
667     lua_getglobal(l, TS_LUA_FUNCTION_G_PRE_REMAP);
668     break;
669 
670   case TS_EVENT_HTTP_POST_REMAP:
671     lua_getglobal(l, TS_LUA_FUNCTION_G_POST_REMAP);
672     break;
673 
674   case TS_EVENT_HTTP_OS_DNS:
675     lua_getglobal(l, TS_LUA_FUNCTION_G_OS_DNS);
676     break;
677 
678   case TS_EVENT_HTTP_READ_CACHE_HDR:
679     lua_getglobal(l, TS_LUA_FUNCTION_G_READ_CACHE);
680     break;
681 
682   case TS_EVENT_HTTP_TXN_CLOSE:
683     lua_getglobal(l, TS_LUA_FUNCTION_G_TXN_CLOSE);
684     break;
685 
686   default:
687     ts_lua_destroy_http_ctx(http_ctx);
688     TSMutexUnlock(main_ctx->mutexp);
689     TSHttpTxnReenable(txnp, TS_EVENT_HTTP_CONTINUE);
690     return 0;
691   }
692 
693   if (lua_type(l, -1) != LUA_TFUNCTION) {
694     lua_pop(l, 1);
695     ts_lua_destroy_http_ctx(http_ctx);
696     TSMutexUnlock(main_ctx->mutexp);
697 
698     TSHttpTxnReenable(txnp, TS_EVENT_HTTP_CONTINUE);
699     return 0;
700   }
701 
702   ts_lua_set_cont_info(l, NULL);
703 
704   if (lua_pcall(l, 0, 1, 0) != 0) {
705     TSError("[ts_lua][%s] lua_pcall failed: %s", __FUNCTION__, lua_tostring(l, -1));
706   }
707 
708   ret = lua_tointeger(l, -1);
709   lua_pop(l, 1);
710 
711   // client response can be changed within a transaction
712   // (e.g. due to the follow redirect feature). So, clearing the pointers
713   // to allow API(s) to fetch the pointers again when it re-enters the hook
714   if (http_ctx->client_response_hdrp != NULL) {
715     TSHandleMLocRelease(http_ctx->client_response_bufp, TS_NULL_MLOC, http_ctx->client_response_hdrp);
716     http_ctx->client_response_bufp = NULL;
717     http_ctx->client_response_hdrp = NULL;
718   }
719 
720   if (http_ctx->has_hook) {
721     // add a hook to release resources for context
722     TSDebug(TS_LUA_DEBUG_TAG, "[%s] has txn hook -> adding txn close hook handler to release resources", __FUNCTION__);
723     TSHttpTxnHookAdd(txnp, TS_HTTP_TXN_CLOSE_HOOK, txn_contp);
724   } else {
725     TSDebug(TS_LUA_DEBUG_TAG, "[%s] no txn hook -> release resources now", __FUNCTION__);
726     ts_lua_destroy_http_ctx(http_ctx);
727   }
728 
729   TSMutexUnlock(main_ctx->mutexp);
730 
731   if (ret) {
732     TSHttpTxnReenable(txnp, TS_EVENT_HTTP_ERROR);
733   } else {
734     TSHttpTxnReenable(txnp, TS_EVENT_HTTP_CONTINUE);
735   }
736 
737   return 0;
738 }
739 
740 void
TSPluginInit(int argc,const char * argv[])741 TSPluginInit(int argc, const char *argv[])
742 {
743   TSPluginRegistrationInfo info;
744 
745   info.plugin_name   = "ts_lua";
746   info.vendor_name   = "Apache Software Foundation";
747   info.support_email = "dev@trafficserver.apache.org";
748 
749   if (TSPluginRegister(&info) != TS_SUCCESS) {
750     TSError("[ts_lua][%s] Plugin registration failed", __FUNCTION__);
751   }
752 
753   if (NULL == ts_lua_g_main_ctx_array) {
754     ts_lua_g_main_ctx_array = create_lua_vms();
755     if (NULL != ts_lua_g_main_ctx_array) {
756       pthread_key_create(&lua_g_state_key, NULL);
757 
758       TSCont const contp = TSContCreate(lifecycleHandler, TSMutexCreate());
759       TSContDataSet(contp, ts_lua_g_main_ctx_array);
760       TSLifecycleHookAdd(TS_LIFECYCLE_MSG_HOOK, contp);
761 
762       ts_lua_plugin_stats *const plugin_stats = create_plugin_stats(ts_lua_g_main_ctx_array, ts_lua_g_stat_strs);
763 
764       if (NULL != plugin_stats) {
765         TSCont const scontp = TSContCreate(statsHandler, TSMutexCreate());
766         TSContDataSet(scontp, plugin_stats);
767         TSContScheduleOnPool(scontp, TS_LUA_STATS_TIMEOUT, TS_THREAD_POOL_TASK);
768       }
769     } else {
770       return;
771     }
772   }
773 
774   int states = ts_lua_max_state_count;
775 
776   int reload                           = 0;
777   static const struct option longopt[] = {
778     {"states", required_argument, 0, 's'},
779     {"enable-reload", no_argument, 0, 'r'},
780     {0, 0, 0, 0},
781   };
782 
783   for (;;) {
784     int opt;
785 
786     opt = getopt_long(argc, (char *const *)argv, "", longopt, NULL);
787     switch (opt) {
788     case 's':
789       states = atoi(optarg);
790       // set state
791       break;
792     case 'r':
793       reload = 1;
794       TSDebug(TS_LUA_DEBUG_TAG, "[%s] enable global plugin reload [%d]", __FUNCTION__, reload);
795       break;
796     }
797 
798     if (opt == -1) {
799       break;
800     }
801   }
802 
803   if (states < 1 || ts_lua_max_state_count < states) {
804     TSError("[ts_lua][%s] invalid # of states from option input. Must be between 1 and %d", __FUNCTION__, ts_lua_max_state_count);
805     return;
806   }
807 
808   if (argc - optind < 1) {
809     TSError("[ts_lua][%s] lua script file required !!", __FUNCTION__);
810     return;
811   }
812 
813   if (strlen(argv[optind]) >= TS_LUA_MAX_SCRIPT_FNAME_LENGTH - 16) {
814     TSError("[ts_lua][%s] lua script file name too long !!", __FUNCTION__);
815     return;
816   }
817 
818   ts_lua_instance_conf *conf = TSmalloc(sizeof(ts_lua_instance_conf));
819   if (!conf) {
820     TSError("[ts_lua][%s] TSmalloc failed !!", __FUNCTION__);
821     return;
822   }
823   memset(conf, 0, sizeof(ts_lua_instance_conf));
824   conf->remap  = 0;
825   conf->states = states;
826 
827   if (argv[optind][0] == '/') {
828     snprintf(conf->script, TS_LUA_MAX_SCRIPT_FNAME_LENGTH, "%s", argv[optind]);
829   } else {
830     snprintf(conf->script, TS_LUA_MAX_SCRIPT_FNAME_LENGTH, "%s/%s", TSConfigDirGet(), argv[optind]);
831   }
832 
833   ts_lua_init_instance(conf);
834 
835   char errbuf[TS_LUA_MAX_STR_LENGTH];
836   int const errbuf_len = sizeof(errbuf);
837   int const ret =
838     ts_lua_add_module(conf, ts_lua_g_main_ctx_array, conf->states, argc - optind, (char **)&argv[optind], errbuf, errbuf_len);
839 
840   if (ret != 0) {
841     TSError(errbuf, NULL);
842     TSError("[ts_lua][%s] ts_lua_add_module failed", __FUNCTION__);
843     return;
844   }
845 
846   TSCont global_contp = TSContCreate(globalHookHandler, NULL);
847   if (!global_contp) {
848     TSError("[ts_lua][%s] could not create transaction start continuation", __FUNCTION__);
849     return;
850   }
851   TSContDataSet(global_contp, conf);
852 
853   // adding hook based on whether the lua global function exists.
854   ts_lua_main_ctx *main_ctx = &ts_lua_g_main_ctx_array[0];
855   ts_lua_http_ctx *http_ctx = ts_lua_create_http_ctx(main_ctx, conf);
856   lua_State *l              = http_ctx->cinfo.routine.lua;
857 
858   lua_getglobal(l, TS_LUA_FUNCTION_G_SEND_REQUEST);
859   if (lua_type(l, -1) == LUA_TFUNCTION) {
860     TSHttpHookAdd(TS_HTTP_SEND_REQUEST_HDR_HOOK, global_contp);
861     TSDebug(TS_LUA_DEBUG_TAG, "send_request_hdr_hook added");
862   }
863   lua_pop(l, 1);
864 
865   lua_getglobal(l, TS_LUA_FUNCTION_G_READ_RESPONSE);
866   if (lua_type(l, -1) == LUA_TFUNCTION) {
867     TSHttpHookAdd(TS_HTTP_READ_RESPONSE_HDR_HOOK, global_contp);
868     TSDebug(TS_LUA_DEBUG_TAG, "read_response_hdr_hook added");
869   }
870   lua_pop(l, 1);
871 
872   lua_getglobal(l, TS_LUA_FUNCTION_G_SEND_RESPONSE);
873   if (lua_type(l, -1) == LUA_TFUNCTION) {
874     TSHttpHookAdd(TS_HTTP_SEND_RESPONSE_HDR_HOOK, global_contp);
875     TSDebug(TS_LUA_DEBUG_TAG, "send_response_hdr_hook added");
876   }
877   lua_pop(l, 1);
878 
879   lua_getglobal(l, TS_LUA_FUNCTION_G_CACHE_LOOKUP_COMPLETE);
880   if (lua_type(l, -1) == LUA_TFUNCTION) {
881     TSHttpHookAdd(TS_HTTP_CACHE_LOOKUP_COMPLETE_HOOK, global_contp);
882     TSDebug(TS_LUA_DEBUG_TAG, "cache_lookup_complete_hook added");
883   }
884   lua_pop(l, 1);
885 
886   lua_getglobal(l, TS_LUA_FUNCTION_G_READ_REQUEST);
887   if (lua_type(l, -1) == LUA_TFUNCTION) {
888     TSHttpHookAdd(TS_HTTP_READ_REQUEST_HDR_HOOK, global_contp);
889     TSDebug(TS_LUA_DEBUG_TAG, "read_request_hdr_hook added");
890   }
891   lua_pop(l, 1);
892 
893   lua_getglobal(l, TS_LUA_FUNCTION_G_TXN_START);
894   if (lua_type(l, -1) == LUA_TFUNCTION) {
895     TSHttpHookAdd(TS_HTTP_TXN_START_HOOK, global_contp);
896     TSDebug(TS_LUA_DEBUG_TAG, "txn_start_hook added");
897   }
898   lua_pop(l, 1);
899 
900   lua_getglobal(l, TS_LUA_FUNCTION_G_PRE_REMAP);
901   if (lua_type(l, -1) == LUA_TFUNCTION) {
902     TSHttpHookAdd(TS_HTTP_PRE_REMAP_HOOK, global_contp);
903     TSDebug(TS_LUA_DEBUG_TAG, "pre_remap_hook added");
904   }
905   lua_pop(l, 1);
906 
907   lua_getglobal(l, TS_LUA_FUNCTION_G_POST_REMAP);
908   if (lua_type(l, -1) == LUA_TFUNCTION) {
909     TSHttpHookAdd(TS_HTTP_POST_REMAP_HOOK, global_contp);
910     TSDebug(TS_LUA_DEBUG_TAG, "post_remap_hook added");
911   }
912   lua_pop(l, 1);
913 
914   lua_getglobal(l, TS_LUA_FUNCTION_G_OS_DNS);
915   if (lua_type(l, -1) == LUA_TFUNCTION) {
916     TSHttpHookAdd(TS_HTTP_OS_DNS_HOOK, global_contp);
917     TSDebug(TS_LUA_DEBUG_TAG, "os_dns_hook added");
918   }
919   lua_pop(l, 1);
920 
921   lua_getglobal(l, TS_LUA_FUNCTION_G_READ_CACHE);
922   if (lua_type(l, -1) == LUA_TFUNCTION) {
923     TSHttpHookAdd(TS_HTTP_READ_CACHE_HDR_HOOK, global_contp);
924     TSDebug(TS_LUA_DEBUG_TAG, "read_cache_hdr_hook added");
925   }
926   lua_pop(l, 1);
927 
928   lua_getglobal(l, TS_LUA_FUNCTION_G_TXN_CLOSE);
929   if (lua_type(l, -1) == LUA_TFUNCTION) {
930     TSHttpHookAdd(TS_HTTP_TXN_CLOSE_HOOK, global_contp);
931     TSDebug(TS_LUA_DEBUG_TAG, "txn_close_hook added");
932   }
933   lua_pop(l, 1);
934 
935   ts_lua_destroy_http_ctx(http_ctx);
936 
937   // support for reload as global plugin
938   if (reload) {
939     TSCont config_contp = TSContCreate(configHandler, NULL);
940     if (!config_contp) {
941       TSError("[ts_lua][%s] could not create configuration continuation", __FUNCTION__);
942       return;
943     }
944     TSContDataSet(config_contp, conf);
945 
946     TSMgmtUpdateRegister(config_contp, "ts_lua");
947   }
948 }
949