From 7271537a1531bad7ca0e2153b1cf188b6462dda4 Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Fri, 24 Sep 2021 18:01:47 +0900
Subject: [PATCH] nghttpx: Add --rlimit-memlock option

---
 docker/README.rst   |  4 ++--
 gennghttpxfun.py    |  1 +
 src/shrpx.cc        | 21 +++++++++++++++++++++
 src/shrpx_config.cc | 22 ++++++++++++++++++++++
 src/shrpx_config.h  |  4 ++++
 5 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/docker/README.rst b/docker/README.rst
index 7f3a6b59..2c32565a 100644
--- a/docker/README.rst
+++ b/docker/README.rst
@@ -20,6 +20,6 @@ certificate in server.key and server.crt respectively :
 .. code-block:: text
 
    $ docker run --rm -it -v $PWD:/shared --net=host --privileged \
-         --ulimit memlock=2048000 nghttp2 nghttpx \
+         nghttp2 nghttpx \
          /shared/server.key /shared/server.crt \
-         -f'*,443;quic'
+         -f'*,443;quic' --rlimit-memlock 262144
diff --git a/gennghttpxfun.py b/gennghttpxfun.py
index 2ee66ae8..a38fbca1 100755
--- a/gennghttpxfun.py
+++ b/gennghttpxfun.py
@@ -194,6 +194,7 @@ OPTIONS = [
     "frontend-quic-congestion-controller",
     "frontend-quic-server-id",
     "frontend-quic-secret-file",
+    "rlimit-memlock",
 ]
 
 LOGVARS = [
diff --git a/src/shrpx.cc b/src/shrpx.cc
index 8fc8893f..ede06a3e 100644
--- a/src/shrpx.cc
+++ b/src/shrpx.cc
@@ -2376,6 +2376,12 @@ Performance:
               If 0 is given, nghttpx does not set the limit.
               Default: )"
       << config->rlimit_nofile << R"(
+  --rlimit-memlock=<N>
+              Set maximum number of bytes of memory that may be locked
+              into  RAM.  If  0 is  given,  nghttpx does  not set  the
+              limit.
+              Default: )"
+      << config->rlimit_memlock << R"(
   --backend-request-buffer=<SIZE>
               Set buffer size used to store backend request.
               Default: )"
@@ -3599,6 +3605,16 @@ int process_options(Config *config,
     }
   }
 
+  if (config->rlimit_memlock) {
+    struct rlimit lim = {static_cast<rlim_t>(config->rlimit_memlock),
+                         static_cast<rlim_t>(config->rlimit_memlock)};
+    if (setrlimit(RLIMIT_MEMLOCK, &lim) != 0) {
+      auto error = errno;
+      LOG(WARN) << "Setting rlimit-memlock failed: "
+                << xsi_strerror(error, errbuf.data(), errbuf.size());
+    }
+  }
+
   auto &fwdconf = config->http.forwarded;
 
   if (fwdconf.by_node_type == ForwardedNode::OBFUSCATED &&
@@ -4080,6 +4096,7 @@ int main(int argc, char **argv) {
          185},
         {SHRPX_OPT_FRONTEND_QUIC_SECRET_FILE.c_str(), required_argument, &flag,
          186},
+        {SHRPX_OPT_RLIMIT_MEMLOCK.c_str(), required_argument, &flag, 187},
         {nullptr, 0, nullptr, 0}};
 
     int option_index = 0;
@@ -4967,6 +4984,10 @@ int main(int argc, char **argv) {
         cmdcfgs.emplace_back(SHRPX_OPT_FRONTEND_QUIC_SECRET_FILE,
                              StringRef{optarg});
         break;
+      case 187:
+        // --rlimit-memlock
+        cmdcfgs.emplace_back(SHRPX_OPT_RLIMIT_MEMLOCK, StringRef{optarg});
+        break;
       default:
         break;
       }
diff --git a/src/shrpx_config.cc b/src/shrpx_config.cc
index 64c401ed..21a2395c 100644
--- a/src/shrpx_config.cc
+++ b/src/shrpx_config.cc
@@ -2057,6 +2057,11 @@ int option_lookup_token(const char *name, size_t namelen) {
         return SHRPX_OPTID_NO_SERVER_PUSH;
       }
       break;
+    case 'k':
+      if (util::strieq_l("rlimit-memloc", name, 13)) {
+        return SHRPX_OPTID_RLIMIT_MEMLOCK;
+      }
+      break;
     case 'p':
       if (util::strieq_l("no-verify-ocs", name, 13)) {
         return SHRPX_OPTID_NO_VERIFY_OCSP;
@@ -4110,6 +4115,23 @@ int parse_config(Config *config, int optid, const StringRef &opt,
 #endif // ENABLE_HTTP3
 
     return 0;
+  case SHRPX_OPTID_RLIMIT_MEMLOCK: {
+    int n;
+
+    if (parse_uint(&n, opt, optarg) != 0) {
+      return -1;
+    }
+
+    if (n < 0) {
+      LOG(ERROR) << opt << ": specify the integer more than or equal to 0";
+
+      return -1;
+    }
+
+    config->rlimit_memlock = n;
+
+    return 0;
+  }
   case SHRPX_OPTID_CONF:
     LOG(WARN) << "conf: ignored";
 
diff --git a/src/shrpx_config.h b/src/shrpx_config.h
index 0c276ddd..d9c6fcef 100644
--- a/src/shrpx_config.h
+++ b/src/shrpx_config.h
@@ -395,6 +395,7 @@ constexpr auto SHRPX_OPT_FRONTEND_QUIC_SERVER_ID =
     StringRef::from_lit("frontend-quic-server-id");
 constexpr auto SHRPX_OPT_FRONTEND_QUIC_SECRET_FILE =
     StringRef::from_lit("frontend-quic-secret-file");
+constexpr auto SHRPX_OPT_RLIMIT_MEMLOCK = StringRef::from_lit("rlimit-memlock");
 
 constexpr size_t SHRPX_OBFUSCATED_NODE_LENGTH = 8;
 
@@ -1064,6 +1065,7 @@ struct Config {
         num_worker{0},
         padding{0},
         rlimit_nofile{0},
+        rlimit_memlock{0},
         uid{0},
         gid{0},
         pid{0},
@@ -1112,6 +1114,7 @@ struct Config {
   size_t num_worker;
   size_t padding;
   size_t rlimit_nofile;
+  size_t rlimit_memlock;
   uid_t uid;
   gid_t gid;
   pid_t pid;
@@ -1281,6 +1284,7 @@ enum {
   SHRPX_OPTID_REDIRECT_HTTPS_PORT,
   SHRPX_OPTID_REQUEST_HEADER_FIELD_BUFFER,
   SHRPX_OPTID_RESPONSE_HEADER_FIELD_BUFFER,
+  SHRPX_OPTID_RLIMIT_MEMLOCK,
   SHRPX_OPTID_RLIMIT_NOFILE,
   SHRPX_OPTID_SERVER_NAME,
   SHRPX_OPTID_SINGLE_PROCESS,