1 /** @file
2 
3   Proxy Side Include plugin (PSI)
4 
5   Synopsis:
6 
7   This plugin allows to insert the content of a file stored on the proxy disk
8   into the body of an html response.
9 
10   The plugin illustrates how to use a pool of threads in order to do blocking
11   calls (here, some disk i/o) in a Traffic Server plugin.
12 
13   Further details: Refer to README file.
14 
15   @section license License
16 
17   Licensed to the Apache Software Foundation (ASF) under one
18   or more contributor license agreements.  See the NOTICE file
19   distributed with this work for additional information
20   regarding copyright ownership.  The ASF licenses this file
21   to you under the Apache License, Version 2.0 (the
22   "License"); you may not use this file except in compliance
23   with the License.  You may obtain a copy of the License at
24 
25       http://www.apache.org/licenses/LICENSE-2.0
26 
27   Unless required by applicable law or agreed to in writing, software
28   distributed under the License is distributed on an "AS IS" BASIS,
29   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
30   See the License for the specific language governing permissions and
31   limitations under the License.
32  */
33 
34 #include <stdio.h>
35 #include <stdlib.h>
36 #include <limits.h>
37 #include <string.h>
38 #include <sys/param.h>
39 
40 #include "ts/ts.h"
41 #include "thread.h"
42 #include "tscore/ink_defs.h"
43 
44 /* This is the number of threads spawned by the plugin.
45    Should be tuned based on performance requirements,
46    blocking calls duration, etc... */
47 #define NB_THREADS 3
48 
49 #define PSI_FILENAME_MAX_SIZE 512
50 #define PSI_PATH_MAX_SIZE 256
51 #define PSI_PATH "include"
52 
53 #define PSI_START_TAG "<!--include="
54 #define PSI_START_TAG_LEN 12
55 
56 #define PSI_END_TAG "-->"
57 #define PSI_END_TAG_LEN 3
58 
59 #define MIME_FIELD_XPSI "X-Psi"
60 
61 typedef enum {
62   STATE_READ_DATA = 1,
63   STATE_READ_PSI  = 2,
64   STATE_DUMP_PSI  = 3,
65 } PluginState;
66 
67 typedef enum {
68   PARSE_SEARCH,
69   PARSE_EXTRACT,
70 } ParseState;
71 
72 typedef struct {
73   unsigned int magic;
74   TSVIO output_vio;
75   TSIOBuffer output_buffer;
76   TSIOBufferReader output_reader;
77 
78   TSIOBuffer psi_buffer;
79   TSIOBufferReader psi_reader;
80   char psi_filename[PSI_FILENAME_MAX_SIZE + 128];
81   int psi_filename_len;
82   int psi_success;
83 
84   ParseState parse_state;
85 
86   PluginState state;
87   int transform_bytes;
88 } ContData;
89 
90 typedef struct {
91   TSCont contp;
92   TSEvent event;
93 } TryLockData;
94 
95 typedef enum {
96   STR_SUCCESS,
97   STR_PARTIAL,
98   STR_FAIL,
99 } StrOperationResult;
100 
101 extern Queue job_queue;
102 
103 static TSTextLogObject log;
104 static char psi_directory[PSI_PATH_MAX_SIZE];
105 
106 static int trylock_handler(TSCont contp, TSEvent event, void *edata);
107 
108 /*-------------------------------------------------------------------------
109   cont_data_alloc
110   Allocate and initialize a ContData structure associated to a transaction
111 
112   Input:
113   Output:
114   Return Value:
115     Pointer on a new allocated ContData structure
116   -------------------------------------------------------------------------*/
117 static ContData *
cont_data_alloc()118 cont_data_alloc()
119 {
120   ContData *data;
121 
122   data                = (ContData *)TSmalloc(sizeof(ContData));
123   data->magic         = MAGIC_ALIVE;
124   data->output_vio    = NULL;
125   data->output_buffer = NULL;
126   data->output_reader = NULL;
127 
128   data->psi_buffer       = NULL;
129   data->psi_reader       = NULL;
130   data->psi_filename[0]  = '\0';
131   data->psi_filename_len = 0;
132   data->psi_success      = 0;
133 
134   data->parse_state = PARSE_SEARCH;
135 
136   data->state           = STATE_READ_DATA;
137   data->transform_bytes = 0;
138 
139   return data;
140 }
141 
142 /*-------------------------------------------------------------------------
143   cont_data_destroy
144   Deallocate ContData structure associated to a transaction
145 
146   Input:
147     data   structure to deallocate
148   Output:
149   Return Value:
150     none
151   -------------------------------------------------------------------------*/
152 static void
cont_data_destroy(ContData * data)153 cont_data_destroy(ContData *data)
154 {
155   TSDebug(PLUGIN_NAME, "Destroying continuation data");
156   if (data) {
157     TSAssert(data->magic == MAGIC_ALIVE);
158     if (data->output_reader) {
159       TSIOBufferReaderFree(data->output_reader);
160       data->output_reader = NULL;
161     }
162     if (data->output_buffer) {
163       TSIOBufferDestroy(data->output_buffer);
164       data->output_buffer = NULL;
165     }
166     if (data->psi_reader) {
167       TSIOBufferReaderFree(data->psi_reader);
168       data->psi_reader = NULL;
169     }
170     if (data->psi_buffer) {
171       TSIOBufferDestroy(data->psi_buffer);
172       data->psi_buffer = NULL;
173     }
174     data->magic = MAGIC_DEAD;
175     TSfree(data);
176   }
177 }
178 
179 /*-------------------------------------------------------------------------
180   strsearch_ioreader
181   Looks for string pattern in an iobuffer
182 
183   Input:
184     reader   reader on a iobuffer
185     pattern  string to look for (nul terminated)
186   Output:
187     nparse   number of chars scanned, excluding the matching pattern
188   Return Value:
189     STR_SUCCESS if pattern found
190     STR_PARTIAL if pattern found partially
191     STR_FAIL    if pattern not found
192   -------------------------------------------------------------------------*/
193 static StrOperationResult
strsearch_ioreader(TSIOBufferReader reader,const char * pattern,int * nparse)194 strsearch_ioreader(TSIOBufferReader reader, const char *pattern, int *nparse)
195 {
196   int index             = 0;
197   TSIOBufferBlock block = TSIOBufferReaderStart(reader);
198   int slen              = strlen(pattern);
199 
200   if (slen <= 0) {
201     return STR_FAIL;
202   }
203 
204   *nparse = 0;
205 
206   /* Loop thru each block while we've not yet found the pattern */
207   while ((block != NULL) && (index < slen)) {
208     int64_t blocklen;
209     const char *blockptr = TSIOBufferBlockReadStart(block, reader, &blocklen);
210     const char *ptr;
211 
212     for (ptr = blockptr; ptr < blockptr + blocklen; ptr++) {
213       (*nparse)++;
214       if (*ptr == pattern[index]) {
215         index++;
216         if (index == slen) {
217           break;
218         }
219       } else {
220         index = 0;
221       }
222     }
223 
224     /* Parse next block */
225     block = TSIOBufferBlockNext(block);
226   }
227 
228   *nparse -= index; /* Adjust nparse so it doesn't include matching chars */
229   if (index == slen) {
230     TSDebug(PLUGIN_NAME, "strfind: match for %s at position %d", pattern, *nparse);
231     return STR_SUCCESS;
232   } else if (index > 0) {
233     TSDebug(PLUGIN_NAME, "strfind: partial match for %s at position %d", pattern, *nparse);
234     return STR_PARTIAL;
235   } else {
236     TSDebug(PLUGIN_NAME, "strfind no match for %s", pattern);
237     return STR_FAIL;
238   }
239 }
240 
241 /*-------------------------------------------------------------------------
242   strextract_ioreader
243   Extract a string from an iobuffer.
244   Start reading at position offset in iobuffer and extract until the
245   string end_pattern is found.
246 
247   Input:
248     reader      reader on a iobuffer
249     offset      position to start reading
250     end_pattern the termination string (nul terminated)
251   Output:
252     buffer      if success, contains the extracted string, nul terminated
253     buflen      if success, contains the buffer length (excluding null char).
254   Return Value:
255     STR_SUCCESS if extraction successful
256     STR_PARTIAL if extraction not yet completed
257     STR_FAIL    if extraction failed
258   -------------------------------------------------------------------------*/
259 static int
strextract_ioreader(TSIOBufferReader reader,int offset,const char * end_pattern,char * buffer,int * buflen)260 strextract_ioreader(TSIOBufferReader reader, int offset, const char *end_pattern, char *buffer, int *buflen)
261 {
262   int buf_idx       = 0;
263   int p_idx         = 0;
264   int nbytes_so_far = 0;
265   int plen          = strlen(end_pattern);
266   const char *ptr;
267   TSIOBufferBlock block = TSIOBufferReaderStart(reader);
268 
269   if (plen <= 0) {
270     return STR_FAIL;
271   }
272 
273   /* Now start extraction */
274   while ((block != NULL) && (p_idx < plen) && (buf_idx < PSI_FILENAME_MAX_SIZE)) {
275     int64_t blocklen;
276     const char *blockptr = TSIOBufferBlockReadStart(block, reader, &blocklen);
277 
278     for (ptr = blockptr; ptr < blockptr + blocklen; ptr++, nbytes_so_far++) {
279       if (nbytes_so_far >= offset) {
280         /* Add a new character to the filename */
281         buffer[buf_idx++] = *ptr;
282 
283         /* If we have reach the end of the filename, we're done */
284         if (end_pattern[p_idx] == *ptr) {
285           p_idx++;
286           if (p_idx == plen) {
287             break;
288           }
289         } else {
290           p_idx = 0;
291         }
292 
293         /* The filename is too long, something is fishy... let's abort extraction */
294         if (buf_idx >= PSI_FILENAME_MAX_SIZE) {
295           break;
296         }
297       }
298     }
299 
300     block = TSIOBufferBlockNext(block);
301   }
302 
303   /* Error, could not read end of filename */
304   if (buf_idx >= PSI_FILENAME_MAX_SIZE) {
305     TSDebug(PLUGIN_NAME, "strextract: filename too long");
306     *buflen = 0;
307     return STR_FAIL;
308   }
309 
310   /* Full Match */
311   if (p_idx == plen) {
312     /* Nul terminate the filename, remove the end_pattern copied into the buffer */
313     *buflen         = buf_idx - plen;
314     buffer[*buflen] = '\0';
315     TSDebug(PLUGIN_NAME, "strextract: filename = |%s|", buffer);
316     return STR_SUCCESS;
317   }
318   /* End of filename not yet reached we need to read some more data */
319   else {
320     TSDebug(PLUGIN_NAME, "strextract: partially extracted filename");
321     *buflen = buf_idx - p_idx;
322     return STR_PARTIAL;
323   }
324 }
325 
326 /*-------------------------------------------------------------------------
327   parse_data
328   Search for psi filename in the data.
329 
330   Input:
331     contp   continuation for the current transaction
332     reader  reader on the iobuffer that contains data
333     avail   amount of data available in the iobuffer
334   Output:
335     towrite    amount of data in the iobuffer that can be written
336                to the downstream vconnection
337     toconsume  amount of data in the iobuffer to consume
338   Return Value:
339     0   if no psi filename found
340     1  if a psi filename was found
341   -------------------------------------------------------------------------*/
342 static int
parse_data(TSCont contp,TSIOBufferReader input_reader,int avail,int * toconsume,int * towrite)343 parse_data(TSCont contp, TSIOBufferReader input_reader, int avail, int *toconsume, int *towrite)
344 {
345   ContData *data;
346   int nparse = 0;
347   int status;
348 
349   data = TSContDataGet(contp);
350   TSAssert(data->magic == MAGIC_ALIVE);
351 
352   if (data->parse_state == PARSE_SEARCH) {
353     /* Search for the start pattern */
354     status = strsearch_ioreader(input_reader, PSI_START_TAG, &nparse);
355     switch (status) {
356     case STR_FAIL:
357       /* We didn't found the pattern */
358       *toconsume        = avail;
359       *towrite          = avail;
360       data->parse_state = PARSE_SEARCH;
361       return 0;
362     case STR_PARTIAL:
363       /* We need to read some more data */
364       *toconsume        = nparse;
365       *towrite          = nparse;
366       data->parse_state = PARSE_SEARCH;
367       return 0;
368     case STR_SUCCESS:
369       /* We found the start_pattern, let's go ahead */
370       data->psi_filename_len = 0;
371       data->psi_filename[0]  = '\0';
372       data->parse_state      = PARSE_EXTRACT;
373       break;
374     default:
375       TSAssert(!"strsearch_ioreader returned unexpected status");
376     }
377   }
378 
379   /* And now let's extract the filename */
380   status = strextract_ioreader(input_reader, nparse + PSI_START_TAG_LEN, PSI_END_TAG, data->psi_filename, &data->psi_filename_len);
381   switch (status) {
382   case STR_FAIL:
383     /* We couldn't extract a valid filename */
384     *toconsume        = nparse;
385     *towrite          = nparse;
386     data->parse_state = PARSE_SEARCH;
387     return 0;
388   case STR_PARTIAL:
389     /* We need to read some more data */
390     *toconsume        = nparse;
391     *towrite          = nparse;
392     data->parse_state = PARSE_EXTRACT;
393     return 0;
394   case STR_SUCCESS:
395     /* We got a valid filename */
396     *toconsume        = nparse + PSI_START_TAG_LEN + data->psi_filename_len + PSI_END_TAG_LEN;
397     *towrite          = nparse;
398     data->parse_state = PARSE_SEARCH;
399     return 1;
400   default:
401     TSAssert(!"strextract_ioreader returned bad status");
402   }
403 
404   return 0;
405 }
406 
407 // TODO: Use libc basename function
408 //
409 /*-------------------------------------------------------------------------
410   strip_path
411   Utility func to strip path from a filename (= _basename cmd on unix)
412   Input:
413     filename
414   Output :
415     None
416   Return Value:
417     Filename with path stripped
418   -------------------------------------------------------------------------*/
419 static const char *
_basename(const char * filename)420 _basename(const char *filename)
421 {
422   char *cptr;
423   const char *ptr = filename;
424 
425   while ((cptr = strchr(ptr, '/')) != NULL) {
426     ptr = cptr + 1;
427   }
428   return ptr;
429 }
430 
431 /*-------------------------------------------------------------------------
432   psi_include
433   Read file to include. Copy its content into an iobuffer.
434 
435   This is the function doing blocking calls and called by the plugin's threads
436 
437   Input:
438     data      continuation for the current transaction
439   Output :
440     data->psi_buffer  contains the file content
441     data->psi_success  0 if include failed, 1 if success
442   Return Value:
443     0  if failure
444     1  if success
445   -------------------------------------------------------------------------*/
446 static int
psi_include(TSCont contp,void * edata ATS_UNUSED)447 psi_include(TSCont contp, void *edata ATS_UNUSED)
448 {
449 #define BUFFER_SIZE 1024
450   ContData *data;
451   TSFile filep;
452   char inc_file[PSI_PATH_MAX_SIZE + PSI_FILENAME_MAX_SIZE];
453 
454   /* We manipulate plugin continuation data from a separate thread.
455      Grab mutex to avoid concurrent access */
456   TSMutexLock(TSContMutexGet(contp));
457   data = TSContDataGet(contp);
458   TSAssert(data->magic == MAGIC_ALIVE);
459 
460   if (!data->psi_buffer) {
461     data->psi_buffer = TSIOBufferCreate();
462     data->psi_reader = TSIOBufferReaderAlloc(data->psi_buffer);
463   }
464 
465   /* For security reason, we do not allow to include files that are
466      not in the directory <plugin_path>/include.
467      Also include file cannot contain any path. */
468   snprintf(inc_file, sizeof(inc_file), "%s/%s", psi_directory, _basename(data->psi_filename));
469 
470   /* Read the include file and copy content into iobuffer */
471   if ((filep = TSfopen(inc_file, "r")) != NULL) {
472     TSDebug(PLUGIN_NAME, "Reading include file %s", inc_file);
473 
474     char buf[BUFFER_SIZE];
475     while (TSfgets(filep, buf, BUFFER_SIZE) != NULL) {
476       int64_t len, ndone, ntodo;
477 
478       len   = strlen(buf);
479       ndone = 0;
480       ntodo = len;
481       while (ntodo > 0) {
482         /* TSIOBufferStart allocates more blocks if required */
483         TSIOBufferBlock block = TSIOBufferStart(data->psi_buffer);
484         int64_t avail;
485         char *ptr_block = TSIOBufferBlockWriteStart(block, &avail);
486         int64_t towrite = MIN(ntodo, avail);
487 
488         memcpy(ptr_block, buf + ndone, towrite);
489         TSIOBufferProduce(data->psi_buffer, towrite);
490         ntodo -= towrite;
491         ndone += towrite;
492       }
493     }
494     TSfclose(filep);
495     data->psi_success = 1;
496     if (log) {
497       TSTextLogObjectWrite(log, "Successfully included file: %s", inc_file);
498     }
499   } else {
500     data->psi_success = 0;
501     if (log) {
502       TSTextLogObjectWrite(log, "Failed to include file: %s", inc_file);
503     }
504   }
505 
506   /* Change state and schedule an event EVENT_IMMEDIATE on the plugin continuation
507      to let it know we're done. */
508 
509   /* Note: if the blocking call was not in the transformation state (i.e. in
510      TS_HTTP_READ_REQUEST_HDR, TS_HTTP_OS_DNS and so on...) we could
511      use TSHttpTxnReenable to wake up the transaction instead of sending an event. */
512 
513   TSContScheduleOnPool(contp, 0, TS_THREAD_POOL_NET);
514   data->psi_success = 0;
515   data->state       = STATE_READ_DATA;
516   TSMutexUnlock(TSContMutexGet(contp));
517 
518   return 0;
519 }
520 
521 /*-------------------------------------------------------------------------
522   wake_up_streams
523   Send an event to the upstream vconnection to either
524     - ask for more data
525     - let it know we're done
526   Reenable the downstream vconnection
527   Input:
528     contp      continuation for the current transaction
529   Output :
530   Return Value:
531    0 if failure
532    1 if success
533   -------------------------------------------------------------------------*/
534 static int
wake_up_streams(TSCont contp)535 wake_up_streams(TSCont contp)
536 {
537   TSVIO input_vio;
538   ContData *data;
539   int ntodo;
540 
541   data = TSContDataGet(contp);
542   TSAssert(data->magic == MAGIC_ALIVE);
543 
544   input_vio = TSVConnWriteVIOGet(contp);
545   ntodo     = TSVIONTodoGet(input_vio);
546 
547   if (ntodo > 0) {
548     TSVIOReenable(data->output_vio);
549     TSContCall(TSVIOContGet(input_vio), TS_EVENT_VCONN_WRITE_READY, input_vio);
550   } else {
551     TSDebug(PLUGIN_NAME, "Total bytes produced by transform = %d", data->transform_bytes);
552     TSVIONBytesSet(data->output_vio, data->transform_bytes);
553     TSVIOReenable(data->output_vio);
554     TSContCall(TSVIOContGet(input_vio), TS_EVENT_VCONN_WRITE_COMPLETE, input_vio);
555   }
556 
557   return 1;
558 }
559 
560 /*-------------------------------------------------------------------------
561   handle_transform
562    Get data from upstream vconn.
563    Parse it.
564    Include file if include tags found.
565    Copy data to downstream vconn.
566    Wake up upstream to get more data.
567 
568   Input:
569     contp      continuation for the current transaction
570   Output :
571   Return Value:
572    0 if failure
573    1 if success
574   -------------------------------------------------------------------------*/
575 static int
handle_transform(TSCont contp)576 handle_transform(TSCont contp)
577 {
578   TSVConn output_conn;
579   TSVIO input_vio;
580   ContData *data;
581   int toread, toconsume = 0, towrite = 0;
582 
583   /* Get the output (downstream) vconnection where we'll write data to. */
584   output_conn = TSTransformOutputVConnGet(contp);
585 
586   /* Get upstream vio */
587   input_vio = TSVConnWriteVIOGet(contp);
588   data      = TSContDataGet(contp);
589   TSAssert(data->magic == MAGIC_ALIVE);
590 
591   if (!data->output_buffer) {
592     data->output_buffer = TSIOBufferCreate();
593     data->output_reader = TSIOBufferReaderAlloc(data->output_buffer);
594 
595     /* INT64_MAX because we don't know yet how much bytes we'll produce */
596     data->output_vio = TSVConnWrite(output_conn, contp, data->output_reader, INT64_MAX);
597   }
598 
599   /* If the input VIO's buffer is NULL, the transformation is over */
600   if (!TSVIOBufferGet(input_vio)) {
601     TSDebug(PLUGIN_NAME, "input_vio NULL, terminating transformation");
602     TSVIONBytesSet(data->output_vio, data->transform_bytes);
603     TSVIOReenable(data->output_vio);
604     return 1;
605   }
606 
607   /* Determine how much data we have left to read. */
608   toread = TSVIONTodoGet(input_vio);
609 
610   if (toread > 0) {
611     TSIOBufferReader input_reader = TSVIOReaderGet(input_vio);
612     int avail                     = TSIOBufferReaderAvail(input_reader);
613 
614     /* There are some data available for reading. Let's parse it */
615     if (avail > 0) {
616       int psi;
617       /* No need to parse data if there are too few bytes left to contain
618          an include command... */
619       if (toread > (PSI_START_TAG_LEN + PSI_END_TAG_LEN)) {
620         psi = parse_data(contp, input_reader, avail, &toconsume, &towrite);
621       } else {
622         towrite   = avail;
623         toconsume = avail;
624         psi       = 0;
625       }
626 
627       if (towrite > 0) {
628         /* Update the total size of the doc so far */
629         data->transform_bytes += towrite;
630 
631         /* Copy the data from the read buffer to the output buffer. */
632         /* TODO: Should we check the return value of TSIOBufferCopy() ? */
633         TSIOBufferCopy(TSVIOBufferGet(data->output_vio), TSVIOReaderGet(input_vio), towrite, 0);
634         /* Reenable the output connection so it can read the data we've produced. */
635         TSVIOReenable(data->output_vio);
636       }
637 
638       if (toconsume > 0) {
639         /* Consume data we've processed an we are no longer interested in */
640         TSIOBufferReaderConsume(input_reader, toconsume);
641 
642         /* Modify the input VIO to reflect how much data we've completed. */
643         TSVIONDoneSet(input_vio, TSVIONDoneGet(input_vio) + toconsume);
644       }
645 
646       /* Did we find a psi filename to execute in the data ? */
647       if (psi) {
648         Job *new_job;
649         /* Add a request to include a file into the jobs queue.. */
650         /* We'll be called back once it's done with an EVENT_IMMEDIATE */
651         TSDebug(PLUGIN_NAME, "Psi filename extracted, adding an include job to thread queue");
652         data->state = STATE_READ_PSI;
653 
654         /* Create a new job request and add it to the queue */
655         new_job = job_create(contp, &psi_include, NULL);
656         add_to_queue(&job_queue, new_job);
657 
658         /* Signal to the threads there is a new job */
659         thread_signal_job();
660 
661         return 1;
662       }
663     }
664   }
665 
666   /* Wake up upstream and downstream vconnections */
667   wake_up_streams(contp);
668 
669   return 1;
670 }
671 
672 /*-------------------------------------------------------------------------
673   dump_psi
674   Dump the psi_output to the downstream vconnection.
675 
676   Input:
677     contp      continuation for the current transaction
678   Output :
679   Return Value:
680    0 if failure
681    1 if success
682   -------------------------------------------------------------------------*/
683 static int
dump_psi(TSCont contp)684 dump_psi(TSCont contp)
685 {
686   ContData *data;
687 
688   data = TSContDataGet(contp);
689   TSAssert(data->magic == MAGIC_ALIVE);
690 
691   /* If script exec succeeded, copy its output to the downstream vconn */
692   if (data->psi_success == 1) {
693     int psi_output_len = TSIOBufferReaderAvail(data->psi_reader);
694 
695     if (psi_output_len > 0) {
696       data->transform_bytes += psi_output_len;
697 
698       TSDebug(PLUGIN_NAME, "Inserting %d bytes from include file", psi_output_len);
699       /* TODO: Should we check the return value of TSIOBufferCopy() ? */
700       TSIOBufferCopy(TSVIOBufferGet(data->output_vio), data->psi_reader, psi_output_len, 0);
701       /* Consume all the output data */
702       TSIOBufferReaderConsume(data->psi_reader, psi_output_len);
703 
704       /* Reenable the output connection so it can read the data we've produced. */
705       TSVIOReenable(data->output_vio);
706     }
707   }
708 
709   /* Change state to finish up reading upstream data */
710   data->state = STATE_READ_DATA;
711   return 0;
712 }
713 
714 /*-------------------------------------------------------------------------
715   transform_handler
716   Handler for all events received during the transformation process
717 
718   Input:
719     contp      continuation for the current transaction
720     event      event received
721     data       pointer on optional data
722   Output :
723   Return Value:
724   -------------------------------------------------------------------------*/
725 static int
transform_handler(TSCont contp,TSEvent event,void * edata ATS_UNUSED)726 transform_handler(TSCont contp, TSEvent event, void *edata ATS_UNUSED)
727 {
728   ContData *data;
729   int state, retval;
730 
731   /* This section will be called by both TS internal
732      and the thread. Protect it with a mutex to avoid
733      concurrent calls. */
734 
735   /* Handle TryLock result */
736   if (TSMutexLockTry(TSContMutexGet(contp)) != TS_SUCCESS) {
737     TSCont c       = TSContCreate(trylock_handler, NULL);
738     TryLockData *d = TSmalloc(sizeof(TryLockData));
739 
740     d->contp = contp;
741     d->event = event;
742     TSContDataSet(c, d);
743     TSContScheduleOnPool(c, 10, TS_THREAD_POOL_NET);
744     return 1;
745   }
746 
747   data = TSContDataGet(contp);
748   TSAssert(data->magic == MAGIC_ALIVE);
749 
750   state = data->state;
751 
752   /* Check to see if the transformation has been closed */
753   retval = TSVConnClosedGet(contp);
754   if (retval) {
755     /* If the thread is still executing its job, we don't want to destroy
756        the continuation right away as the thread will call us back
757        on this continuation. */
758     if (state == STATE_READ_PSI) {
759       TSContScheduleOnPool(contp, 10, TS_THREAD_POOL_NET);
760     } else {
761       TSMutexUnlock(TSContMutexGet(contp));
762       cont_data_destroy(TSContDataGet(contp));
763       TSContDestroy(contp);
764       return 1;
765     }
766   } else {
767     TSVIO input_vio;
768     switch (event) {
769     case TS_EVENT_ERROR:
770       input_vio = TSVConnWriteVIOGet(contp);
771       TSContCall(TSVIOContGet(input_vio), TS_EVENT_ERROR, input_vio);
772       break;
773 
774     case TS_EVENT_VCONN_WRITE_COMPLETE:
775       TSVConnShutdown(TSTransformOutputVConnGet(contp), 0, 1);
776       break;
777 
778     case TS_EVENT_VCONN_WRITE_READY:
779       /* downstream vconnection is done reading data we've write into it.
780          let's read some more data from upstream if we're in read state. */
781       if (state == STATE_READ_DATA) {
782         handle_transform(contp);
783       }
784       break;
785 
786     case TS_EVENT_IMMEDIATE:
787       if (state == STATE_READ_DATA) {
788         /* upstream vconnection signals some more data ready to be read
789            let's try to transform some more data */
790         handle_transform(contp);
791       } else if (state == STATE_DUMP_PSI) {
792         /* The thread scheduled an event on our continuation to let us
793            know it has completed its job
794            Let's dump the include content to the output vconnection */
795         dump_psi(contp);
796         wake_up_streams(contp);
797       }
798       break;
799 
800     default:
801       TSAssert(!"Unexpected event");
802       break;
803     }
804   }
805 
806   TSMutexUnlock(TSContMutexGet(contp));
807   return 1;
808 }
809 
810 /*-------------------------------------------------------------------------
811   trylock_handler
812   Small handler to handle TSMutexLockTry failures
813 
814   Input:
815     contp      continuation for the current transaction
816     event      event received
817     data       pointer on optional data
818   Output :
819   Return Value:
820   -------------------------------------------------------------------------*/
821 static int
trylock_handler(TSCont contp,TSEvent event ATS_UNUSED,void * edata ATS_UNUSED)822 trylock_handler(TSCont contp, TSEvent event ATS_UNUSED, void *edata ATS_UNUSED)
823 {
824   TryLockData *data = TSContDataGet(contp);
825   transform_handler(data->contp, data->event, NULL);
826   TSfree(data);
827   TSContDestroy(contp);
828   return 0;
829 }
830 
831 /*-------------------------------------------------------------------------
832   transformable
833   Determine if the current transaction should be transformed or not
834 
835   Input:
836     txnp      current transaction
837   Output :
838   Return Value:
839     1  if transformable
840     0  if not
841   -------------------------------------------------------------------------*/
842 static int
transformable(TSHttpTxn txnp)843 transformable(TSHttpTxn txnp)
844 {
845   /*  We are only interested in transforming "200 OK" responses
846      with a Content-Type: text/ header and with X-Psi header */
847   TSMBuffer bufp;
848   TSMLoc hdr_loc;
849   TSHttpStatus resp_status;
850 
851   if (TS_SUCCESS == TSHttpTxnServerRespGet(txnp, &bufp, &hdr_loc)) {
852     resp_status = TSHttpHdrStatusGet(bufp, hdr_loc);
853     if (resp_status != TS_HTTP_STATUS_OK) {
854       TSHandleMLocRelease(bufp, TS_NULL_MLOC, hdr_loc);
855       return 0;
856     }
857 
858     TSMLoc field_loc;
859     field_loc = TSMimeHdrFieldFind(bufp, hdr_loc, TS_MIME_FIELD_CONTENT_TYPE, -1);
860     if (field_loc == TS_NULL_MLOC) {
861       TSError("[%s] Unable to search Content-Type field", PLUGIN_NAME);
862       TSHandleMLocRelease(bufp, TS_NULL_MLOC, hdr_loc);
863       return 0;
864     }
865 
866     const char *value = TSMimeHdrFieldValueStringGet(bufp, hdr_loc, field_loc, -1, NULL);
867     if ((value == NULL) || (strncasecmp(value, "text/", sizeof("text/") - 1) != 0)) {
868       TSHandleMLocRelease(bufp, hdr_loc, field_loc);
869       TSHandleMLocRelease(bufp, TS_NULL_MLOC, hdr_loc);
870       return 0;
871     }
872 
873     TSHandleMLocRelease(bufp, hdr_loc, field_loc);
874 
875     field_loc = TSMimeHdrFieldFind(bufp, hdr_loc, MIME_FIELD_XPSI, -1);
876 
877     TSHandleMLocRelease(bufp, hdr_loc, field_loc);
878     TSHandleMLocRelease(bufp, TS_NULL_MLOC, hdr_loc);
879   }
880 
881   return 1;
882 }
883 
884 /*-------------------------------------------------------------------------
885   transform_add
886   Create a transformation and alloc data structure
887 
888   Input:
889     txnp      current transaction
890   Output :
891   Return Value:
892     1  if transformation created
893     0  if not
894   -------------------------------------------------------------------------*/
895 static int
transform_add(TSHttpTxn txnp)896 transform_add(TSHttpTxn txnp)
897 {
898   TSCont contp;
899   ContData *data;
900 
901   contp = TSTransformCreate(transform_handler, txnp);
902   data  = cont_data_alloc();
903   TSContDataSet(contp, data);
904 
905   TSHttpTxnHookAdd(txnp, TS_HTTP_RESPONSE_TRANSFORM_HOOK, contp);
906   return 1;
907 }
908 
909 /*-------------------------------------------------------------------------
910   read_response_handler
911   Handler for events related to hook READ_RESPONSE
912 
913   Input:
914     contp      continuation for the current transaction
915     event      event received
916     data       pointer on eventual data
917   Output :
918   Return Value:
919   -------------------------------------------------------------------------*/
920 static int
read_response_handler(TSCont contp ATS_UNUSED,TSEvent event,void * edata)921 read_response_handler(TSCont contp ATS_UNUSED, TSEvent event, void *edata)
922 {
923   TSHttpTxn txnp = (TSHttpTxn)edata;
924 
925   switch (event) {
926   case TS_EVENT_HTTP_READ_RESPONSE_HDR:
927     if (transformable(txnp)) {
928       TSDebug(PLUGIN_NAME, "Add a transformation");
929       transform_add(txnp);
930     }
931     TSHttpTxnReenable(txnp, TS_EVENT_HTTP_CONTINUE);
932     return 0;
933   default:
934     break;
935   }
936 
937   return 0;
938 }
939 
940 /*-------------------------------------------------------------------------
941   TSPluginInit
942   Function called at plugin init time
943 
944   Input:
945     argc  number of args
946     argv  list vof args
947   Output :
948   Return Value:
949   -------------------------------------------------------------------------*/
950 void
TSPluginInit(int argc ATS_UNUSED,const char * argv[]ATS_UNUSED)951 TSPluginInit(int argc ATS_UNUSED, const char *argv[] ATS_UNUSED)
952 {
953   TSPluginRegistrationInfo info;
954   int i;
955   TSReturnCode retval;
956 
957   info.plugin_name   = "psi";
958   info.vendor_name   = "Apache";
959   info.support_email = "";
960 
961   if (TSPluginRegister(&info) != TS_SUCCESS) {
962     TSError("[%s] Plugin registration failed", PLUGIN_NAME);
963   }
964 
965   /* Initialize the psi directory = <plugin_path>/include */
966   snprintf(psi_directory, sizeof(psi_directory), "%s/%s", TSPluginDirGet(), PSI_PATH);
967 
968   /* create an TSTextLogObject to log any psi include */
969   retval = TSTextLogObjectCreate("psi", TS_LOG_MODE_ADD_TIMESTAMP, &log);
970   if (retval == TS_ERROR) {
971     TSError("[%s] Failed creating log for psi plugin", PLUGIN_NAME);
972     log = NULL;
973   }
974 
975   /* Create working threads */
976   thread_init();
977   init_queue(&job_queue);
978 
979   for (i = 0; i < NB_THREADS; i++) {
980     char *thread_name = (char *)TSmalloc(64);
981     sprintf(thread_name, "Thread[%d]", i);
982     if (!TSThreadCreate(thread_loop, thread_name)) {
983       TSError("[%s] Failed creating threads", PLUGIN_NAME);
984       return;
985     }
986   }
987 
988   TSHttpHookAdd(TS_HTTP_READ_RESPONSE_HDR_HOOK, TSContCreate(read_response_handler, TSMutexCreate()));
989   TSDebug(PLUGIN_NAME, "Plugin started");
990 }
991