Add support for injecting a custom attribute matcher.

PiperOrigin-RevId: 769263582

Author: jnthntatum

base/attribute_set.h

@@ -15,8 +15,6 @@
 #ifndef THIRD_PARTY_CEL_CPP_BASE_ATTRIBUTE_SET_H_
 #define THIRD_PARTY_CEL_CPP_BASE_ATTRIBUTE_SET_H_
 
-#include <vector>
-
 #include "absl/container/btree_set.h"
 #include "absl/types/span.h"
 #include "base/attribute.h"

eval/eval/BUILD

@@ -48,6 +48,7 @@ cc_library(
         "//runtime",
         "//runtime:activation_interface",
         "//runtime:runtime_options",
+        "//runtime/internal:activation_attribute_matcher_access",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/log:absl_check",
@@ -1026,6 +1027,8 @@ cc_library(
         "//common:value",
         "//eval/internal:errors",
         "//internal:status_macros",
+        "//runtime/internal:attribute_matcher",
+        "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/types:optional",
         "@com_google_absl//absl/types:span",
@@ -1039,6 +1042,7 @@ cc_test(
         "attribute_utility_test.cc",
     ],
     deps = [
+        ":attribute_trail",
         ":attribute_utility",
         "//base:attributes",
         "//common:unknown",
@@ -1048,6 +1052,8 @@ cc_test(
         "//eval/public:unknown_attribute_set",
         "//eval/public:unknown_set",
         "//internal:testing",
+        "//runtime/internal:attribute_matcher",
+        "@com_google_absl//absl/types:span",
         "@com_google_protobuf//:protobuf",
     ],
 )

eval/eval/attribute_utility.cc

@@ -19,9 +19,12 @@
 #include "eval/eval/attribute_trail.h"
 #include "eval/internal/errors.h"
 #include "internal/status_macros.h"
+#include "runtime/internal/attribute_matcher.h"
 
 namespace google::api::expr::runtime {
 
+using ::cel::Attribute;
+using ::cel::AttributePattern;
 using ::cel::AttributeSet;
 using ::cel::Cast;
 using ::cel::ErrorValue;
@@ -31,25 +34,55 @@ using ::cel::InstanceOf;
 using ::cel::UnknownValue;
 using ::cel::Value;
 using ::cel::base_internal::UnknownSet;
+using ::cel::runtime_internal::AttributeMatcher;
 
 using Accumulator = AttributeUtility::Accumulator;
+using MatchResult = AttributeMatcher::MatchResult;
+
+DefaultAttributeMatcher::DefaultAttributeMatcher(
+    absl::Span<const AttributePattern> unknown_patterns,
+    absl::Span<const AttributePattern> missing_patterns)
+    : unknown_patterns_(unknown_patterns),
+      missing_patterns_(missing_patterns) {}
+
+DefaultAttributeMatcher::DefaultAttributeMatcher() = default;
+
+AttributeMatcher::MatchResult MatchAgainstPatterns(
+    absl::Span<const AttributePattern> patterns, const Attribute& attr) {
+  MatchResult result = MatchResult::NONE;
+  for (const auto& pattern : patterns) {
+    auto current_match = pattern.IsMatch(attr);
+    if (current_match == cel::AttributePattern::MatchType::FULL) {
+      return MatchResult::FULL;
+    }
+    if (current_match == cel::AttributePattern::MatchType::PARTIAL) {
+      result = MatchResult::PARTIAL;
+    }
+  }
+  return result;
+}
+
+DefaultAttributeMatcher::MatchResult DefaultAttributeMatcher::CheckForUnknown(
+    const Attribute& attr) const {
+  return MatchAgainstPatterns(unknown_patterns_, attr);
+}
+
+DefaultAttributeMatcher::MatchResult DefaultAttributeMatcher::CheckForMissing(
+    const Attribute& attr) const {
+  return MatchAgainstPatterns(missing_patterns_, attr);
+}
 
 bool AttributeUtility::CheckForMissingAttribute(
     const AttributeTrail& trail) const {
   if (trail.empty()) {
     return false;
   }
-
-  for (const auto& pattern : missing_attribute_patterns_) {
-    // (b/161297249) Preserving existing behavior for now, will add a streamz
-    // for partial match, follow up with tightening up which fields are exposed
-    // to the condition (w/ ajay and jim)
-    if (pattern.IsMatch(trail.attribute()) ==
-        cel::AttributePattern::MatchType::FULL) {
-      return true;
-    }
-  }
-  return false;
+  // Missing attributes are only treated as errors if the attribute exactly
+  // matches (so no guard against passing partial state to a function as with
+  // unknowns). This was initially a design oversight, but is difficult to
+  // change now.
+  return matcher_->CheckForMissing(trail.attribute()) ==
+         AttributeMatcher::MatchResult::FULL;
 }
 
 // Checks whether particular corresponds to any patterns that define unknowns.
@@ -58,13 +91,11 @@ bool AttributeUtility::CheckForUnknown(const AttributeTrail& trail,
   if (trail.empty()) {
     return false;
   }
-  for (const auto& pattern : unknown_patterns_) {
-    auto current_match = pattern.IsMatch(trail.attribute());
-    if (current_match == cel::AttributePattern::MatchType::FULL ||
-        (use_partial &&
-         current_match == cel::AttributePattern::MatchType::PARTIAL)) {
-      return true;
-    }
+  MatchResult result = matcher_->CheckForUnknown(trail.attribute());
+
+  if (result == MatchResult::FULL ||
+      (use_partial && result == MatchResult::PARTIAL)) {
+    return true;
   }
   return false;
 }

eval/eval/attribute_utility.h

@@ -1,17 +1,43 @@
 #ifndef THIRD_PARTY_CEL_CPP_EVAL_EVAL_UNKNOWNS_UTILITY_H_
 #define THIRD_PARTY_CEL_CPP_EVAL_EVAL_UNKNOWNS_UTILITY_H_
 
+#include <cstdint>
+
+#include "absl/base/nullability.h"
 #include "absl/status/statusor.h"
+#include "absl/types/optional.h"
 #include "absl/types/span.h"
 #include "base/attribute.h"
 #include "base/attribute_set.h"
 #include "base/function_result_set.h"
 #include "common/function_descriptor.h"
 #include "common/value.h"
 #include "eval/eval/attribute_trail.h"
+#include "runtime/internal/attribute_matcher.h"
 
 namespace google::api::expr::runtime {
 
+// Default implementation of the attribute matcher.
+// Scans the attribute trail against a list of unknown or missing patterns.
+class DefaultAttributeMatcher : public cel::runtime_internal::AttributeMatcher {
+ private:
+  using MatchResult = cel::runtime_internal::AttributeMatcher::MatchResult;
+
+ public:
+  DefaultAttributeMatcher(
+      absl::Span<const cel::AttributePattern> unknown_patterns,
+      absl::Span<const cel::AttributePattern> missing_patterns);
+
+  DefaultAttributeMatcher();
+
+  MatchResult CheckForUnknown(const cel::Attribute& attr) const override;
+  MatchResult CheckForMissing(const cel::Attribute& attr) const override;
+
+ private:
+  absl::Span<const cel::AttributePattern> unknown_patterns_;
+  absl::Span<const cel::AttributePattern> missing_patterns_;
+};
+
 // Helper class for handling unknowns and missing attribute logic. Provides
 // helpers for merging unknown sets from arguments on the stack and for
 // identifying unknown/missing attributes based on the patterns for a given
@@ -61,11 +87,14 @@ class AttributeUtility {
     bool unknown_present_;
   };
 
-  AttributeUtility(
-      absl::Span<const cel::AttributePattern> unknown_patterns,
-      absl::Span<const cel::AttributePattern> missing_attribute_patterns)
-      : unknown_patterns_(unknown_patterns),
-        missing_attribute_patterns_(missing_attribute_patterns) {}
+  AttributeUtility(absl::Span<const cel::AttributePattern> unknown_patterns,
+                   absl::Span<const cel::AttributePattern> missing_patterns)
+      : default_matcher_(unknown_patterns, missing_patterns),
+        matcher_(&default_matcher_) {}
+
+  explicit AttributeUtility(
+      const cel::runtime_internal::AttributeMatcher* ABSL_NONNULL matcher)
+      : matcher_(matcher) {}
 
   AttributeUtility(const AttributeUtility&) = delete;
   AttributeUtility& operator=(const AttributeUtility&) = delete;
@@ -131,13 +160,18 @@ class AttributeUtility {
     return Accumulator(*this);
   }
 
+  void set_matcher(
+      const cel::runtime_internal::AttributeMatcher* ABSL_NONNULL matcher) {
+    matcher_ = matcher;
+  }
+
  private:
   // Workaround friend visibility.
   void Add(Accumulator& a, const cel::UnknownValue& v) const;
   void Add(Accumulator& a, const AttributeTrail& attr) const;
 
-  absl::Span<const cel::AttributePattern> unknown_patterns_;
-  absl::Span<const cel::AttributePattern> missing_attribute_patterns_;
+  DefaultAttributeMatcher default_matcher_;
+  const cel::runtime_internal::AttributeMatcher* ABSL_NONNULL matcher_;
 };
 
 }  // namespace google::api::expr::runtime

eval/eval/attribute_utility_test.cc

@@ -1,15 +1,20 @@
 #include "eval/eval/attribute_utility.h"
 
+#include <string>
 #include <vector>
 
+#include "absl/types/span.h"
+#include "base/attribute.h"
 #include "base/attribute_set.h"
 #include "common/unknown.h"
 #include "common/value.h"
+#include "eval/eval/attribute_trail.h"
 #include "eval/public/cel_attribute.h"
 #include "eval/public/cel_value.h"
 #include "eval/public/unknown_attribute_set.h"
 #include "eval/public/unknown_set.h"
 #include "internal/testing.h"
+#include "runtime/internal/attribute_matcher.h"
 #include "google/protobuf/arena.h"
 
 namespace google::api::expr::runtime {
@@ -30,6 +35,8 @@ class AttributeUtilityTest : public ::testing::Test {
   google::protobuf::Arena arena_;
 };
 
+absl::Span<const CelAttributePattern> NoPatterns() { return {}; }
+
 TEST_F(AttributeUtilityTest, UnknownsUtilityCheckUnknowns) {
   std::vector<CelAttributePattern> unknown_patterns = {
       CelAttributePattern("unknown0", {CreateCelAttributeQualifierPattern(
@@ -160,4 +167,50 @@ TEST_F(AttributeUtilityTest, CreateUnknownSet) {
   EXPECT_EQ(elem, "destination.ip");
 }
 
+class FakeMatcher : public cel::runtime_internal::AttributeMatcher {
+ private:
+  using MatchResult = cel::runtime_internal::AttributeMatcher::MatchResult;
+
+ public:
+  MatchResult CheckForUnknown(const cel::Attribute& attr) const override {
+    std::string attr_str = attr.AsString().value_or("");
+    if (attr_str == "device.foo") {
+      return MatchResult::FULL;
+    } else if (attr_str == "device") {
+      return MatchResult::PARTIAL;
+    }
+    return MatchResult::NONE;
+  }
+
+  MatchResult CheckForMissing(const cel::Attribute& attr) const override {
+    std::string attr_str = attr.AsString().value_or("");
+
+    if (attr_str == "device2.foo") {
+      return MatchResult::FULL;
+    } else if (attr_str == "device2") {
+      return MatchResult::PARTIAL;
+    }
+    return MatchResult::NONE;
+  }
+};
+
+TEST_F(AttributeUtilityTest, CustomMatcher) {
+  AttributeTrail trail("device");
+
+  AttributeUtility utility(NoPatterns(), NoPatterns());
+  FakeMatcher matcher;
+  utility.set_matcher(&matcher);
+  EXPECT_TRUE(utility.CheckForUnknownPartial(trail));
+  EXPECT_FALSE(utility.CheckForUnknownExact(trail));
+
+  trail = trail.Step(cel::AttributeQualifier::OfString("foo"));
+  EXPECT_TRUE(utility.CheckForUnknownExact(trail));
+  EXPECT_TRUE(utility.CheckForUnknownPartial(trail));
+
+  trail = AttributeTrail("device2");
+  EXPECT_FALSE(utility.CheckForMissingAttribute(trail));
+  trail = trail.Step(cel::AttributeQualifier::OfString("foo"));
+  EXPECT_TRUE(utility.CheckForMissingAttribute(trail));
+}
+
 }  // namespace google::api::expr::runtime

eval/eval/evaluator_core.h

@@ -36,6 +36,7 @@
 #include "eval/eval/evaluator_stack.h"
 #include "eval/eval/iterator_stack.h"
 #include "runtime/activation_interface.h"
+#include "runtime/internal/activation_attribute_matcher_access.h"
 #include "runtime/runtime.h"
 #include "runtime/runtime_options.h"
 #include "google/protobuf/arena.h"
@@ -172,7 +173,15 @@ class ExecutionFrameBase {
                            activation.GetMissingAttributes()),
         slots_(&ComprehensionSlots::GetEmptyInstance()),
         max_iterations_(options.comprehension_max_iterations),
-        iterations_(0) {}
+        iterations_(0) {
+    if (unknown_processing_enabled()) {
+      if (auto matcher = cel::runtime_internal::
+              ActivationAttributeMatcherAccess::GetAttributeMatcher(activation);
+          matcher != nullptr) {
+        attribute_utility_.set_matcher(matcher);
+      }
+    }
+  }
 
   ExecutionFrameBase(const cel::ActivationInterface& activation,
                      EvaluationListener callback,
@@ -193,7 +202,15 @@ class ExecutionFrameBase {
                            activation.GetMissingAttributes()),
         slots_(&slots),
         max_iterations_(options.comprehension_max_iterations),
-        iterations_(0) {}
+        iterations_(0) {
+    if (unknown_processing_enabled()) {
+      if (auto matcher = cel::runtime_internal::
+              ActivationAttributeMatcherAccess::GetAttributeMatcher(activation);
+          matcher != nullptr) {
+        attribute_utility_.set_matcher(matcher);
+      }
+    }
+  }
 
   const cel::ActivationInterface& activation() const { return *activation_; }
 

eval/internal/BUILD

@@ -89,6 +89,8 @@ cc_library(
         "//internal:status_macros",
         "//runtime:activation_interface",
         "//runtime:function_overload_reference",
+        "//runtime/internal:activation_attribute_matcher_access",
+        "//runtime/internal:attribute_matcher",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

eval/internal/adapter_activation_impl.cc

@@ -25,6 +25,8 @@
 #include "eval/public/cel_value.h"
 #include "internal/status_macros.h"
 #include "runtime/function_overload_reference.h"
+#include "runtime/internal/activation_attribute_matcher_access.h"
+#include "runtime/internal/attribute_matcher.h"
 #include "google/protobuf/arena.h"
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/message.h"
@@ -76,4 +78,10 @@ absl::Span<const AttributePattern> AdapterActivationImpl::GetMissingAttributes()
   return legacy_activation_.missing_attribute_patterns();
 }
 
+const runtime_internal::AttributeMatcher* ABSL_NULLABLE
+AdapterActivationImpl::GetAttributeMatcher() const {
+  return runtime_internal::ActivationAttributeMatcherAccess::
+      GetAttributeMatcher(legacy_activation_);
+}
+
 }  // namespace cel::interop_internal

eval/internal/adapter_activation_impl.h

@@ -26,6 +26,7 @@
 #include "eval/public/base_activation.h"
 #include "runtime/activation_interface.h"
 #include "runtime/function_overload_reference.h"
+#include "runtime/internal/attribute_matcher.h"
 #include "google/protobuf/arena.h"
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/message.h"
@@ -56,6 +57,9 @@ class AdapterActivationImpl : public ActivationInterface {
   absl::Span<const cel::AttributePattern> GetMissingAttributes() const override;
 
  private:
+  const runtime_internal::AttributeMatcher* ABSL_NULLABLE GetAttributeMatcher()
+      const override;
+
   const google::api::expr::runtime::BaseActivation& legacy_activation_;
 };