From eda95b0bdb3acc1ad96ee37b469c77d3278e9768 Mon Sep 17 00:00:00 2001
From: Sokwhan Huh <sokwhan@google.com>
Date: Thu, 4 Apr 2024 14:29:32 -0700
Subject: [PATCH] Fix optional field selection behavior on structs

PiperOrigin-RevId: 621971615
---
 .../extensions/CelOptionalLibraryTest.java    | 25 +++++++++++++++++
 .../runtime/DescriptorMessageProvider.java    | 28 +++++++++++--------
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java b/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java
index be215b645..12eab47c4 100644
--- a/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java
+++ b/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java
@@ -541,6 +541,31 @@ public void optionalFieldSelection_onProtoMessage_returnsOptionalValue() throws
     assertThat(result).hasValue(5L);
   }
 
+  @Test
+  public void optionalFieldSelection_onProtoMessage_listValue() throws Exception {
+    Cel cel = newCelBuilder().build();
+    CelAbstractSyntaxTree ast =
+        cel.compile("optional.of(TestAllTypes{repeated_string: ['foo']}).?repeated_string.value()")
+            .getAst();
+
+    List<String> result = (List<String>) cel.createProgram(ast).eval();
+
+    assertThat(result).containsExactly("foo");
+  }
+
+  @Test
+  public void optionalFieldSelection_onProtoMessage_indexValue() throws Exception {
+    Cel cel = newCelBuilder().build();
+    CelAbstractSyntaxTree ast =
+        cel.compile(
+                "optional.of(TestAllTypes{repeated_string: ['foo']}).?repeated_string[0].value()")
+            .getAst();
+
+    String result = (String) cel.createProgram(ast).eval();
+
+    assertThat(result).isEqualTo("foo");
+  }
+
   @Test
   public void optionalFieldSelection_onProtoMessage_chainedSuccess() throws Exception {
     Cel cel =
diff --git a/runtime/src/main/java/dev/cel/runtime/DescriptorMessageProvider.java b/runtime/src/main/java/dev/cel/runtime/DescriptorMessageProvider.java
index e6782a206..e8f87f118 100644
--- a/runtime/src/main/java/dev/cel/runtime/DescriptorMessageProvider.java
+++ b/runtime/src/main/java/dev/cel/runtime/DescriptorMessageProvider.java
@@ -140,28 +140,32 @@ public Object createMessage(String messageName, Map<String, Object> values) {
   @Nullable
   @SuppressWarnings("unchecked")
   public Object selectField(Object message, String fieldName) {
+    boolean isOptionalMessage = false;
     if (message instanceof Optional) {
-      Optional<Map<?, ?>> optionalMap = (Optional<Map<?, ?>>) message;
-      if (!optionalMap.isPresent()) {
-        return Optional.empty();
-      }
-
-      Map<?, ?> unwrappedMap = optionalMap.get();
-      if (!unwrappedMap.containsKey(fieldName)) {
+      isOptionalMessage = true;
+      Optional<Object> optionalMessage = (Optional<Object>) message;
+      if (!optionalMessage.isPresent()) {
         return Optional.empty();
       }
 
-      return Optional.of(unwrappedMap.get(fieldName));
+      message = optionalMessage.get();
     }
 
     if (message instanceof Map) {
       Map<?, ?> map = (Map<?, ?>) message;
       if (map.containsKey(fieldName)) {
-        return map.get(fieldName);
+        Object mapValue = map.get(fieldName);
+        return isOptionalMessage ? Optional.of(mapValue) : mapValue;
+      }
+
+      if (isOptionalMessage) {
+        return Optional.empty();
+      } else {
+        throw new CelRuntimeException(
+            new IllegalArgumentException(
+                String.format("key '%s' is not present in map.", fieldName)),
+            CelErrorCode.ATTRIBUTE_NOT_FOUND);
       }
-      throw new CelRuntimeException(
-          new IllegalArgumentException(String.format("key '%s' is not present in map.", fieldName)),
-          CelErrorCode.ATTRIBUTE_NOT_FOUND);
     }
 
     MessageOrBuilder typedMessage = assertFullProtoMessage(message);