[PATCH 04/13] fixup! KVM: selftests: Add memory size parameter to the demand paging test

[Andrew Jones <drjones@xxxxxxxxxx>]

[Rewrote parse_size() to simplify and provide user more flexibility as
 to how sizes are input. Also fixed size overflow assert.]
Signed-off-by: Andrew Jones <drjones@xxxxxxxxxx>
---
 tools/testing/selftests/kvm/lib/test_util.c | 76 +++++++++------------
 1 file changed, 33 insertions(+), 43 deletions(-)

diff --git a/tools/testing/selftests/kvm/lib/test_util.c b/tools/testing/selftests/kvm/lib/test_util.c
index 706e0f963a44..cbd7f51b07a1 100644
--- a/tools/testing/selftests/kvm/lib/test_util.c
+++ b/tools/testing/selftests/kvm/lib/test_util.c
@@ -4,58 +4,48 @@
  *
  * Copyright (C) 2020, Google LLC.
  */
-
-#include "test_util.h"
-
+#include <stdlib.h>
 #include <ctype.h>
+#include <limits.h>
+#include "test_util.h"
 
 /*
  * Parses "[0-9]+[kmgt]?".
  */
 size_t parse_size(const char *size)
 {
-	size_t len = strlen(size);
-	size_t i;
-	size_t scale_shift = 0;
 	size_t base;
-
-	TEST_ASSERT(len > 0, "Need at least 1 digit in '%s'", size);
-
-	/* Find the first letter in the string, indicating scale. */
-	for (i = 0; i < len; i++) {
-		if (!isdigit(size[i])) {
-			TEST_ASSERT(i > 0, "Need at least 1 digit in '%s'",
-				    size);
-			TEST_ASSERT(i == len - 1,
-				    "Expected letter at the end in '%s'.",
-				    size);
-			switch (tolower(size[i])) {
-			case 't':
-				scale_shift = 40;
-				break;
-			case 'g':
-				scale_shift = 30;
-				break;
-			case 'm':
-				scale_shift = 20;
-				break;
-			case 'k':
-				scale_shift = 10;
-				break;
-			default:
-				TEST_ASSERT(false, "Unknown size letter %c",
-					    size[i]);
-			}
-		}
+	char *scale;
+	int shift = 0;
+
+	TEST_ASSERT(size && isdigit(size[0]), "Need at least one digit in '%s'", size);
+
+	base = strtoull(size, &scale, 0);
+
+	TEST_ASSERT(base != ULLONG_MAX, "Overflow parsing size!");
+
+	switch (tolower(*scale)) {
+	case 't':
+		shift = 40;
+		break;
+	case 'g':
+		shift = 30;
+		break;
+	case 'm':
+		shift = 20;
+		break;
+	case 'k':
+		shift = 10;
+		break;
+	case 'b':
+	case '\0':
+		shift = 0;
+		break;
+	default:
+		TEST_ASSERT(false, "Unknown size letter %c", *scale);
 	}
 
-	TEST_ASSERT(scale_shift < 8 * sizeof(size_t),
-		    "Overflow parsing scale!");
-
-	base = atoi(size);
-
-	TEST_ASSERT(!(base & ~((1 << (sizeof(size_t) - scale_shift)) - 1)),
-	       "Overflow parsing size!");
+	TEST_ASSERT((base << shift) >> shift == base, "Overflow scaling size!");
 
-	return base << scale_shift;
+	return base << shift;
 }
-- 
2.21.1