units.parse_bytes: add E* and P* units (#3692)

Previous changes to the built-in function to use big.Int et al. had been made without knowledge of these missing units. So, here they are. Fixes #2911. Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
open-policy-agent · Jul 31, 2021 · 8fc2ba1 · 8fc2ba1
1 parent ab3152f
commit 8fc2ba1
Show file tree

Hide file tree

Showing 2 changed files with 127 additions and 18 deletions.
diff --git a/test/cases/testdata/units/test-parse-bytes.yaml b/test/cases/testdata/units/test-parse-bytes.yaml
@@ -356,4 +356,100 @@ cases:
   note: parse_bytes/100 tebibytes as tib
   query: data.test.p = x
   want_result:
+  - x: true
+
+  - data:
+  modules:
+  - |
+    package test
+    p {
+        units.parse_bytes("100p") == 100000000000000000
+    }
+  note: parse_bytes/100 petabytes as p
+  query: data.test.p = x
+  want_result:
+  - x: true
+
+- data:
+  modules:
+  - |
+    package test
+    p {
+        units.parse_bytes("100pb") == 100000000000000000
+    }
+  note: parse_bytes/100 petabytes as pb
+  query: data.test.p = x
+  want_result:
+  - x: true
+
+- data:
+  modules:
+  - |
+    package test
+    p {
+        units.parse_bytes("100pi") == 112589990684262400
+    }
+  note: parse_bytes/100 pebibytes as pi
+  query: data.test.p = x
+  want_result:
+  - x: true
+
+- data:
+  modules:
+  - |
+    package test
+    p {
+        units.parse_bytes("100pib") == 112589990684262400
+    }
+  note: parse_bytes/100 pebibytes as pib
+  query: data.test.p = x
+  want_result:
+  - x: true
+
+  - data:
+  modules:
+  - |
+    package test
+    p {
+        units.parse_bytes("10e") == 10000000000000000000
+    }
+  note: parse_bytes/10 etabytes as e
+  query: data.test.p = x
+  want_result:
+  - x: true
+
+- data:
+  modules:
+  - |
+    package test
+    p {
+        units.parse_bytes("10eb") == 10000000000000000000
+    }
+  note: parse_bytes/10 etabytes as eb
+  query: data.test.p = x
+  want_result:
+  - x: true
+
+- data:
+  modules:
+  - |
+    package test
+    p {
+        units.parse_bytes("10ei") == 11529215046068469760
+    }
+  note: parse_bytes/10 ebibytes as ei
+  query: data.test.p = x
+  want_result:
+  - x: true
+
+- data:
+  modules:
+  - |
+    package test
+    p {
+        units.parse_bytes("10eib") == 11529215046068469760
+    }
+  note: parse_bytes/10 ebibytes as eib
+  query: data.test.p = x
+  want_result:
   - x: true
diff --git a/topdown/parse_bytes.go b/topdown/parse_bytes.go
@@ -15,15 +15,20 @@ import (
 )
 
 const (
-	none int64 = 1
-	kb         = 1000
-	ki         = 1024
-	mb         = kb * 1000
-	mi         = ki * 1024
-	gb         = mb * 1000
-	gi         = mi * 1024
-	tb         = gb * 1000
-	ti         = gi * 1024
+	none uint64 = 1 << (10 * iota)
+	ki
+	mi
+	gi
+	ti
+	pi
+	ei
+
+	kb uint64 = 1000
+	mb        = kb * 1000
+	gb        = mb * 1000
+	tb        = gb * 1000
+	pb        = tb * 1000
+	eb        = pb * 1000
 )
 
 func parseNumBytesError(msg string) error {
@@ -61,23 +66,31 @@ func builtinNumBytes(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.T
 
 	switch unit {
 	case "":
-		m.SetInt64(none)
+		m.SetUint64(none)
 	case "kb", "k":
-		m.SetInt64(kb)
+		m.SetUint64(kb)
 	case "kib", "ki":
-		m.SetInt64(ki)
+		m.SetUint64(ki)
 	case "mb", "m":
-		m.SetInt64(mb)
+		m.SetUint64(mb)
 	case "mib", "mi":
-		m.SetInt64(mi)
+		m.SetUint64(mi)
 	case "gb", "g":
-		m.SetInt64(gb)
+		m.SetUint64(gb)
 	case "gib", "gi":
-		m.SetInt64(gi)
+		m.SetUint64(gi)
 	case "tb", "t":
-		m.SetInt64(tb)
+		m.SetUint64(tb)
 	case "tib", "ti":
-		m.SetInt64(ti)
+		m.SetUint64(ti)
+	case "pb", "p":
+		m.SetUint64(pb)
+	case "pib", "pi":
+		m.SetUint64(pi)
+	case "eb", "e":
+		m.SetUint64(eb)
+	case "eib", "ei":
+		m.SetUint64(ei)
 	default:
 		return errUnitNotRecognized(unit)
 	}