Fix problems with "Probe" option for X-Debug MIME header field.

plugins/xdebug/Cleanup.h

@@ -0,0 +1,186 @@
+/**
+  Licensed to the Apache Software Foundation (ASF) under one
+  or more contributor license agreements.  See the NOTICE file
+  distributed with this work for additional information
+  regarding copyright ownership.  The ASF licenses this file
+  to you under the Apache License, Version 2.0 (the
+  "License"); you may not use this file except in compliance
+  with the License.  You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+ */
+
+/**
+ * @file Cleanup.h
+ * @brief Easy-to-use utilities to avoid resource leaks or double-releases of resources.  Independent of the rest
+ *        of the CPPAPI.
+ */
+
+#pragma once
+
+#include <type_traits>
+#include <memory>
+
+#include <ts/ts.h>
+
+namespace atscppapi
+{
+// For TS API types TSXxx with a TSXxxDestroy function, define standard deleter TSXxxDeleter, and use it to
+// define TSXxxUniqPtr (specialization of std::unique_ptr).  X() is used when the destroy function returns void,
+// Y() is used when the destroy function returns TSReturnCode.
+
+#if defined(X)
+#error "X defined as preprocessor symbol"
+#endif
+
+#define X(NAME_SEGMENT)               \
+  struct TS##NAME_SEGMENT##Deleter {  \
+    void                              \
+    operator()(TS##NAME_SEGMENT ptr)  \
+    {                                 \
+      TS##NAME_SEGMENT##Destroy(ptr); \
+    }                                 \
+  };                                  \
+  using TS##NAME_SEGMENT##UniqPtr = std::unique_ptr<std::remove_pointer_t<TS##NAME_SEGMENT>, TS##NAME_SEGMENT##Deleter>;
+
+#if defined(Y)
+#error "Y defined as preprocessor symbol"
+#endif
+
+#define Y(NAME_SEGMENT)                                       \
+  struct TS##NAME_SEGMENT##Deleter {                          \
+    void                                                      \
+    operator()(TS##NAME_SEGMENT ptr)                          \
+    {                                                         \
+      TSAssert(TS##NAME_SEGMENT##Destroy(ptr) == TS_SUCCESS); \
+    }                                                         \
+  };                                                          \
+  using TS##NAME_SEGMENT##UniqPtr = std::unique_ptr<std::remove_pointer_t<TS##NAME_SEGMENT>, TS##NAME_SEGMENT##Deleter>;
+
+Y(MBuffer)       // Defines TSMBufferDeleter and TSMBufferUniqPtr.
+X(MimeParser)    // Defines TSMimeParserDeleter and TSMimeParserUniqPtr.
+X(Thread)        // Defines TSThreadDeleter and TSThreadUniqPtr.
+X(Mutex)         // Defines TSMutexDeleter and TSMutexUniqPtr.
+Y(CacheKey)      // Defines TSCacheKeyDeleter and TSCacheKeyUniqPtr.
+X(Cont)          // Defines TSContDeleter and TSContUniqPtr.
+X(SslContext)    // Defines TSSslContextDeleter and TSSslContextUniqPtr.
+X(IOBuffer)      // Defines TSIOBufferDeleter and TSIOBufferUniqPtr.
+Y(TextLogObject) // Defines TSTextLogObjectDeleter and TSTextLogObjectUniqPtr.
+X(Uuid)          // Defines TSUuidDeleter and TSUuidUniqPtr.
+
+#undef X
+#undef Y
+
+// Deleter and unique pointer for memory buffer returned by TSalloc(), TSrealloc(), Tstrdup(), TSsrtndup().
+//
+struct TSMemDeleter {
+  void
+  operator()(void *ptr)
+  {
+    TSfree(ptr);
+  }
+};
+using TSMemUniqPtr = std::unique_ptr<void, TSMemDeleter>;
+
+// Deleter and unique pointer for TSIOBufferReader.  Care must be taken that the reader is deleted before the
+// TSIOBuffer to which it refers is deleted.
+//
+struct TSIOBufferReaderDeleter {
+  void
+  operator()(TSIOBufferReader ptr)
+  {
+    TSIOBufferReaderFree(ptr);
+  }
+};
+using TSIOBufferReaderUniqPtr = std::unique_ptr<std::remove_pointer_t<TSIOBufferReader>, TSIOBufferReaderDeleter>;
+
+class TxnAuxDataMgrBase
+{
+protected:
+  struct MgrData_ {
+    TSCont txnCloseContp = nullptr;
+    int txnArgIndex      = -1;
+  };
+
+public:
+  class MgrData : private MgrData_
+  {
+    friend class TxnAuxDataMgrBase;
+  };
+
+protected:
+  static MgrData_ &
+  access(MgrData &md)
+  {
+    return md;
+  }
+};
+
+using TxnAuxMgrData = TxnAuxDataMgrBase::MgrData;
+
+// Class to manage auxilliary data for a transaction.  If an instance is created for the transaction, the instance
+// will be deleted on the TXN_CLOSE transaction hook (which is always triggered for all transactions).
+// The TxnAuxData class must have a public default constructor.
+//
+template <class TxnAuxData, TxnAuxMgrData &MDRef> class TxnAuxDataMgr : private TxnAuxDataMgrBase
+{
+public:
+  using Data = TxnAuxData;
+
+  // This must be called from the plugin init function.  arg_name is the name for the transaction argument used
+  // to store the pointer to the auxiliary data class instance.  Repeated calls are ignored.
+  //
+  static void
+  init(char const *arg_name, char const *arg_desc = "per-transaction auxiliary data")
+  {
+    MgrData_ &md = access(MDRef);
+
+    if (md.txnArgIndex >= 0) {
+      return;
+    }
+
+    TSReleaseAssert(TSHttpTxnArgIndexReserve(arg_name, arg_desc, &md.txnArgIndex) == TS_SUCCESS);
+    TSReleaseAssert(md.txnCloseContp = TSContCreate(_deleteAuxData, nullptr));
+  }
+
+  // Get a reference to the auxiliary data for a transaction.
+  //
+  static TxnAuxData &
+  data(TSHttpTxn txn)
+  {
+    MgrData_ &md = access(MDRef);
+
+    TSAssert(md.txnArgIndex >= 0);
+
+    auto d = static_cast<TxnAuxData *>(TSHttpTxnArgGet(txn, md.txnArgIndex));
+    if (!d) {
+      d = new TxnAuxData;
+
+      TSHttpTxnArgSet(txn, md.txnArgIndex, d);
+
+      TSHttpTxnHookAdd(txn, TS_HTTP_TXN_CLOSE_HOOK, md.txnCloseContp);
+    }
+    return *d;
+  }
+
+private:
+  static int
+  _deleteAuxData(TSCont, TSEvent, void *edata)
+  {
+    MgrData_ &md = access(MDRef);
+
+    auto txn  = static_cast<TSHttpTxn>(edata);
+    auto data = static_cast<TxnAuxData *>(TSHttpTxnArgGet(txn, md.txnArgIndex));
+    delete data;
+    TSHttpTxnReenable(txn, TS_EVENT_HTTP_CONTINUE);
+    return 0;
+  };
+};
+
+} // end namespace atscppapi

plugins/xdebug/Makefile.inc

@@ -15,4 +15,6 @@
 #  limitations under the License.
 
 pkglib_LTLIBRARIES += xdebug/xdebug.la
-xdebug_xdebug_la_SOURCES = xdebug/xdebug.cc
+xdebug_xdebug_la_SOURCES = \
+xdebug/Cleanup.h \
+xdebug/xdebug.cc

plugins/xdebug/xdebug.cc

@@ -22,6 +22,8 @@
 #include <strings.h>
 #include <sstream>
 #include <cstring>
+#include <atomic>
+#include <memory>
 #include <getopt.h>
 #include <cstdint>
 #include <cinttypes>
@@ -32,6 +34,34 @@
 #include "tscore/ink_defs.h"
 #include "tscpp/util/PostScript.h"
 #include "tscpp/util/TextView.h"
+#include "Cleanup.h"
+
+namespace
+{
+struct BodyBuilder {
+  atscppapi::TSContUniqPtr transform_connp;
+  atscppapi::TSIOBufferUniqPtr output_buffer;
+  // It's important that output_reader comes after output_buffer so it will be deleted first.
+  atscppapi::TSIOBufferReaderUniqPtr output_reader;
+  TSVIO output_vio   = nullptr;
+  bool wrote_prebody = false;
+  bool wrote_body    = false;
+  bool hdr_ready     = false;
+  std::atomic_flag wrote_postbody;
+
+  int64_t nbytes = 0;
+};
+
+struct XDebugTxnAuxData {
+  std::unique_ptr<BodyBuilder> body_builder;
+  unsigned xheaders = 0;
+};
+
+atscppapi::TxnAuxMgrData mgrData;
+
+using AuxDataMgr = atscppapi::TxnAuxDataMgr<XDebugTxnAuxData, mgrData>;
+
+} // end anonymous namespace
 
 #include "xdebug_headers.cc"
 #include "xdebug_transforms.cc"
@@ -52,8 +82,6 @@ enum {
   XHEADER_X_PROBE_HEADERS  = 1u << 9,
 };
 
-static int XArgIndex              = 0;
-static int BodyBuilderArgIndex    = 0;
 static TSCont XInjectHeadersCont  = nullptr;
 static TSCont XDeleteDebugHdrCont = nullptr;
 
@@ -334,7 +362,7 @@ XInjectResponseHeaders(TSCont /* contp */, TSEvent event, void *edata)
 
   TSReleaseAssert(event == TS_EVENT_HTTP_SEND_RESPONSE_HDR);
 
-  uintptr_t xheaders = reinterpret_cast<uintptr_t>(TSHttpTxnArgGet(txn, XArgIndex));
+  unsigned xheaders = AuxDataMgr::data(txn).xheaders;
   if (xheaders == 0) {
     goto done;
   }
@@ -374,14 +402,14 @@ XInjectResponseHeaders(TSCont /* contp */, TSEvent event, void *edata)
   }
 
   if (xheaders & XHEADER_X_PROBE_HEADERS) {
-    BodyBuilder *data = static_cast<BodyBuilder *>(TSHttpTxnArgGet(txn, BodyBuilderArgIndex));
+    BodyBuilder *data = AuxDataMgr::data(txn).body_builder.get();
     TSDebug("xdebug_transform", "XInjectResponseHeaders(): client resp header ready");
     if (data == nullptr) {
       TSHttpTxnReenable(txn, TS_EVENT_HTTP_ERROR);
       return TS_ERROR;
     }
     data->hdr_ready = true;
-    writePostBody(data);
+    writePostBody(txn, data);
   }
 
 done:
@@ -441,9 +469,9 @@ isFwdFieldValue(std::string_view value, intmax_t &fwdCnt)
 static int
 XScanRequestHeaders(TSCont /* contp */, TSEvent event, void *edata)
 {
-  TSHttpTxn txn      = static_cast<TSHttpTxn>(edata);
-  uintptr_t xheaders = 0;
-  intmax_t fwdCnt    = 0;
+  TSHttpTxn txn     = static_cast<TSHttpTxn>(edata);
+  unsigned xheaders = 0;
+  intmax_t fwdCnt   = 0;
   TSMLoc field, next;
   TSMBuffer buffer;
   TSMLoc hdr;
@@ -498,26 +526,17 @@ XScanRequestHeaders(TSCont /* contp */, TSEvent event, void *edata)
       } else if (header_field_eq("probe", value, vsize)) {
         xheaders |= XHEADER_X_PROBE_HEADERS;
 
+        auto &auxData = AuxDataMgr::data(txn);
+
         // prefix request headers and postfix response headers
         BodyBuilder *data = new BodyBuilder();
-        data->txn         = txn;
+        auxData.body_builder.reset(data);
 
         TSVConn connp = TSTransformCreate(body_transform, txn);
-        TSContDataSet(connp, data);
+        data->transform_connp.reset(connp);
+        TSContDataSet(connp, txn);
         TSHttpTxnHookAdd(txn, TS_HTTP_RESPONSE_TRANSFORM_HOOK, connp);
 
-        // store data pointer in txnarg to use in global cont XInjectResponseHeaders
-        TSHttpTxnArgSet(txn, BodyBuilderArgIndex, data);
-
-        // create a self-cleanup on close
-        auto cleanupBodyBuilder = [](TSCont /* contp */, TSEvent event, void *edata) -> int {
-          TSHttpTxn txn     = static_cast<TSHttpTxn>(edata);
-          BodyBuilder *data = static_cast<BodyBuilder *>(TSHttpTxnArgGet(txn, BodyBuilderArgIndex));
-          delete data;
-          return TS_EVENT_NONE;
-        };
-        TSHttpTxnHookAdd(txn, TS_HTTP_TXN_CLOSE_HOOK, TSContCreate(cleanupBodyBuilder, nullptr));
-
         // disable writing to cache because we are injecting data into the body.
         TSHttpTxnReqCacheableSet(txn, 0);
         TSHttpTxnRespCacheableSet(txn, 0);
@@ -553,7 +572,7 @@ XScanRequestHeaders(TSCont /* contp */, TSEvent event, void *edata)
     TSDebug("xdebug", "adding response hook for header mask %p and forward count %" PRIiMAX, reinterpret_cast<void *>(xheaders),
             fwdCnt);
     TSHttpTxnHookAdd(txn, TS_HTTP_SEND_RESPONSE_HDR_HOOK, XInjectHeadersCont);
-    TSHttpTxnArgSet(txn, XArgIndex, reinterpret_cast<void *>(xheaders));
+    AuxDataMgr::data(txn).xheaders = xheaders;
 
     if (fwdCnt == 0) {
       // X-Debug header has to be deleted, but not too soon for other plugins to see it.
@@ -633,9 +652,9 @@ TSPluginInit(int argc, const char *argv[])
   }
   xDebugHeader.len = strlen(xDebugHeader.str);
 
+  AuxDataMgr::init("xdebug");
+
   // Setup the global hook
-  TSReleaseAssert(TSHttpTxnArgIndexReserve("xdebug", "xdebug header requests", &XArgIndex) == TS_SUCCESS);
-  TSReleaseAssert(TSHttpTxnArgIndexReserve("bodyTransform", "BodyBuilder*", &XArgIndex) == TS_SUCCESS);
   TSReleaseAssert(XInjectHeadersCont = TSContCreate(XInjectResponseHeaders, nullptr));
   TSReleaseAssert(XDeleteDebugHdrCont = TSContCreate(XDeleteDebugHdr, nullptr));
   TSHttpHookAdd(TS_HTTP_READ_REQUEST_HDR_HOOK, TSContCreate(XScanRequestHeaders, nullptr));