Prevent planning createMap for heterogeneous duplicate keys

PiperOrigin-RevId: 875399236

Author: l46kok

common/src/main/java/dev/cel/common/exceptions/CelInvalidArgumentException.java

@@ -24,4 +24,8 @@ public final class CelInvalidArgumentException extends CelRuntimeException {
   public CelInvalidArgumentException(Throwable cause) {
     super(cause, CelErrorCode.INVALID_ARGUMENT);
   }
+
+  public CelInvalidArgumentException(String message) {
+    super(message, CelErrorCode.INVALID_ARGUMENT);
+  }
 }

conformance/src/test/java/dev/cel/conformance/BUILD.bazel

@@ -175,9 +175,6 @@ _TESTS_TO_SKIP_PLANNER = [
     "timestamps/timestamp_range/sub_time_duration_under",
 
     # Skip until fixed.
-    "fields/qualified_identifier_resolution/map_key_float",
-    "fields/qualified_identifier_resolution/map_key_null",
-    "fields/qualified_identifier_resolution/map_value_repeat_key_heterogeneous",
     "optionals/optionals/map_null_entry_no_such_key",
     "optionals/optionals/map_present_key_invalid_field",
     "parse/receiver_function_names",

runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel

@@ -333,7 +333,9 @@ java_library(
         ":execution_frame",
         ":localized_evaluation_exception",
         ":planned_interpretable",
+        "//common:error_codes",
         "//common/exceptions:duplicate_key",
+        "//common/exceptions:invalid_argument",
         "//runtime:evaluation_exception",
         "//runtime:interpretable",
         "@maven//:com_google_errorprone_error_prone_annotations",

runtime/src/main/java/dev/cel/runtime/planner/EvalCreateMap.java

@@ -17,8 +17,10 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Sets;
+import com.google.common.primitives.UnsignedLong;
 import com.google.errorprone.annotations.Immutable;
 import dev.cel.common.exceptions.CelDuplicateKeyException;
+import dev.cel.common.exceptions.CelInvalidArgumentException;
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.runtime.GlobalResolver;
 import java.util.HashSet;
@@ -46,13 +48,37 @@ public Object eval(GlobalResolver resolver, ExecutionFrame frame) throws CelEval
     HashSet<Object> keysSeen = Sets.newHashSetWithExpectedSize(keys.length);
 
     for (int i = 0; i < keys.length; i++) {
-      Object key = keys[i].eval(resolver, frame);
-      Object val = values[i].eval(resolver, frame);
+      PlannedInterpretable keyInterpretable = keys[i];
+      Object key = keyInterpretable.eval(resolver, frame);
+      if (!(key instanceof String
+          || key instanceof Long
+          || key instanceof UnsignedLong
+          || key instanceof Boolean)) {
+        throw new LocalizedEvaluationException(
+            new CelInvalidArgumentException("Unsupported key type: " + key),
+            keyInterpretable.exprId());
+      }
 
-      if (!keysSeen.add(key)) {
-        throw new LocalizedEvaluationException(CelDuplicateKeyException.of(key), keys[i].exprId());
+      boolean isDuplicate = !keysSeen.add(key);
+      if (!isDuplicate) {
+        if (key instanceof Long) {
+          long longVal = (Long) key;
+          if (longVal >= 0) {
+            isDuplicate = keysSeen.contains(UnsignedLong.valueOf(longVal));
+          }
+        } else if (key instanceof UnsignedLong) {
+          UnsignedLong ulongVal = (UnsignedLong) key;
+          isDuplicate = keysSeen.contains(ulongVal.longValue());
+        }
       }
 
+      if (isDuplicate) {
+        throw new LocalizedEvaluationException(
+            CelDuplicateKeyException.of(key), keyInterpretable.exprId());
+      }
+
+      Object val = values[i].eval(resolver, frame);
+
       if (isOptional[i]) {
         if (!(val instanceof Optional)) {
           throw new IllegalArgumentException(

runtime/src/test/java/dev/cel/runtime/planner/ProgramPlannerTest.java

@@ -390,6 +390,17 @@ public void plan_createMap_containsDuplicateKey_throws() throws Exception {
         .contains("evaluation error at <input>:20: duplicate map key [true]");
   }
 
+  @Test
+  public void plan_createMap_unsupportedKeyType_throws() throws Exception {
+    CelAbstractSyntaxTree ast = compile("{1.0: 'foo'}");
+    Program program = PLANNER.plan(ast);
+
+    CelEvaluationException e = assertThrows(CelEvaluationException.class, program::eval);
+    assertThat(e)
+        .hasMessageThat()
+        .contains("evaluation error at <input>:1: Unsupported key type: 1.0");
+  }
+
   @Test
   public void plan_createStruct() throws Exception {
     CelAbstractSyntaxTree ast = compile("cel.expr.conformance.proto3.TestAllTypes{}");