INTERNAL: Add max_stats_prefixes limit to stats prefix command

engines/default/default_engine.c

@@ -160,6 +160,13 @@ static int check_configuration(struct engine_config *conf)
                 conf->scrub_count, MINIMUM_SCRUB_COUNT, MAXIMUM_SCRUB_COUNT);
         return -1;
     }
+    if (conf->max_stats_prefixes < 0 ||
+        conf->max_stats_prefixes > UINT32_MAX) {
+        logger->log(EXTENSION_LOG_WARNING, NULL,
+                "default engine: max_stats_prefixes(%u) is out of range(%u~%u).\n",
+                conf->max_stats_prefixes, 0, UINT32_MAX);
+        return -1;
+    }
 #ifdef ENABLE_PERSISTENCE
     if (conf->use_persistence) {
         /* check data & logs directory path. */
@@ -214,6 +221,7 @@ initialize_configuration(struct default_engine *se, const char *cfg_str)
         { .key = "max_btree_size",    .datatype = DT_UINT32, .value.dt_uint32 = &se->config.max_btree_size },
         { .key = "max_element_bytes", .datatype = DT_UINT32, .value.dt_uint32 = &se->config.max_element_bytes },
         { .key = "scrub_count",       .datatype = DT_UINT32, .value.dt_uint32 = &se->config.scrub_count},
+        { .key = "max_stats_prefixes",.datatype = DT_UINT32, .value.dt_uint32 = &se->config.max_stats_prefixes},
 #ifdef ENABLE_PERSISTENCE
         { .key = "use_persistence",   .datatype = DT_BOOL,   .value.dt_bool = &se->config.use_persistence },
         { .key = "data_path",         .datatype = DT_STRING, .value.dt_string = &se->config.data_path },
@@ -1489,6 +1497,17 @@ default_set_config(ENGINE_HANDLE* handle, const void* cookie,
         }
         pthread_mutex_unlock(&engine->cache_lock);
     }
+    else if (strcmp(config_key, "max_stats_prefixes") == 0) {
+        uint32_t new_maxstatsprefixes = *(uint32_t*)config_value;
+        pthread_mutex_lock(&engine->cache_lock);
+        if (new_maxstatsprefixes >= 0 &&
+            new_maxstatsprefixes <= UINT32_MAX) {
+            engine->config.max_stats_prefixes = new_maxstatsprefixes;
+        } else {
+            ret = ENGINE_EBADVALUE;
+        }
+        pthread_mutex_unlock(&engine->cache_lock);
+    }
     else if (strcmp(config_key, "verbosity") == 0) {
         pthread_mutex_lock(&engine->cache_lock);
         engine->config.verbose = *(size_t*)config_value;
@@ -1580,6 +1599,11 @@ default_get_config(ENGINE_HANDLE* handle, const void* cookie,
         *(uint32_t*)config_value = engine->config.scrub_count;
         pthread_mutex_unlock(&engine->cache_lock);
     }
+    else if (strcmp(config_key, "max_stats_prefixes") == 0) {
+        pthread_mutex_lock(&engine->cache_lock);
+        *(uint32_t*)config_value = engine->config.max_stats_prefixes;
+        pthread_mutex_unlock(&engine->cache_lock);
+    }
     else if (strcmp(config_key, "verbosity") == 0) {
         pthread_mutex_lock(&engine->cache_lock);
         *(size_t*)config_value = engine->config.verbose;
@@ -2063,6 +2087,7 @@ create_instance(uint64_t interface, GET_SERVER_API get_server_api,
          .max_btree_size = DEFAULT_MAX_BTREE_SIZE,
          .max_element_bytes = DEFAULT_MAX_ELEMENT_BYTES,
          .scrub_count = DEFAULT_SCRUB_COUNT,
+         .max_stats_prefixes = DEFAULT_MAX_STATS_PREFIXES,
 #ifdef ENABLE_PERSISTENCE
          .use_persistence = false,
          .async_logging = false, /* default, sync logging */

engines/default/default_engine.conf

@@ -18,7 +18,9 @@ max_element_bytes=16KB
 # Scrub count (default: 96, min: 16, max: 320)
 # Count of scrubbing items at each try.
 scrub_count=96
-
+#
+# Max prefixes to get stats (default: 50, min: 0, max: 4294967295)
+max_stats_prefixes=50
 #
 # Persistence configuration
 #

engines/default/default_engine.h

@@ -66,6 +66,7 @@ struct engine_config {
    uint32_t   max_btree_size;
    uint32_t   max_element_bytes;
    uint32_t   scrub_count;
+   uint32_t   max_stats_prefixes;
 #ifdef ENABLE_PERSISTENCE
    bool       use_persistence;
    bool       async_logging;

engines/default/prefix.c

@@ -633,6 +633,12 @@ char *prefix_dump_stats(token_t *tokens, const size_t ntokens, int *length)
             sum_nameleng += tokens[i].length;
     }
 
+    if (num_prefixes > config->max_stats_prefixes) {
+        *length = -2;
+        logger->log(EXTENSION_LOG_WARNING, NULL, "Too many prefixes. Operation blocked.\n");
+        return NULL;
+    }
+
     /* Allocate stats buffer: <length, prefix stats list, tail>.
      * Check the count of "%llu" and "%02d" in the above format string.
      *   - 12 : the count of "%llu" strings.

engines/default/prefix.h

@@ -20,6 +20,8 @@
 
 #include <memcached/types.h>
 
+#define DEFAULT_MAX_STATS_PREFIXES 50
+
 typedef struct _prefix_t {
 #ifdef SCAN_COMMAND
     uint16_t refcount; /* reference count */

memcached.c

@@ -359,6 +359,7 @@ static void settings_init(void)
     settings.max_btree_size = 50000; /* DEFAULT_MAX_BTREE_SIZE */
     settings.max_element_bytes = 16 * 1024; /* DEFAULT_MAX_ELEMENT_BYTES */
     settings.scrub_count = 96; /* DEFAULT_SCRUB_COUNT */
+    settings.max_stats_prefixes = 50; /* DEFAULT_MAX_STATS_PREFIXES */
     settings.topkeys = 0;
     settings.require_sasl = false;
     settings.extensions.logger = get_stderr_logger();
@@ -7838,7 +7839,10 @@ inline static void process_stats_detail(conn *c, const char *command)
             int len;
             char *stats = stats_prefix_dump(NULL, 0, &len);
             if (stats == NULL) {
-                out_string(c, "SERVER_ERROR no more memory");
+                if (len == -2)
+                    out_string(c, "SERVER_ERROR too many prefixes");
+                else
+                    out_string(c, "SERVER_ERROR no more memory");
                 return;
             }
             write_and_free(c, stats, len);
@@ -8148,6 +8152,7 @@ static void process_stats_settings(ADD_STAT add_stats, void *c)
     APPEND_STAT("max_btree_size", "%u", settings.max_btree_size);
     APPEND_STAT("max_element_bytes", "%u", settings.max_element_bytes);
     APPEND_STAT("scrub_count", "%u", settings.scrub_count);
+    APPEND_STAT("max_stats_prefixes", "%u", settings.max_stats_prefixes);
     APPEND_STAT("topkeys", "%d", settings.topkeys);
 #ifdef ENABLE_ZK_INTEGRATION
     APPEND_STAT("hb_timeout", "%u", hb_confs.timeout);
@@ -8243,6 +8248,8 @@ static void process_stats_command(conn *c, token_t *tokens, const size_t ntokens
         if (stats == NULL) {
             if (len == -1)
                 out_string(c, "NOT_SUPPORTED");
+            else if (len == -2)
+                out_string(c, "SERVER_ERROR too many prefixes");
             else
                 out_string(c, "SERVER_ERROR no more memory");
             return;
@@ -8269,6 +8276,8 @@ static void process_stats_command(conn *c, token_t *tokens, const size_t ntokens
         if (stats == NULL) {
             if (len == -1)
                 out_string(c, "NOT_SUPPORTED");
+            else if (len == -2)
+                out_string(c, "SERVER_ERROR too many prefixes");
             else
                 out_string(c, "SERVER_ERROR no more memory");
             return;
@@ -9031,6 +9040,33 @@ static void process_scrubcount_command(conn *c, token_t *tokens, const size_t nt
     }
 }
 
+static void process_maxstatsprefixes_command(conn *c, token_t *tokens, const size_t ntokens)
+{
+    char *config_key = tokens[SUBCOMMAND_TOKEN].value;
+    char *config_val = tokens[SUBCOMMAND_TOKEN+1].value;
+    uint32_t new_max_stats_prefixes;
+
+    if (ntokens == 3) {
+        char buf[50];
+        sprintf(buf, "max_stats_prefixes %u\r\nEND", settings.max_stats_prefixes);
+        out_string(c, buf);
+    } else if (ntokens == 4 && safe_strtoul(config_val, &new_max_stats_prefixes)) {
+        ENGINE_ERROR_CODE ret;
+        LOCK_SETTING();
+        ret = mc_engine.v1->set_config(mc_engine.v0, c, config_key, (void*)&new_max_stats_prefixes);
+        if (ret == ENGINE_SUCCESS) {
+            settings.max_stats_prefixes = new_max_stats_prefixes;
+        }
+        UNLOCK_SETTING();
+        if (ret == ENGINE_SUCCESS)        out_string(c, "END");
+        else if (ret == ENGINE_EBADVALUE) out_string(c, "CLIENT_ERROR bad value");
+        else handle_unexpected_errorcode_ascii(c, __func__, ret);
+    } else {
+        print_invalid_command(c, tokens, ntokens);
+        out_string(c, "CLIENT_ERROR bad command line format");
+    }
+}
+
 static void process_verbosity_command(conn *c, token_t *tokens, const size_t ntokens)
 {
     assert(c != NULL);
@@ -9190,6 +9226,9 @@ static void process_config_command(conn *c, token_t *tokens, const size_t ntoken
     else if (strcmp(config_key, "scrub_count") == 0) {
         process_scrubcount_command(c, tokens, ntokens);
     }
+    else if (strcmp(config_key, "max_stats_prefixes") == 0) {
+        process_maxstatsprefixes_command(c, tokens, ntokens);
+    }
 #ifdef ENABLE_ZK_INTEGRATION
     else if (strcmp(config_key, "zkfailstop") == 0) {
         process_zkfailstop_command(c, tokens, ntokens);
@@ -9471,6 +9510,7 @@ static void process_help_command(conn *c, token_t *tokens, const size_t ntokens)
         "\t" "config max_btree_size [<maxsize>]\\r\\n" "\n"
         "\t" "config max_element_bytes [<maxbytes>]\\r\\n" "\n"
         "\t" "config scrub_count [<count>]\\r\\n" "\n"
+        "\t" "config max_stats_prefixes [<value>]\\r\\n" "\n"
 #ifdef ENABLE_ZK_INTEGRATION
         "\t" "config hbtimeout [<hbtimeout>]\\r\\n" "\n"
         "\t" "config hbfailstop [<hbfailstop>]\\r\\n" "\n"
@@ -15002,6 +15042,7 @@ static void settings_reload_engine_config(void)
     uint32_t maxsize;
     uint32_t maxbytes;
     uint32_t scrubcount;
+    uint32_t maxstatsprefixes;
 
     /* Following settings are loaded by getting engine config */
 
@@ -15029,6 +15070,10 @@ static void settings_reload_engine_config(void)
     if (ret == ENGINE_SUCCESS) {
         settings.scrub_count = scrubcount;
     }
+    ret = mc_engine.v1->get_config(mc_engine.v0, NULL, "max_stats_prefixes", (void*)&maxstatsprefixes);
+    if (ret == ENGINE_SUCCESS) {
+        settings.max_stats_prefixes = maxstatsprefixes;
+    }
 }
 
 static void close_listen_sockets(void)

memcached.h

@@ -244,6 +244,7 @@ struct settings {
     uint32_t max_btree_size;     /* Maximum elements in b+tree collection */
     uint32_t max_element_bytes;  /* Maximum element bytes of collections */
     uint32_t scrub_count;        /* count of scrubbing items at each try */
+    uint32_t max_stats_prefixes; /* Maximum prefixes to view stats */
     int topkeys;            /* Number of top keys to track */
     struct {
         EXTENSION_DAEMON_DESCRIPTOR *daemons;

stats_prefix.c

@@ -1247,6 +1247,12 @@ char *stats_prefix_dump(token_t *tokens, const size_t ntokens, int *length)
         nprefixes = num_prefixes;
     }
 
+    if (nprefixes > settings.max_stats_prefixes) {
+        *length = -2;
+        UNLOCK_STATS();
+        return NULL;
+    }
+
     size = strlen(format) + prefix_name_size +
            nprefixes * (strlen(format) - 2 /* %s */
                            + 54 * (20 - 4)) /* %llu replaced by 20-digit num */

t/flush-prefix.t

@@ -1,7 +1,7 @@
 #!/usr/bin/perl
 
 use strict;
-use Test::More tests => 6001;
+use Test::More tests => 301;
 use FindBin qw($Bin);
 use lib "$Bin/lib";
 use MemcachedTest;
@@ -14,7 +14,7 @@ my $val;
 my $rst;
 my $size;
 my $count;
-my $prefix_size = 1000;
+my $prefix_size = 50;
 
 sub prefix_insert {
   for ($size = 0; $size < $prefix_size; $size++) {