From a9cd4f3ca4b92467e6a8bde51833122ee047ea21 Mon Sep 17 00:00:00 2001
From: Sokwhan Huh <sokwhan@google.com>
Date: Tue, 6 May 2025 15:27:36 -0700
Subject: [PATCH] Assert the key type for JSON struct to be a string

PiperOrigin-RevId: 755547701
---
 .../common/values/BaseProtoCelValueConverter.java   | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/common/src/main/java/dev/cel/common/values/BaseProtoCelValueConverter.java b/common/src/main/java/dev/cel/common/values/BaseProtoCelValueConverter.java
index 3310b3dad..ef45a622b 100644
--- a/common/src/main/java/dev/cel/common/values/BaseProtoCelValueConverter.java
+++ b/common/src/main/java/dev/cel/common/values/BaseProtoCelValueConverter.java
@@ -158,13 +158,20 @@ private ListValue<CelValue> adaptJsonListToCelValue(com.google.protobuf.ListValu
             .collect(toImmutableList()));
   }
 
-  // TODO: Investigate changing MapValue key to StringValue
-  private MapValue<CelValue, CelValue> adaptJsonStructToCelValue(Struct struct) {
+  private MapValue<dev.cel.common.values.StringValue, CelValue> adaptJsonStructToCelValue(
+      Struct struct) {
     return ImmutableMapValue.create(
         struct.getFieldsMap().entrySet().stream()
             .collect(
                 toImmutableMap(
-                    e -> fromJavaObjectToCelValue(e.getKey()),
+                    e -> {
+                      CelValue key = fromJavaObjectToCelValue(e.getKey());
+                      if (!(key instanceof dev.cel.common.values.StringValue)) {
+                        throw new IllegalStateException(
+                            "Expected a string type for the key, but instead got: " + key);
+                      }
+                      return (dev.cel.common.values.StringValue) key;
+                    },
                     e -> adaptJsonValueToCelValue(e.getValue()))));
   }