From 933e0f40bb21c0252762e1a3a12cc97dc05e020e Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Tue, 19 Jan 2016 19:24:15 +0900
Subject: [PATCH] h2load: Support UNIX domain socket

---
 src/h2load.cc    | 62 ++++++++++++++++++++++++++++++++++++++++++++----
 src/h2load.h     |  6 +++++
 src/util.h       |  5 ++++
 src/util_test.cc |  3 +++
 4 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/src/h2load.cc b/src/h2load.cc
index 62d65bd3..d2aad6bf 100644
--- a/src/h2load.cc
+++ b/src/h2load.cc
@@ -86,7 +86,11 @@ Config::Config()
       port(0), default_port(0), verbose(false), timing_script(false) {}
 
 Config::~Config() {
-  freeaddrinfo(addrs);
+  if (base_uri_unix) {
+    delete addrs;
+  } else {
+    freeaddrinfo(addrs);
+  }
 
   if (data_fd != -1) {
     close(data_fd);
@@ -1301,6 +1305,18 @@ process_time_stats(const std::vector<std::unique_ptr<Worker>> &workers) {
 
 namespace {
 void resolve_host() {
+  if (config.base_uri_unix) {
+    auto res = make_unique<addrinfo>();
+    res->ai_family = config.unix_addr.sun_family;
+    res->ai_socktype = SOCK_STREAM;
+    res->ai_addrlen = sizeof(config.unix_addr);
+    res->ai_addr =
+        static_cast<struct sockaddr *>(static_cast<void *>(&config.unix_addr));
+
+    config.addrs = res.release();
+    return;
+  };
+
   int rv;
   addrinfo hints{}, *res;
 
@@ -1357,6 +1373,10 @@ int client_select_next_proto_cb(SSL *ssl, unsigned char **out,
 }
 } // namespace
 
+namespace {
+constexpr char UNIX_PATH_PREFIX[] = "unix:";
+} // namespace
+
 namespace {
 bool parse_base_uri(std::string base_uri) {
   http_parser_url u{};
@@ -1636,11 +1656,16 @@ Options:
               contained  in  other  URIs,  if  present,  are  ignored.
               Definition of a  base URI overrides all  scheme, host or
               port values.
-  -B, --base-uri=<URI>
+  -B, --base-uri=(<URI>|unix:<PATH>)
               Specify URI from which the scheme, host and port will be
               used  for  all requests.   The  base  URI overrides  all
               values  defined either  at  the command  line or  inside
-              input files.
+              input files.  If argument  starts with "unix:", then the
+              rest  of the  argument will  be treated  as UNIX  domain
+              socket path.   The connection is made  through that path
+              instead of TCP.   In this case, scheme  is inferred from
+              the first  URI appeared  in the  command line  or inside
+              input files as usual.
   --npn-list=<LIST>
               Comma delimited list of  ALPN protocol identifier sorted
               in the  order of preference.  That  means most desirable
@@ -1819,13 +1844,40 @@ int main(int argc, char **argv) {
         exit(EXIT_FAILURE);
       }
       break;
-    case 'B':
+    case 'B': {
+      config.base_uri = "";
+      config.base_uri_unix = false;
+
+      if (util::istarts_with_l(optarg, UNIX_PATH_PREFIX)) {
+        // UNIX domain socket path
+        sockaddr_un un;
+
+        auto path = optarg + str_size(UNIX_PATH_PREFIX);
+        auto pathlen = strlen(optarg) - str_size(UNIX_PATH_PREFIX);
+
+        if (pathlen == 0 || pathlen + 1 > sizeof(un.sun_path)) {
+          std::cerr << "--base-uri: invalid UNIX domain socket path: " << optarg
+                    << std::endl;
+          exit(EXIT_FAILURE);
+        }
+
+        config.base_uri_unix = true;
+
+        auto &unix_addr = config.unix_addr;
+        std::copy_n(path, pathlen + 1, unix_addr.sun_path);
+        unix_addr.sun_family = AF_UNIX;
+
+        break;
+      }
+
       if (!parse_base_uri(optarg)) {
-        std::cerr << "invalid base URI: " << optarg << std::endl;
+        std::cerr << "--base-uri: invalid base URI: " << optarg << std::endl;
         exit(EXIT_FAILURE);
       }
+
       config.base_uri = optarg;
       break;
+    }
     case 'v':
       config.verbose = true;
       break;
diff --git a/src/h2load.h b/src/h2load.h
index 0d95b8cd..c053b9c8 100644
--- a/src/h2load.h
+++ b/src/h2load.h
@@ -34,6 +34,7 @@
 #ifdef HAVE_NETDB_H
 #include <netdb.h>
 #endif // HAVE_NETDB_H
+#include <sys/un.h>
 
 #include <vector>
 #include <string>
@@ -100,6 +101,11 @@ struct Config {
   bool verbose;
   bool timing_script;
   std::string base_uri;
+  // true if UNIX domain socket is used.  In this case, base_uri is
+  // not used in usual way.
+  bool base_uri_unix;
+  // used when UNIX domain socket is used (base_uri_unix is true).
+  sockaddr_un unix_addr;
   // list of supported NPN/ALPN protocol strings in the order of
   // preference.
   std::vector<std::string> npn_list;
diff --git a/src/util.h b/src/util.h
index 669cc638..fd88f318 100644
--- a/src/util.h
+++ b/src/util.h
@@ -212,6 +212,11 @@ bool istarts_with(InputIt a, size_t an, const char *b) {
 
 bool istarts_with(const char *a, const char *b);
 
+template <size_t N>
+bool istarts_with_l(const std::string &a, const char(&b)[N]) {
+  return istarts_with(std::begin(a), std::end(a), b, b + N - 1);
+}
+
 template <typename InputIterator1, typename InputIterator2>
 bool ends_with(InputIterator1 first1, InputIterator1 last1,
                InputIterator2 first2, InputIterator2 last2) {
diff --git a/src/util_test.cc b/src/util_test.cc
index 97a35c4e..440674b0 100644
--- a/src/util_test.cc
+++ b/src/util_test.cc
@@ -362,6 +362,9 @@ void test_util_starts_with(void) {
   CU_ASSERT(util::starts_with("ofoo", ""));
   CU_ASSERT(util::istarts_with("fOOo", "Foo"));
   CU_ASSERT(!util::istarts_with("ofoo", "foo"));
+
+  CU_ASSERT(util::istarts_with_l("fOOo", "Foo"));
+  CU_ASSERT(!util::istarts_with_l("ofoo", "foo"));
 }
 
 void test_util_ends_with(void) {