add support for overload-ignorelist (#237)

Snapchat · Nov 17, 2023 · c889505 · c889505
1 parent 0b5b0c0
commit c889505
Show file tree

Hide file tree

Showing 5 changed files with 112 additions and 39 deletions.
diff --git a/src/config.cpp b/src/config.cpp
@@ -779,6 +779,15 @@ void loadServerConfigFromString(char *config) {
             }
             for (int i = 1; i < argc; i++)
                 g_pserver->tls_auditlog_blocklist.emplace(argv[i], strlen(argv[i]));
+        } else if (!strcasecmp(argv[0], "overload-ignorelist")) {
+            if (argc < 2) {
+                err = "must supply at least one element in the ignore list"; goto loaderr;
+            }
+            if (!g_pserver->overload_ignorelist.empty()) {
+                err = "overload-ignorelist may only be set once"; goto loaderr;
+            }
+            for (int i = 1; i < argc; i++)
+                g_pserver->overload_ignorelist.emplace(argv[i], strlen(argv[i]));
         } else if (!strcasecmp(argv[0], "version-override") && argc == 2) {
             KEYDB_SET_VERSION = zstrdup(argv[1]);
             serverLog(LL_WARNING, "Warning version is overriden to: %s\n", KEYDB_SET_VERSION);
@@ -908,8 +917,8 @@ void configSetCommand(client *c) {
 
     if (c->argc < 4 || c->argc > 4) {
         o = nullptr;
-        // Variadic set is only supported for tls-allowlist
-        if (strcasecmp(szFromObj(c->argv[2]), "tls-allowlist")) {
+        // Variadic set is only supported for tls-allowlist, tls-auditlog-blocklist and overload-ignorelist
+        if (strcasecmp(szFromObj(c->argv[2]), "tls-allowlist") && strcasecmp(szFromObj(c->argv[2]), "tls-auditlog-blocklist") && strcasecmp(szFromObj(c->argv[2]), "overload-ignorelist")) {
             addReplySubcommandSyntaxError(c);
             return;
         }
@@ -1084,6 +1093,18 @@ void configSetCommand(client *c) {
             robj *val = c->argv[i];
             g_pserver->tls_allowlist.emplace(szFromObj(val), sdslen(szFromObj(val)));
         }
+    } config_set_special_field("tls-auditlog-blocklist") {
+        g_pserver->tls_auditlog_blocklist.clear();
+        for (int i = 3; i < c->argc; ++i) {
+            robj *val = c->argv[i];
+            g_pserver->tls_auditlog_blocklist.emplace(szFromObj(val), sdslen(szFromObj(val)));
+        }
+    } config_set_special_field("overload-ignorelist") {
+        g_pserver->overload_ignorelist.clear();
+        for (int i = 3; i < c->argc; ++i) {
+            robj *val = c->argv[i];
+            g_pserver->overload_ignorelist.emplace(szFromObj(val), sdslen(szFromObj(val)));
+        }
     /* Everything else is an error... */
     } config_set_else {
         addReplyErrorFormat(c,"Unsupported CONFIG parameter: %s",
@@ -1295,6 +1316,22 @@ void configGetCommand(client *c) {
         }
         matches++;
     }
+    if (stringmatch(pattern, "tls-auditlog-blocklist", 1)) {
+        addReplyBulkCString(c,"tls-auditlog-blocklist");
+        addReplyArrayLen(c, (long)g_pserver->tls_auditlog_blocklist.size());
+        for (auto &elem : g_pserver->tls_auditlog_blocklist) {
+            addReplyBulkCBuffer(c, elem.get(), elem.size()); // addReplyBulkSds will free which we absolutely don't want
+        }
+        matches++;
+    }
+    if (stringmatch(pattern, "overload-ignorelist", 1)) {
+        addReplyBulkCString(c,"overload-ignorelist");
+        addReplyArrayLen(c, (long)g_pserver->overload_ignorelist.size());
+        for (auto &elem : g_pserver->overload_ignorelist) {
+            addReplyBulkCBuffer(c, elem.get(), elem.size()); // addReplyBulkSds will free which we absolutely don't want
+        }
+        matches++;
+    }
 
     setDeferredMapLen(c,replylen,matches);
 }
@@ -1978,6 +2015,34 @@ int rewriteConfig(char *path, int force_all) {
         rewriteConfigMarkAsProcessed(state, "tls-allowlist"); // ensure the line is removed if it existed
     }
 
+    if (!g_pserver->tls_auditlog_blocklist.empty()) {
+        sds conf = sdsnew("tls-auditlog-blocklist ");
+        for (auto &elem : g_pserver->tls_auditlog_blocklist) {
+            conf = sdscatsds(conf, (sds)elem.get());
+            conf = sdscat(conf, " ");
+        }
+        // trim the trailing space
+        sdsrange(conf, 0, -1);
+        rewriteConfigRewriteLine(state,"tls-auditlog-blocklist",conf,1 /*force*/);
+        // note: conf is owned by rewriteConfigRewriteLine - no need to free
+    } else {
+        rewriteConfigMarkAsProcessed(state, "tls-auditlog-blocklist"); // ensure the line is removed if it existed
+    }
+
+    if (!g_pserver->overload_ignorelist.empty()) {
+        sds conf = sdsnew("overload-ignorelist ");
+        for (auto &elem : g_pserver->overload_ignorelist) {
+            conf = sdscatsds(conf, (sds)elem.get());
+            conf = sdscat(conf, " ");
+        }
+        // trim the trailing space
+        sdsrange(conf, 0, -1);
+        rewriteConfigRewriteLine(state,"overload-ignorelist",conf,1 /*force*/);
+        // note: conf is owned by rewriteConfigRewriteLine - no need to free
+    } else {
+        rewriteConfigMarkAsProcessed(state, "overload-ignorelist"); // ensure the line is removed if it existed
+    }
+
     /* Rewrite Sentinel config if in Sentinel mode. */
     if (g_pserver->sentinel_mode) rewriteConfigSentinelOption(state);
 

diff --git a/src/networking.cpp b/src/networking.cpp
@@ -1193,47 +1193,51 @@ void clientAcceptHandler(connection *conn) {
     if (cserver.fThreadAffinity)
         connSetThreadAffinity(conn, c->iel);
 
+    char cip[NET_IP_STR_LEN+1] = { 0 };
+    connPeerToString(conn, cip, sizeof(cip)-1, NULL);
+
     /* If the server is running in protected mode (the default) and there
      * is no password set, nor a specific interface is bound, we don't accept
      * requests from non loopback interfaces. Instead we try to explain the
      * user what to do to fix it if needed. */
     if (g_pserver->protected_mode &&
         g_pserver->bindaddr_count == 0 &&
         DefaultUser->flags & USER_FLAG_NOPASS &&
-        !(c->flags & CLIENT_UNIX_SOCKET))
+        !(c->flags & CLIENT_UNIX_SOCKET) &&
+        strcmp(cip,"127.0.0.1") && strcmp(cip,"::1"))
     {
-        char cip[NET_IP_STR_LEN+1] = { 0 };
-        connPeerToString(conn, cip, sizeof(cip)-1, NULL);
-
-        if (strcmp(cip,"127.0.0.1") && strcmp(cip,"::1")) {
-            const char *err =
-                "-DENIED KeyDB is running in protected mode because protected "
-                "mode is enabled, no bind address was specified, no "
-                "authentication password is requested to clients. In this mode "
-                "connections are only accepted from the loopback interface. "
-                "If you want to connect from external computers to KeyDB you "
-                "may adopt one of the following solutions: "
-                "1) Just disable protected mode sending the command "
-                "'CONFIG SET protected-mode no' from the loopback interface "
-                "by connecting to KeyDB from the same host the server is "
-                "running, however MAKE SURE KeyDB is not publicly accessible "
-                "from internet if you do so. Use CONFIG REWRITE to make this "
-                "change permanent. "
-                "2) Alternatively you can just disable the protected mode by "
-                "editing the KeyDB configuration file, and setting the protected "
-                "mode option to 'no', and then restarting the server "
-                "3) If you started the server manually just for testing, restart "
-                "it with the '--protected-mode no' option. "
-                "4) Setup a bind address or an authentication password. "
-                "NOTE: You only need to do one of the above things in order for "
-                "the server to start accepting connections from the outside.\r\n";
-            if (connWrite(c->conn,err,strlen(err)) == -1) {
-                /* Nothing to do, Just to avoid the warning... */
-            }
-            g_pserver->stat_rejected_conn++;
-            freeClientAsync(c);
-            return;
+        const char *err =
+            "-DENIED KeyDB is running in protected mode because protected "
+            "mode is enabled, no bind address was specified, no "
+            "authentication password is requested to clients. In this mode "
+            "connections are only accepted from the loopback interface. "
+            "If you want to connect from external computers to KeyDB you "
+            "may adopt one of the following solutions: "
+            "1) Just disable protected mode sending the command "
+            "'CONFIG SET protected-mode no' from the loopback interface "
+            "by connecting to KeyDB from the same host the server is "
+            "running, however MAKE SURE KeyDB is not publicly accessible "
+            "from internet if you do so. Use CONFIG REWRITE to make this "
+            "change permanent. "
+            "2) Alternatively you can just disable the protected mode by "
+            "editing the KeyDB configuration file, and setting the protected "
+            "mode option to 'no', and then restarting the server "
+            "3) If you started the server manually just for testing, restart "
+            "it with the '--protected-mode no' option. "
+            "4) Setup a bind address or an authentication password. "
+            "NOTE: You only need to do one of the above things in order for "
+            "the server to start accepting connections from the outside.\r\n";
+        if (connWrite(c->conn,err,strlen(err)) == -1) {
+            /* Nothing to do, Just to avoid the warning... */
         }
+        g_pserver->stat_rejected_conn++;
+        freeClientAsync(c);
+        return;
+    }
+
+    if (g_pserver->overload_ignorelist.find(sdsstring(sdsnew(cip))) != g_pserver->overload_ignorelist.end())
+    {
+        c->flags |= CLIENT_IGNORE_OVERLOAD;
     }
 
     g_pserver->stat_numconnections++;

diff --git a/src/server.cpp b/src/server.cpp
@@ -1965,7 +1965,7 @@ int closeClientOnOverload(client *c) {
     if (g_pserver->overload_closed_clients > MAX_CLIENTS_SHED_PER_PERIOD) return false;
     if (!g_pserver->is_overloaded) return false;
     // Don't close masters, replicas, or pub/sub clients
-    if (c->flags & (CLIENT_MASTER | CLIENT_SLAVE | CLIENT_PENDING_WRITE | CLIENT_PUBSUB | CLIENT_BLOCKED)) return false;
+    if (c->flags & (CLIENT_MASTER | CLIENT_SLAVE | CLIENT_PENDING_WRITE | CLIENT_PUBSUB | CLIENT_BLOCKED | CLIENT_IGNORE_OVERLOAD)) return false;
     freeClient(c);
     ++g_pserver->overload_closed_clients;
     return true;

diff --git a/src/server.h b/src/server.h
@@ -546,6 +546,7 @@ extern int configOOMScoreAdjValuesDefaults[CONFIG_OOM_COUNT];
                                           RDB without replication buffer. */
 #define CLIENT_FORCE_REPLY (1ULL<<44) /* Should addReply be forced to write the text? */
 #define CLIENT_AUDIT_LOGGING (1ULL<<45) /* Client commands required audit logging */
+#define CLIENT_IGNORE_OVERLOAD (1ULL<<46) /* Client that should not be disconnected by overload protection */
 
 /* Client block type (btype field in client structure)
  * if CLIENT_BLOCKED flag is set. */
@@ -2714,6 +2715,7 @@ struct redisServer {
 
     std::set<sdsstring> tls_auditlog_blocklist; /* Certificates that can be excluded from audit logging */
     std::set<sdsstring> tls_allowlist;
+    std::set<sdsstring> overload_ignorelist;
     redisTLSContextConfig tls_ctx_config;
 
     /* cpu affinity */

diff --git a/tests/unit/introspection.tcl b/tests/unit/introspection.tcl
@@ -162,15 +162,17 @@ start_server {tags {"introspection"}} {
             aof_rewrite_cpulist
             time-thread-priority
             bgsave_cpulist
-	    storage-cache-mode
-	    storage-provider-options
-	    use-fork
+            storage-cache-mode
+            storage-provider-options
+            use-fork
             multi-master
             active-replica
             bind
             set-proc-title
             repl-backlog-disk-reserve
-	    tls-allowlist
+            tls-allowlist
+            tls-auditlog-blocklist
+            overload-ignorelist
             db-s3-object
         }