|
22 | 22 | #include <strings.h> |
23 | 23 | #include <sstream> |
24 | 24 | #include <cstring> |
| 25 | +#include <atomic> |
| 26 | +#include <memory> |
25 | 27 | #include <getopt.h> |
26 | 28 | #include <cstdint> |
27 | 29 | #include <cinttypes> |
|
32 | 34 | #include "tscore/ink_defs.h" |
33 | 35 | #include "tscpp/util/PostScript.h" |
34 | 36 | #include "tscpp/util/TextView.h" |
| 37 | +#include "Cleanup.h" |
| 38 | + |
| 39 | +namespace |
| 40 | +{ |
| 41 | +struct BodyBuilder { |
| 42 | + atscppapi::TSContUniqPtr transform_connp; |
| 43 | + atscppapi::TSIOBufferUniqPtr output_buffer; |
| 44 | + // It's important that output_reader comes after output_buffer so it will be deleted first. |
| 45 | + atscppapi::TSIOBufferReaderUniqPtr output_reader; |
| 46 | + TSVIO output_vio = nullptr; |
| 47 | + bool wrote_prebody = false; |
| 48 | + bool wrote_body = false; |
| 49 | + bool hdr_ready = false; |
| 50 | + std::atomic_flag wrote_postbody; |
| 51 | + |
| 52 | + int64_t nbytes = 0; |
| 53 | +}; |
| 54 | + |
| 55 | +struct XDebugTxnAuxData { |
| 56 | + std::unique_ptr<BodyBuilder> body_builder; |
| 57 | + unsigned xheaders = 0; |
| 58 | +}; |
| 59 | + |
| 60 | +atscppapi::TxnAuxMgrData mgrData; |
| 61 | + |
| 62 | +using AuxDataMgr = atscppapi::TxnAuxDataMgr<XDebugTxnAuxData, mgrData>; |
| 63 | + |
| 64 | +} // end anonymous namespace |
35 | 65 |
|
36 | 66 | #include "xdebug_headers.cc" |
37 | 67 | #include "xdebug_transforms.cc" |
|
53 | 83 | XHEADER_X_PSELECT_KEY = 1u << 10, |
54 | 84 | }; |
55 | 85 |
|
56 | | -static int XArgIndex = 0; |
57 | | -static int BodyBuilderArgIndex = 0; |
58 | 86 | static TSCont XInjectHeadersCont = nullptr; |
59 | 87 | static TSCont XDeleteDebugHdrCont = nullptr; |
60 | 88 |
|
@@ -382,7 +410,7 @@ XInjectResponseHeaders(TSCont /* contp */, TSEvent event, void *edata) |
382 | 410 |
|
383 | 411 | TSReleaseAssert(event == TS_EVENT_HTTP_SEND_RESPONSE_HDR); |
384 | 412 |
|
385 | | - uintptr_t xheaders = reinterpret_cast<uintptr_t>(TSHttpTxnArgGet(txn, XArgIndex)); |
| 413 | + unsigned xheaders = AuxDataMgr::data(txn).xheaders; |
386 | 414 | if (xheaders == 0) { |
387 | 415 | goto done; |
388 | 416 | } |
@@ -422,14 +450,14 @@ XInjectResponseHeaders(TSCont /* contp */, TSEvent event, void *edata) |
422 | 450 | } |
423 | 451 |
|
424 | 452 | if (xheaders & XHEADER_X_PROBE_HEADERS) { |
425 | | - BodyBuilder *data = static_cast<BodyBuilder *>(TSHttpTxnArgGet(txn, BodyBuilderArgIndex)); |
| 453 | + BodyBuilder *data = AuxDataMgr::data(txn).body_builder.get(); |
426 | 454 | TSDebug("xdebug_transform", "XInjectResponseHeaders(): client resp header ready"); |
427 | 455 | if (data == nullptr) { |
428 | 456 | TSHttpTxnReenable(txn, TS_EVENT_HTTP_ERROR); |
429 | 457 | return TS_ERROR; |
430 | 458 | } |
431 | 459 | data->hdr_ready = true; |
432 | | - writePostBody(data); |
| 460 | + writePostBody(txn, data); |
433 | 461 | } |
434 | 462 |
|
435 | 463 | if (xheaders & XHEADER_X_PSELECT_KEY) { |
@@ -493,9 +521,9 @@ isFwdFieldValue(std::string_view value, intmax_t &fwdCnt) |
493 | 521 | static int |
494 | 522 | XScanRequestHeaders(TSCont /* contp */, TSEvent event, void *edata) |
495 | 523 | { |
496 | | - TSHttpTxn txn = static_cast<TSHttpTxn>(edata); |
497 | | - uintptr_t xheaders = 0; |
498 | | - intmax_t fwdCnt = 0; |
| 524 | + TSHttpTxn txn = static_cast<TSHttpTxn>(edata); |
| 525 | + unsigned xheaders = 0; |
| 526 | + intmax_t fwdCnt = 0; |
499 | 527 | TSMLoc field, next; |
500 | 528 | TSMBuffer buffer; |
501 | 529 | TSMLoc hdr; |
@@ -550,26 +578,17 @@ XScanRequestHeaders(TSCont /* contp */, TSEvent event, void *edata) |
550 | 578 | } else if (header_field_eq("probe", value, vsize)) { |
551 | 579 | xheaders |= XHEADER_X_PROBE_HEADERS; |
552 | 580 |
|
| 581 | + auto &auxData = AuxDataMgr::data(txn); |
| 582 | + |
553 | 583 | // prefix request headers and postfix response headers |
554 | 584 | BodyBuilder *data = new BodyBuilder(); |
555 | | - data->txn = txn; |
| 585 | + auxData.body_builder.reset(data); |
556 | 586 |
|
557 | 587 | TSVConn connp = TSTransformCreate(body_transform, txn); |
558 | | - TSContDataSet(connp, data); |
| 588 | + data->transform_connp.reset(connp); |
| 589 | + TSContDataSet(connp, txn); |
559 | 590 | TSHttpTxnHookAdd(txn, TS_HTTP_RESPONSE_TRANSFORM_HOOK, connp); |
560 | 591 |
|
561 | | - // store data pointer in txnarg to use in global cont XInjectResponseHeaders |
562 | | - TSHttpTxnArgSet(txn, BodyBuilderArgIndex, data); |
563 | | - |
564 | | - // create a self-cleanup on close |
565 | | - auto cleanupBodyBuilder = [](TSCont /* contp */, TSEvent event, void *edata) -> int { |
566 | | - TSHttpTxn txn = static_cast<TSHttpTxn>(edata); |
567 | | - BodyBuilder *data = static_cast<BodyBuilder *>(TSHttpTxnArgGet(txn, BodyBuilderArgIndex)); |
568 | | - delete data; |
569 | | - return TS_EVENT_NONE; |
570 | | - }; |
571 | | - TSHttpTxnHookAdd(txn, TS_HTTP_TXN_CLOSE_HOOK, TSContCreate(cleanupBodyBuilder, nullptr)); |
572 | | - |
573 | 592 | // disable writing to cache because we are injecting data into the body. |
574 | 593 | TSHttpTxnReqCacheableSet(txn, 0); |
575 | 594 | TSHttpTxnRespCacheableSet(txn, 0); |
@@ -608,7 +627,7 @@ XScanRequestHeaders(TSCont /* contp */, TSEvent event, void *edata) |
608 | 627 | TSDebug("xdebug", "adding response hook for header mask %p and forward count %" PRIiMAX, reinterpret_cast<void *>(xheaders), |
609 | 628 | fwdCnt); |
610 | 629 | TSHttpTxnHookAdd(txn, TS_HTTP_SEND_RESPONSE_HDR_HOOK, XInjectHeadersCont); |
611 | | - TSHttpTxnArgSet(txn, XArgIndex, reinterpret_cast<void *>(xheaders)); |
| 630 | + AuxDataMgr::data(txn).xheaders = xheaders; |
612 | 631 |
|
613 | 632 | if (fwdCnt == 0) { |
614 | 633 | // X-Debug header has to be deleted, but not too soon for other plugins to see it. |
@@ -688,9 +707,9 @@ TSPluginInit(int argc, const char *argv[]) |
688 | 707 | } |
689 | 708 | xDebugHeader.len = strlen(xDebugHeader.str); |
690 | 709 |
|
| 710 | + AuxDataMgr::init("xdebug"); |
| 711 | + |
691 | 712 | // Setup the global hook |
692 | | - TSReleaseAssert(TSHttpTxnArgIndexReserve("xdebug", "xdebug header requests", &XArgIndex) == TS_SUCCESS); |
693 | | - TSReleaseAssert(TSHttpTxnArgIndexReserve("bodyTransform", "BodyBuilder*", &XArgIndex) == TS_SUCCESS); |
694 | 713 | TSReleaseAssert(XInjectHeadersCont = TSContCreate(XInjectResponseHeaders, nullptr)); |
695 | 714 | TSReleaseAssert(XDeleteDebugHdrCont = TSContCreate(XDeleteDebugHdr, nullptr)); |
696 | 715 | TSHttpHookAdd(TS_HTTP_READ_REQUEST_HDR_HOOK, TSContCreate(XScanRequestHeaders, nullptr)); |
|
0 commit comments