Merge branch 'kennypeng-header_add_override'

src/h2load.cc

@@ -33,6 +33,7 @@
 #include <cassert>
 #include <cstdlib>
 #include <iostream>
+#include <fstream>
 #include <chrono>
 #include <thread>
 #include <future>
@@ -605,7 +606,8 @@ void print_version(std::ostream& out)
 namespace {
 void print_usage(std::ostream& out)
 {
-  out << R"(Usage: h2load [OPTIONS]... <URI>...
+  out << R"(
+Usage: h2load [OPTIONS]... [URI]...
 benchmarking tool for HTTP/2 and SPDY server)" << std::endl;
 }
 } // namespace
@@ -630,6 +632,15 @@ Options:
       << config.nclients << R"(
   -t, --threads=<N>  Number of native threads. Default: )"
       << config.nthreads << R"(
+  -i, --input-file=<FILE>
+                     Path of a file with multiple URIs  are  seperated
+                     by EOLs.    This option will disable URIs getting    
+                     from stdin. URIs are used  in this order for each
+                     client.   All URIs  are used,  then first  URI is
+                     used and  then 2nd URI,  and so on.   The scheme,
+                     host and port in the subsequent URIs, if present,
+                     are  ignored.  Those  in the  first URI  are used
+                     solely.
   -m, --max-concurrent-streams=(auto|<N>)
                      Max concurrent streams to  issue per session.  If
                      "auto"  is given,  the  number of  given URIs  is
@@ -642,6 +653,8 @@ Options:
                      (2**<N>)-1.  For  SPDY, if  <N> is  strictly less
                      than  16,  this  option  is  ignored.   Otherwise
                      2**<N> is used for SPDY.
+  -H, --header=<HEADER>
+                     Add/Override a header to the requests.
   -p, --no-tls-proto=<PROTOID>
                      Specify  ALPN identifier  of the  protocol to  be
                      used  when accessing  http  URI without  SSL/TLS.)";
@@ -674,6 +687,8 @@ int main(int argc, char **argv)
       {"max-concurrent-streams", required_argument, nullptr, 'm'},
       {"window-bits", required_argument, nullptr, 'w'},
       {"connection-window-bits", required_argument, nullptr, 'W'},
+      {"input-file", required_argument, nullptr, 'i'},
+      {"header", required_argument, nullptr, 'H'},
       {"no-tls-proto", required_argument, nullptr, 'p'},
       {"verbose", no_argument, nullptr, 'v'},
       {"help", no_argument, nullptr, 'h'},
@@ -681,7 +696,7 @@ int main(int argc, char **argv)
       {nullptr, 0, nullptr, 0 }
     };
     int option_index = 0;
-    auto c = getopt_long(argc, argv, "hvW:c:m:n:p:t:w:", long_options,
+    auto c = getopt_long(argc, argv, "hvW:c:m:n:p:t:w:H:i:", long_options,
                          &option_index);
     if(c == -1) {
       break;
@@ -727,6 +742,35 @@ int main(int argc, char **argv)
       }
       break;
     }
+    case 'H': {
+      char *header = optarg;
+      // Skip first possible ':' in the header name
+      char *value = strchr( optarg + 1, ':' );
+      if ( ! value || (header[0] == ':' && header + 1 == value)) {
+        std::cerr << "-H: invalid header: " << optarg
+                  << std::endl;
+        exit(EXIT_FAILURE);
+      }
+      *value = 0;
+      value++;
+      while( isspace( *value ) ) { value++; }
+      if ( *value == 0 ) {
+        // This could also be a valid case for suppressing a header
+        // similar to curl
+        std::cerr << "-H: invalid header - value missing: " << optarg
+                  << std::endl;
+        exit(EXIT_FAILURE);
+      }
+      // Note that there is no processing currently to handle multiple
+      // message-header fields with the same field name
+      config.custom_headers.emplace_back(header, value);
+      util::inp_strlower(config.custom_headers.back().first);
+      break;
+    }
+    case 'i': {
+      config.ifile = std::string(optarg);
+      break;
+    }
     case 'p':
       if(util::strieq(NGHTTP2_CLEARTEXT_PROTO_VERSION_ID, optarg)) {
         config.no_tls_proto = Config::PROTO_HTTP2;
@@ -766,9 +810,11 @@ int main(int argc, char **argv)
   }
 
   if(argc == optind) {
-    std::cerr << "no URI given" << std::endl;
+    if(config.ifile.empty()){
+      std::cerr << "no URI or input file given" << std::endl;
       exit(EXIT_FAILURE);
     }
+  }
 
   if(config.nreqs == 0) {
     std::cerr << "-n: the number of requests must be strictly greater than 0."
@@ -828,7 +874,25 @@ int main(int argc, char **argv)
   // this URI and ignore those in the remaining URIs if present.
   http_parser_url u;
   memset(&u, 0, sizeof(u));
-  auto uri = argv[optind];
+  auto uri = "";
+
+  std::ifstream uri_file;
+  std::string line_uri;
+
+  if (config.ifile.empty()) {
+    uri = argv[optind];
+  } else {
+    if (std::ifstream(config.ifile)) {
+      uri_file.open(config.ifile, std::ifstream::in);
+      // get first line as first uri
+      std::getline (uri_file, line_uri);
+      uri = (char *)line_uri.c_str();
+    } else {
+      std::cerr << "cannot read input file: " << config.ifile << std::endl;
+      exit(EXIT_FAILURE);
+    }
+  }
+
   if(http_parser_parse_url(uri, strlen(uri), 0, &u) != 0 ||
      !util::has_uri_field(u, UF_SCHEMA) || !util::has_uri_field(u, UF_HOST)) {
     std::cerr << "invalid URI: " << uri << std::endl;
@@ -847,6 +911,7 @@ int main(int argc, char **argv)
 
   reqlines.push_back(get_reqline(uri, u));
 
+  if (config.ifile.empty()) {
     ++optind;
     for(int i = optind; i < argc; ++i) {
       memset(&u, 0, sizeof(u));
@@ -860,6 +925,22 @@ int main(int argc, char **argv)
 
       reqlines.push_back(get_reqline(uri, u));
     }
+  } else {
+    if(uri_file.is_open()) {
+      //load rest uris from file 
+      while(std::getline (uri_file, line_uri)) {
+        auto uri = (char *)line_uri.c_str();
+
+        if(http_parser_parse_url(uri, strlen(uri), 0, &u) != 0) {
+          std::cerr << "invalid URI in input file: " << uri << std::endl;
+          exit(EXIT_FAILURE);
+        }
+
+        reqlines.push_back(get_reqline(uri, u));
+      }
+      uri_file.close();
+    }
+  }
 
   if(config.max_concurrent_streams == -1) {
     config.max_concurrent_streams = reqlines.size();
@@ -875,6 +956,24 @@ int main(int argc, char **argv)
   }
   shared_nva.emplace_back(":method", "GET");
 
+  //list overridalbe headers
+  std::vector<std::string> override_hdrs = {":host", ":scheme", ":method"};
+
+  for(auto& kv : config.custom_headers) {
+    if(std::find(override_hdrs.begin(), override_hdrs.end(), kv.first) != override_hdrs.end()) {
+      // override header
+      for(auto& nv : shared_nva) {
+        if( (nv.name == ":authority" && kv.first == ":host") 
+           || (nv.name == kv.first) ) {
+          nv.value = kv.second;
+        }
+      }
+    } else {
+      // add additional headers
+      shared_nva.emplace_back(kv.first.c_str(), kv.second.c_str());
+    }
+  }
+
   for(auto& req : reqlines) {
     // For nghttp2
     std::vector<nghttp2_nv> nva;

src/h2load.h

@@ -50,8 +50,10 @@ class Session;
 struct Config {
   std::vector<std::vector<nghttp2_nv>> nva;
   std::vector<std::vector<const char*>> nv;
+  std::vector<std::pair<std::string, std::string>> custom_headers;
   std::string scheme;
   std::string host;
+  std::string ifile;
   addrinfo *addrs;
   size_t nreqs;
   size_t nclients;