slice plugin: add --include-regex, --exclude-regex parameters (#6701)

traeak · web-flow · commit 8fdf694436e6 · 2020-04-28T13:30:54.000-06:00
diff --git a/doc/admin-guide/plugins/slice.en.rst b/doc/admin-guide/plugins/slice.en.rst
@@ -88,6 +88,17 @@ The slice plugin supports the following options::
         Disable writing block stitch errors to the error log.
         -d for short
 
+    --exclude-regex=<regex> (optional)
+        If provided, only slice what matches.
+        If not provided will always slice
+        Cannot be used with --include-regex
+        -e for short
+
+    --include-regex=<regex> (optional)
+        If provided, only slice what matches.
+        If not provided will always slice
+        Cannot be used with --exclude-regex
+        -i for short
 
 Examples::
 
@@ -122,6 +133,16 @@ After modifying :file:`remap.config`, restart or reload traffic server
 (sudo traffic_ctl config reload) or (sudo traffic_ctl server restart)
 to activate the new configuration values.
 
+Don't slice txt files::
+
+  slice.so --exclude-regex=\\.txt
+  slice.so -e \\.txt
+
+Slice only mp4 files::
+
+  slice.so --include-regex=\\.mp4
+  slice.so -i \\.mp4
+
 Debug Options
 -------------
 
diff --git a/plugins/experimental/slice/Config.cc b/plugins/experimental/slice/Config.cc
@@ -26,6 +26,20 @@
 
 #include "ts/experimental.h"
 
+Config::~Config()
+{
+  if (nullptr != m_regex_extra) {
+#ifndef PCRE_STUDY_JIT_COMPILE
+    pcre_free(m_regex_extra);
+#else
+    pcre_free_study(m_regex_extra);
+#endif
+  }
+  if (nullptr != m_regex) {
+    pcre_free(m_regex);
+  }
+}
+
 int64_t
 Config::bytesFrom(char const *const valstr)
 {
@@ -94,18 +108,20 @@ Config::fromArgs(int const argc, char const *const argv[])
   // standard parsing
   constexpr struct option longopts[] = {
     {const_cast<char *>("blockbytes"), required_argument, nullptr, 'b'},
-    {const_cast<char *>("blockbytes-test"), required_argument, nullptr, 't'},
-    {const_cast<char *>("remap-host"), required_argument, nullptr, 'r'},
-    {const_cast<char *>("pace-errorlog"), required_argument, nullptr, 'p'},
     {const_cast<char *>("disable-errorlog"), no_argument, nullptr, 'd'},
+    {const_cast<char *>("exclude-regex"), required_argument, nullptr, 'e'},
+    {const_cast<char *>("include-regex"), required_argument, nullptr, 'i'},
     {const_cast<char *>("throttle"), no_argument, nullptr, 'o'},
+    {const_cast<char *>("pace-errorlog"), required_argument, nullptr, 'p'},
+    {const_cast<char *>("remap-host"), required_argument, nullptr, 'r'},
+    {const_cast<char *>("blockbytes-test"), required_argument, nullptr, 't'},
     {nullptr, 0, nullptr, 0},
   };
 
   // getopt assumes args start at '1' so this hack is needed
   char *const *argvp = (const_cast<char *const *>(argv) - 1);
   for (;;) {
-    int const opt = getopt_long(argc + 1, argvp, "b:t:r:p:do", longopts, nullptr);
+    int const opt = getopt_long(argc + 1, argvp, "b:de:i:op:r:t:", longopts, nullptr);
     if (-1 == opt) {
       break;
     }
@@ -122,6 +138,61 @@ Config::fromArgs(int const argc, char const *const argv[])
         ERROR_LOG("Invalid blockbytes: %s", optarg);
       }
     } break;
+    case 'd': {
+      m_paceerrsecs = -1;
+    } break;
+    case 'e': {
+      if (None != m_regex_type) {
+        ERROR_LOG("Regex already specified!");
+        break;
+      }
+
+      const char *errptr;
+      int erroffset;
+      m_regexstr = optarg;
+      m_regex    = pcre_compile(m_regexstr.c_str(), 0, &errptr, &erroffset, NULL);
+      if (nullptr == m_regex) {
+        ERROR_LOG("Invalid regex: '%s'", m_regexstr.c_str());
+      } else {
+        m_regex_type  = Exclude;
+        m_regex_extra = pcre_study(m_regex, 0, &errptr);
+        DEBUG_LOG("Using regex for url exclude: '%s'", m_regexstr.c_str());
+      }
+    } break;
+    case 'i': {
+      if (None != m_regex_type) {
+        ERROR_LOG("Regex already specified!");
+        break;
+      }
+
+      const char *errptr;
+      int erroffset;
+      m_regexstr = optarg;
+      m_regex    = pcre_compile(m_regexstr.c_str(), 0, &errptr, &erroffset, NULL);
+      if (nullptr == m_regex) {
+        ERROR_LOG("Invalid regex: '%s'", m_regexstr.c_str());
+      } else {
+        m_regex_type  = Include;
+        m_regex_extra = pcre_study(m_regex, 0, &errptr);
+        DEBUG_LOG("Using regex for url include: '%s'", m_regexstr.c_str());
+      }
+    } break;
+    case 'o': {
+      m_throttle = true;
+      DEBUG_LOG("Enabling internal block throttling");
+    } break;
+    case 'p': {
+      int const secsread = atoi(optarg);
+      if (0 < secsread) {
+        m_paceerrsecs = std::min(secsread, 60);
+      } else {
+        ERROR_LOG("Ignoring pace-errlog argument");
+      }
+    } break;
+    case 'r': {
+      m_remaphost = optarg;
+      DEBUG_LOG("Using loopback remap host override: %s", m_remaphost.c_str());
+    } break;
     case 't': {
       if (0 == blockbytes) {
         int64_t const bytesread = bytesFrom(optarg);
@@ -135,25 +206,6 @@ Config::fromArgs(int const argc, char const *const argv[])
         DEBUG_LOG("Skipping blockbytes-test in favor of blockbytes");
       }
     } break;
-    case 'r':
-      m_remaphost = optarg;
-      DEBUG_LOG("Using loopback remap host override: %s", m_remaphost.c_str());
-      break;
-    case 'p': {
-      int const secsread = atoi(optarg);
-      if (0 < secsread) {
-        m_paceerrsecs = std::min(secsread, 60);
-      } else {
-        ERROR_LOG("Ignoring pace-errlog argument");
-      }
-    } break;
-    case 'd':
-      m_paceerrsecs = -1;
-      break;
-    case 'o':
-      m_throttle = true;
-      DEBUG_LOG("Enabling internal block throttling");
-      break;
     default:
       break;
     }
@@ -204,3 +256,26 @@ Config::canLogError()
 
   return true;
 }
+
+bool
+Config::matchesRegex(char const *const url, int const urllen) const
+{
+  bool matches = true;
+
+  switch (m_regex_type) {
+  case Exclude: {
+    if (0 <= pcre_exec(m_regex, m_regex_extra, url, urllen, 0, 0, NULL, 0)) {
+      matches = false;
+    }
+  } break;
+  case Include: {
+    if (pcre_exec(m_regex, m_regex_extra, url, urllen, 0, 0, NULL, 0) < 0) {
+      matches = false;
+    }
+  } break;
+  default:
+    break;
+  }
+
+  return matches;
+}
diff --git a/plugins/experimental/slice/Config.h b/plugins/experimental/slice/Config.h
@@ -20,6 +20,12 @@
 
 #include "slice.h"
 
+#ifdef HAVE_PCRE_PCRE_H
+#include <pcre/pcre.h>
+#else
+#include <pcre.h>
+#endif
+
 #include <mutex>
 
 // Data Structures and Classes
@@ -30,18 +36,36 @@ struct Config {
 
   int64_t m_blockbytes{blockbytesdefault};
   std::string m_remaphost; // remap host to use for loopback slice GET
-  bool m_throttle{false};  // internal block throttling
-  int m_paceerrsecs{0};    // -1 disable logging, 0 no pacing, max 60s
+  std::string m_regexstr;  // regex string for things to slice (default all)
+  enum RegexType { None, Include, Exclude };
+  RegexType m_regex_type{None};
+  pcre *m_regex{nullptr};
+  pcre_extra *m_regex_extra{nullptr};
+  bool m_throttle{false}; // internal block throttling
+  int m_paceerrsecs{0};   // -1 disable logging, 0 no pacing, max 60s
 
   // Convert optarg to bytes
   static int64_t bytesFrom(char const *const valstr);
 
+  // clean up pcre if applicable
+  ~Config();
+
   // Parse from args, ast one wins
   bool fromArgs(int const argc, char const *const argv[]);
 
   // Check if the error should can be logged, if sucessful may update m_nexttime
   bool canLogError();
 
+  // Check if regex supplied
+  bool
+  hasRegex() const
+  {
+    return None != m_regex_type;
+  }
+
+  // If no null reg, true, otherwise check against regex
+  bool matchesRegex(char const *const url, int const urllen) const;
+
 private:
   TSHRTime m_nextlogtime{0}; // next time to log in ns
   std::mutex m_mutex;
diff --git a/plugins/experimental/slice/HttpHeader.cc b/plugins/experimental/slice/HttpHeader.cc
@@ -54,18 +54,20 @@ HttpHeader::setStatus(TSHttpStatus const newstatus)
 }
 
 char *
-HttpHeader ::urlString(int *const urllen) const
+HttpHeader::urlString(int *const urllen) const
 {
   char *urlstr = nullptr;
   TSAssert(nullptr != urllen);
 
   TSMLoc locurl            = nullptr;
   TSReturnCode const rcode = TSHttpHdrUrlGet(m_buffer, m_lochdr, &locurl);
-  if (TS_SUCCESS == rcode && nullptr != locurl) {
-    urlstr = TSUrlStringGet(m_buffer, locurl, urllen);
+  if (nullptr != locurl) {
+    if (TS_SUCCESS == rcode) {
+      urlstr = TSUrlStringGet(m_buffer, locurl, urllen);
+    } else {
+      *urllen = 0;
+    }
     TSHandleMLocRelease(m_buffer, m_lochdr, locurl);
-  } else {
-    *urllen = 0;
   }
 
   return urlstr;
diff --git a/plugins/experimental/slice/Makefile.inc b/plugins/experimental/slice/Makefile.inc
@@ -49,16 +49,22 @@ experimental_slice_test_content_range_SOURCES = \
   experimental/slice/unit-tests/test_content_range.cc \
   experimental/slice/ContentRange.cc
 
+experimental_slice_test_content_range_LDADD = @LIBPCRE@
+
 check_PROGRAMS += experimental/slice/test_range
 
 experimental_slice_test_range_CPPFLAGS = $(AM_CPPFLAGS) -I$(abs_top_srcdir)/tests/include -DUNITTEST
 experimental_slice_test_range_SOURCES = \
   experimental/slice/unit-tests/test_range.cc \
   experimental/slice/Range.cc
 
+experimental_slice_test_range_LDADD = @LIBPCRE@
+
 check_PROGRAMS += experimental/slice/test_config
 
 experimental_slice_test_config_CPPFLAGS = $(AM_CPPFLAGS) -I$(abs_top_srcdir)/tests/include -DUNITTEST
 experimental_slice_test_config_SOURCES = \
   experimental/slice/unit-tests/test_config.cc \
   experimental/slice/Config.cc
+
+experimental_slice_test_config_LDADD = @LIBPCRE@
diff --git a/plugins/experimental/slice/slice.cc b/plugins/experimental/slice/slice.cc
@@ -59,10 +59,26 @@ read_request(TSHttpTxn txnp, Config *const config)
       // check if any previous plugin has monkeyed with the transaction status
       TSHttpStatus const txnstat = TSHttpTxnStatusGet(txnp);
       if (0 != (int)txnstat) {
-        DEBUG_LOG("slice: txn status change detected (%d), skipping plugin\n", (int)txnstat);
+        DEBUG_LOG("txn status change detected (%d), skipping plugin\n", (int)txnstat);
         return false;
       }
 
+      if (config->hasRegex()) {
+        int urllen         = 0;
+        char *const urlstr = TSHttpTxnEffectiveUrlStringGet(txnp, &urllen);
+        if (nullptr != urlstr) {
+          bool const shouldslice = config->matchesRegex(urlstr, urllen);
+          if (!shouldslice) {
+            DEBUG_LOG("request failed regex, not slicing: '%.*s'", urllen, urlstr);
+            TSfree(urlstr);
+            return false;
+          }
+
+          DEBUG_LOG("request passed regex, slicing: '%.*s'", urllen, urlstr);
+          TSfree(urlstr);
+        }
+      }
+
       // turn off any and all transaction caching (shouldn't matter)
       TSHttpTxnServerRespNoStoreSet(txnp, 1);
       TSHttpTxnRespCacheableSet(txnp, 0);
@@ -99,8 +115,8 @@ read_request(TSHttpTxn txnp, Config *const config)
       // is the plugin configured to use a remap host?
       std::string const &newhost = config->m_remaphost;
       if (newhost.empty()) {
-        TSMBuffer urlbuf;
-        TSMLoc urlloc;
+        TSMBuffer urlbuf   = nullptr;
+        TSMLoc urlloc      = nullptr;
         TSReturnCode rcode = TSHttpTxnPristineUrlGet(txnp, &urlbuf, &urlloc);
 
         if (TS_SUCCESS == rcode) {
diff --git a/tests/gold_tests/pluginTest/slice/slice_regex.test.py b/tests/gold_tests/pluginTest/slice/slice_regex.test.py
