YARN-7455. quote_and_append_arg can overflow buffer. Contributed by Jim Brennan

(cherry picked from commit 60f95fb719f00a718b484c36a823ec5aa8b099f4)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/native/container-executor/impl/util.c

@@ -21,6 +21,7 @@
 #include <string.h>
 #include <ctype.h>
 #include <regex.h>
+#include <stdio.h>
 
 char** split_delimiter(char *value, const char *delim) {
   char **return_values = NULL;
@@ -176,17 +177,19 @@ char* escape_single_quote(const char *str) {
 
 void quote_and_append_arg(char **str, size_t *size, const char* param, const char *arg) {
   char *tmp = escape_single_quote(arg);
-  int alloc_block = 1024;
-  strcat(*str, param);
-  strcat(*str, "'");
-  if (strlen(*str) + strlen(tmp) > *size) {
-    *size = (strlen(*str) + strlen(tmp) + alloc_block) * sizeof(char);
-    *str = (char *) realloc(*str, *size);
-    if (*str == NULL) {
-      exit(OUT_OF_MEMORY);
-    }
+  const char *append_format = "%s'%s' ";
+  size_t append_size = snprintf(NULL, 0, append_format, param, tmp);
+  append_size += 1;   // for the terminating NUL
+  size_t len_str = strlen(*str);
+  size_t new_size = len_str + append_size;
+  if (new_size > *size) {
+      *size = new_size + QUOTE_AND_APPEND_ARG_GROWTH;
+      *str = (char *) realloc(*str, *size);
+      if (*str == NULL) {
+          exit(OUT_OF_MEMORY);
+      }
   }
-  strcat(*str, tmp);
-  strcat(*str, "' ");
+  char *cur_ptr = *str + len_str;
+  sprintf(cur_ptr, append_format, param, tmp);
   free(tmp);
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/native/container-executor/impl/util.h

@@ -141,12 +141,13 @@ char* escape_single_quote(const char *str);
 
 /**
  * Helper function to quote the argument to a parameter and then append it to the provided string.
- * @param str Buffer to which the param='arg' string must be appended
+ * @param str Buffer to which the param'arg' string must be appended
  * @param size Size of the buffer
  * @param param Param name
  * @param arg Argument to be quoted
  */
 void quote_and_append_arg(char **str, size_t *size, const char* param, const char *arg);
+#define QUOTE_AND_APPEND_ARG_GROWTH (1024) // how much to increase str buffer by if needed
 
 /**
  * Helper function to allocate and clear a block of memory. It'll call exit if the allocation fails.

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/native/container-executor/test/test_util.cc

@@ -149,25 +149,155 @@ namespace ContainerExecutor {
     }
   }
 
+  /**
+   * Internal function for testing quote_and_append_arg()
+   */
+  void test_quote_and_append_arg_function_internal(char **str, size_t *size, const char* param, const char *arg, const char *expected_result) {
+    const size_t alloc_block_size = QUOTE_AND_APPEND_ARG_GROWTH;
+    size_t orig_size = *size;
+    size_t expected_size = strlen(expected_result) + 1;
+    if (expected_size > orig_size) {
+      expected_size += alloc_block_size;
+    } else {
+      expected_size = orig_size; // fits in original string
+    }
+    quote_and_append_arg(str, size, param, arg);
+    ASSERT_STREQ(*str, expected_result);
+    ASSERT_EQ(*size, expected_size);
+    return;
+  }
+
   TEST_F(TestUtil, test_quote_and_append_arg) {
+    size_t str_real_size = 32;
+
+    // Simple test - size = 32, result = 16
+    size_t str_size = str_real_size;
+    char *str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    strcpy(str, "ssss");
+    const char *param = "pppp";
+    const char *arg = "aaaa";
+    const char *expected_result = "sssspppp'aaaa' ";
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
+
+    // Original test - size = 32, result = 19
+    str_size = str_real_size;
+    str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    param = "param=";
+    arg = "argument1";
+    expected_result = "param='argument1' ";
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
+
+    // Original test - size = 32 and result = 21
+    str_size = str_real_size;
+    str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    param = "param=";
+    arg = "ab'cd";
+    expected_result = "param='ab'\"'\"'cd' "; // 21 characters
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
+
+    // Lie about size of buffer so we don't crash from an actual buffer overflow
+    // Original Test - Size = 4 and result = 19
+    str_size = 4;
+    str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    param = "param=";
+    arg = "argument1";
+    expected_result = "param='argument1' "; // 19 characters
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
+
+    // Size = 8 and result = 7
+    str_size = 8;
+    str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    strcpy(str, "s");
+    param = "p";
+    arg = "a";
+    expected_result = "sp'a' "; // 7 characters
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
+
+    // Size = 8 and result = 7
+    str_size = 8;
+    str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    strcpy(str, "s");
+    param = "p";
+    arg = "a";
+    expected_result = "sp'a' "; // 7 characters
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
+
+    // Size = 8 and result = 8
+    str_size = 8;
+    str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    strcpy(str, "ss");
+    param = "p";
+    arg = "a";
+    expected_result = "ssp'a' "; // 8 characters
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
+
+    // size = 8, result = 9 (should grow buffer)
+    str_size = 8;
+    str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    strcpy(str, "ss");
+    param = "pp";
+    arg = "a";
+    expected_result = "sspp'a' "; // 9 characters
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
 
-    char *tmp = static_cast<char *>(malloc(4096));
-    size_t tmp_size = 4096;
+    // size = 8, result = 10 (should grow buffer)
+    str_size = 8;
+    str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    strcpy(str, "ss");
+    param = "pp";
+    arg = "aa";
+    expected_result = "sspp'aa' "; // 10 characters
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
 
-    memset(tmp, 0, tmp_size);
-    quote_and_append_arg(&tmp, &tmp_size, "param=", "argument1");
-    ASSERT_STREQ("param='argument1' ", tmp);
+    // size = 8, result = 11 (should grow buffer)
+    str_size = 8;
+    str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    strcpy(str, "sss");
+    param = "pp";
+    arg = "aa";
+    expected_result = "ssspp'aa' "; // 11 characters
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
 
-    memset(tmp, 0, tmp_size);
-    quote_and_append_arg(&tmp, &tmp_size, "param=", "ab'cd");
-    ASSERT_STREQ("param='ab'\"'\"'cd' ", tmp);
-    free(tmp);
+    // One with quotes - size = 32, result = 17
+    str_size = 32;
+    str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    strcpy(str, "s");
+    param = "p";
+    arg = "'a'";
+    expected_result = "sp''\"'\"'a'\"'\"'' "; // 17 characters
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
 
-    tmp = static_cast<char *>(malloc(4));
-    tmp_size = 4;
-    memset(tmp, 0, tmp_size);
-    quote_and_append_arg(&tmp, &tmp_size, "param=", "argument1");
-    ASSERT_STREQ("param='argument1' ", tmp);
-    ASSERT_EQ(1040, tmp_size);
+    // One with quotes - size = 16, result = 17
+    str_size = 16;
+    str = (char *) malloc(str_real_size);
+    memset(str, 0, str_real_size);
+    strcpy(str, "s");
+    param = "p";
+    arg = "'a'";
+    expected_result = "sp''\"'\"'a'\"'\"'' "; // 17 characters
+    test_quote_and_append_arg_function_internal(&str, &str_size, param, arg, expected_result);
+    free(str);
   }
 }