Limit exponent size in parse_bytes/parse_units to prevent timeout bypass (#8524)

Timeouts passed via context are not honored because builtins that
parse exponents dont include context and will run with unrestricted
size on the exponents. Added an exponent limit checker based
on rough benchmakrking data of whats a reasonabile limit.

Fixes: #8326

Signed-off-by: Isaiah Vita <vita.isaiah@gmail.com>

v1/test/cases/testdata/v0/units/test-parse-bytes-errors.yaml

@@ -84,3 +84,15 @@ cases:
     want_error_code: eval_builtin_error
     want_error: "units.parse_bytes: spaces not allowed in resource strings"
     strict_error: true
+  - data:
+    modules:
+      - |
+        package test
+        p {
+            units.parse_bytes("10e10000000EiB")
+        }
+    note: units_parse_bytes/exponent too large
+    query: data.test.p = x
+    want_error_code: eval_builtin_error
+    want_error: "units.parse_bytes: exponent too large"
+    strict_error: true

v1/test/cases/testdata/v0/units/test-parse-units-errors.yaml

@@ -84,3 +84,15 @@ cases:
     want_error_code: eval_builtin_error
     want_error: "units.parse: spaces not allowed in resource strings"
     strict_error: true
+  - data:
+    modules:
+      - |
+        package test
+        p {
+            units.parse("10e10000000Ei")
+        }
+    note: units_parse/exponent too large
+    query: data.test.p = x
+    want_error_code: eval_builtin_error
+    want_error: "units.parse: exponent too large"
+    strict_error: true

v1/test/cases/testdata/v1/units/test-parse-bytes-errors.yaml

@@ -84,3 +84,15 @@ cases:
     want_error_code: eval_builtin_error
     want_error: "units.parse_bytes: spaces not allowed in resource strings"
     strict_error: true
+  - note: units_parse_bytes/exponent too large
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+        	units.parse_bytes("10e10000000EiB")
+        }
+    want_error_code: eval_builtin_error
+    want_error: "units.parse_bytes: exponent too large"
+    strict_error: true

v1/test/cases/testdata/v1/units/test-parse-units-errors.yaml

@@ -84,3 +84,15 @@ cases:
     want_error_code: eval_builtin_error
     want_error: "units.parse: spaces not allowed in resource strings"
     strict_error: true
+  - note: units_parse/exponent too large
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+        	units.parse("10e10000000Ei")
+        }
+    want_error_code: eval_builtin_error
+    want_error: "units.parse: exponent too large"
+    strict_error: true

v1/topdown/parse_bytes.go

@@ -5,6 +5,7 @@
 package topdown
 
 import (
+	"errors"
 	"fmt"
 	"math/big"
 	"strings"
@@ -43,8 +44,13 @@ var (
 	errBytesValueNoAmount       = parseNumBytesError("no byte amount provided")
 	errBytesValueNumConv        = parseNumBytesError("could not parse byte amount to a number")
 	errBytesValueIncludesSpaces = parseNumBytesError("spaces not allowed in resource strings")
+	errBytesExponentTooLarge    = parseNumBytesError("exponent too large")
 )
 
+// maxExponentDigits limits the number of digits allowed in the exponent of
+// scientific notation input.
+const maxExponentDigits = 6
+
 func builtinNumBytes(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
 	var m big.Float
 
@@ -59,7 +65,10 @@ func builtinNumBytes(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term
 		return errBytesValueIncludesSpaces
 	}
 
-	num, unit := extractNumAndUnit(s)
+	num, unit, err := extractNumAndUnit(s)
+	if err != nil {
+		return errBytesExponentTooLarge
+	}
 	if num == "" {
 		return errBytesValueNoAmount
 	}
@@ -115,7 +124,9 @@ func formatString(s ast.String) string {
 // Splits the string into a number string à la "10" or "10.2" and a unit
 // string à la "gb" or "MiB" or "foo". Either can be an empty string
 // (error handling is provided elsewhere).
-func extractNumAndUnit(s string) (string, string) {
+// Returns an error if the exponent in scientific notation exceeds
+// maxExponentDigits digits.
+func extractNumAndUnit(s string) (string, string, error) {
 	isNum := func(r rune) bool {
 		return unicode.IsDigit(r) || r == '.'
 	}
@@ -138,18 +149,28 @@ func extractNumAndUnit(s string) (string, string) {
 			if idx+1 < len(s) && (s[idx+1] == '+' || s[idx+1] == '-') {
 				idx++
 			}
+
+			// Count the digits in the exponent and reject if too large.
+			expStart := idx + 1
+			expEnd := expStart
+			for expEnd < len(s) && unicode.IsDigit(rune(s[expEnd])) {
+				expEnd++
+			}
+			if expEnd-expStart > maxExponentDigits {
+				return "", "", errors.New("exponent too large")
+			}
 		}
 	}
 
 	if firstNonNumIdx == -1 { // only digits, '.', or valid scientific notation
-		return s, ""
+		return s, "", nil
 	}
 	if firstNonNumIdx == 0 { // only units (starts with non-digit)
-		return "", s
+		return "", s, nil
 	}
 
 	// Return the number and the rest as the unit
-	return s[:firstNonNumIdx], s[firstNonNumIdx:]
+	return s[:firstNonNumIdx], s[firstNonNumIdx:], nil
 }
 
 func init() {

v1/topdown/parse_units.go

@@ -34,9 +34,10 @@ func errUnitNotRecognized(unit string) error {
 }
 
 var (
-	errNoAmount       = parseUnitsError("no amount provided")
-	errNumConv        = parseUnitsError("could not parse amount to a number")
-	errIncludesSpaces = parseUnitsError("spaces not allowed in resource strings")
+	errNoAmount              = parseUnitsError("no amount provided")
+	errNumConv               = parseUnitsError("could not parse amount to a number")
+	errIncludesSpaces        = parseUnitsError("spaces not allowed in resource strings")
+	errUnitsExponentTooLarge = parseUnitsError("exponent too large")
 )
 
 // Accepts both normal SI and binary SI units.
@@ -56,7 +57,10 @@ func builtinUnits(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) e
 		return errIncludesSpaces
 	}
 
-	num, unit := extractNumAndUnit(s)
+	num, unit, err := extractNumAndUnit(s)
+	if err != nil {
+		return errUnitsExponentTooLarge
+	}
 	if num == "" {
 		return errNoAmount
 	}