optimize custom-response filter for ai-fallback

higress-group · Nov 25, 2024 · 0cf0aea · 0cf0aea
1 parent e9302f5
commit 0cf0aea
Show file tree

Hide file tree

Showing 8 changed files with 61 additions and 35 deletions.
diff --git a/api/envoy/extensions/http/custom_response/redirect_policy/v3/redirect_policy.proto b/api/envoy/extensions/http/custom_response/redirect_policy/v3/redirect_policy.proto
@@ -51,6 +51,8 @@ message RedirectPolicy {
     // - `regex_rewrite`
     config.route.v3.RedirectAction redirect_action = 2;
 
+    string uri_from_response_header = 106 [(validate.rules).string = {min_len: 1}];
+
     google.protobuf.BoolValue use_original_request_uri = 107;
   }
 

diff --git a/source/common/http/conn_manager_impl.cc b/source/common/http/conn_manager_impl.cc
@@ -1349,6 +1349,13 @@ void ConnectionManagerImpl::ActiveStream::decodeHeaders(RequestHeaderMapSharedPt
 
     filter_manager_.setDownstreamRemoteAddress(mutate_result.final_remote_address);
   }
+#if defined(HIGRESS)
+  else {
+    request_headers_->setReferenceKey(
+        Http::CustomHeaders::get().AliExtendedValues.XEnvoyInternalRoute,
+        Http::CustomHeaders::get().EnvoyIntenralRouteValues.True);
+  }
+#endif
   ASSERT(filter_manager_.streamInfo().downstreamAddressProvider().remoteAddress() != nullptr);
 
   ASSERT(!cached_route_);

diff --git a/source/common/http/conn_manager_utility.cc b/source/common/http/conn_manager_utility.cc
@@ -325,6 +325,9 @@ void ConnectionManagerUtility::cleanInternalHeaders(
   request_headers.removeEnvoyIpTags();
   request_headers.removeEnvoyOriginalUrl();
   request_headers.removeEnvoyHedgeOnPerTryTimeout();
+#if defined(HIGRESS)
+  request_headers.remove(Http::CustomHeaders::get().AliExtendedValues.XEnvoyInternalRoute);
+#endif
 
   for (const LowerCaseString& header : internal_only_headers) {
     request_headers.remove(header);

diff --git a/source/common/http/headers.h b/source/common/http/headers.h
@@ -138,7 +138,12 @@ class CustomHeaderValues {
     const LowerCaseString EnvoyOriginalHost{"original-host"};
     const LowerCaseString XEnvoyOriginalHost{"x-envoy-original-host"};
     const LowerCaseString XEnvoyRouteIdentifier{"x-envoy-route-identifier"};
+    const LowerCaseString XEnvoyInternalRoute{"x-envoy-internal-route"};
   } AliExtendedValues;
+
+  struct {
+    const std::string True{"true"};
+  } EnvoyIntenralRouteValues;
 #endif
 };
 

diff --git a/source/extensions/filters/http/custom_response/custom_response_filter.cc b/source/extensions/filters/http/custom_response/custom_response_filter.cc
@@ -20,11 +20,7 @@ Http::FilterHeadersStatus CustomResponseFilter::decodeHeaders(Http::RequestHeade
   downstream_headers_ = &header_map;
   const FilterConfig* config = nullptr;
   if (decoder_callbacks_ && decoder_callbacks_->route()) {
-    const auto* config =
-        Http::Utility::resolveMostSpecificPerFilterConfig<FilterConfig>(decoder_callbacks_);
-    if (config != nullptr) {
-      has_rules_ = true;
-    }
+    config = Http::Utility::resolveMostSpecificPerFilterConfig<FilterConfig>(decoder_callbacks_);
   }
   if (config == nullptr) {
     config = config_.get();
@@ -77,24 +73,19 @@ Http::FilterHeadersStatus CustomResponseFilter::encodeHeaders(Http::ResponseHead
   // policy. Note that since the traversal is least to most specific, we can't
   // return early when a match is found.
   PolicySharedPtr policy;
-#if defined(HIGRESS)
-  if (has_rules_) {
-#endif
-    decoder_callbacks_->traversePerFilterConfig(
-        [&policy, &headers, this](const Router::RouteSpecificFilterConfig& config) {
-          const FilterConfig* typed_config = dynamic_cast<const FilterConfig*>(&config);
-          if (typed_config) {
-            // Check if a match is found first to avoid overwriting policy with an
-            // empty shared_ptr.
-            auto maybe_policy = typed_config->getPolicy(headers, encoder_callbacks_->streamInfo());
-            if (maybe_policy) {
-              policy = maybe_policy;
-            }
+  decoder_callbacks_->traversePerFilterConfig(
+      [&policy, &headers, this](const Router::RouteSpecificFilterConfig& config) {
+        const FilterConfig* typed_config = dynamic_cast<const FilterConfig*>(&config);
+        if (typed_config) {
+          // Check if a match is found first to avoid overwriting policy with an
+          // empty shared_ptr.
+          auto maybe_policy = typed_config->getPolicy(headers, encoder_callbacks_->streamInfo());
+          if (maybe_policy) {
+            policy = maybe_policy;
           }
-        });
-#if defined(HIGRESS)
-  }
-#endif
+        }
+      });
+
   if (!policy) {
     policy = config_->getPolicy(headers, encoder_callbacks_->streamInfo());
   }

diff --git a/source/extensions/filters/http/custom_response/custom_response_filter.h b/source/extensions/filters/http/custom_response/custom_response_filter.h
@@ -50,9 +50,6 @@ class CustomResponseFilter : public ::Envoy::Http::PassThroughFilter,
   const std::shared_ptr<const FilterConfig> config_;
   ::Envoy::Http::RequestHeaderMap* downstream_headers_ = nullptr;
   bool on_local_reply_called_ = false;
-#if defined(HIGRESS)
-  bool has_rules_ = false;
-#endif
 };
 
 } // namespace CustomResponse

diff --git a/source/extensions/http/custom_response/redirect_policy/redirect_policy.cc b/source/extensions/http/custom_response/redirect_policy/redirect_policy.cc
@@ -64,6 +64,9 @@ RedirectPolicy::RedirectPolicy(
                                  createRedirectConfig(config.redirect_action()))
                            : nullptr},
 #if defined(HIGRESS)
+      uri_from_response_header_{config.has_uri_from_response_header()
+                                    ? absl::optional<std::string>(config.uri_from_response_header())
+                                    : absl::nullopt},
       use_original_request_uri_(
           PROTOBUF_GET_WRAPPED_OR_DEFAULT(config, use_original_request_uri, false)),
       keep_original_response_code_(
@@ -85,7 +88,8 @@ RedirectPolicy::RedirectPolicy(
       modify_request_headers_action_(createModifyRequestHeadersAction(config, context)) {
 #if defined(HIGRESS)
   // Ensure that exactly one of uri_ or redirect_action_ or use_original_request_uri_ is specified.
-  ASSERT(int(uri_ != nullptr) + int(redirect_action_ != nullptr) + int(use_original_request_uri_) ==
+  ASSERT(int(uri_ != nullptr) + int(redirect_action_ != nullptr) + int(use_original_request_uri_) +
+             int(uri_from_response_header_.has_value()) ==
          1);
 #else
   // Ensure that exactly one of uri_ or redirect_action_ is specified.
@@ -236,17 +240,33 @@ ::Envoy::Http::FilterHeadersStatus RedirectPolicy::encodeHeaders(
     downstream_headers->setHost(real_original_host);
     downstream_headers->setPath(real_original_path);
   } else {
+    std::string uri;
+    if (uri_from_response_header_.has_value()) {
+      auto custom_location = headers.get(Envoy::Http::LowerCaseString(*uri_from_response_header_));
+      uri = custom_location.empty() ? "" : custom_location[0]->value().getStringView();
+      if (uri == "" || !absolute_url.initialize(uri, false)) {
+        stats_.custom_response_invalid_uri_.inc();
+        ENVOY_LOG(debug, "uri specified in response header is invalid");
+        return ::Envoy::Http::FilterHeadersStatus::Continue;
+      }
+    } else {
+      uri = uri_ ? *uri_ : ::Envoy::Http::Utility::newUri(*redirect_action_, *downstream_headers);
+#else
+  std::string uri(uri_ ? *uri_
+                       : ::Envoy::Http::Utility::newUri(*redirect_action_, *downstream_headers));
 #endif
-    std::string uri(uri_ ? *uri_
-                         : ::Envoy::Http::Utility::newUri(*redirect_action_, *downstream_headers));
-    if (!absolute_url.initialize(uri, false)) {
-      stats_.custom_response_invalid_uri_.inc();
-      // We could potentially get an invalid url only if redirect_action_ was specified instead
-      // of uri_. Hence, assert that uri_ is not set.
-      ENVOY_BUG(!static_cast<bool>(uri_),
-                "uri should not be invalid as this was already validated during config load");
-      return ::Envoy::Http::FilterHeadersStatus::Continue;
+
+      if (!absolute_url.initialize(uri, false)) {
+        stats_.custom_response_invalid_uri_.inc();
+        // We could potentially get an invalid url only if redirect_action_ was specified instead
+        // of uri_. Hence, assert that uri_ is not set.
+        ENVOY_BUG(!static_cast<bool>(uri_),
+                  "uri should not be invalid as this was already validated during config load");
+        return ::Envoy::Http::FilterHeadersStatus::Continue;
+      }
+#if defined(HIGRESS)
     }
+#endif
     downstream_headers->setScheme(absolute_url.scheme());
     downstream_headers->setHost(absolute_url.hostAndPort());
 

diff --git a/source/extensions/http/custom_response/redirect_policy/redirect_policy.h b/source/extensions/http/custom_response/redirect_policy/redirect_policy.h
@@ -71,6 +71,7 @@ class RedirectPolicy : public Extensions::HttpFilters::CustomResponse::Policy,
   const std::unique_ptr<const std::string> uri_;
   const std::unique_ptr<const ::Envoy::Http::Utility::RedirectConfig> redirect_action_;
 #if defined(HIGRESS)
+  absl::optional<std::string> uri_from_response_header_;
   bool use_original_request_uri_;
   bool keep_original_response_code_;
   uint32_t max_internal_redirects_;