(impala) 09/12: IMPALA-15219: Enforce strict validation for AI endpoint hostnames

[mi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

MikaelSmith pushed a commit to branch branch-4.5.2
in repository https://gitbox.apache.org/repos/asf/impala.git

commit 9e92b6f80f44b273967ed991f68f77bf417aa6af
Author: yida.wu <yi...@cloudera.com>
AuthorDate: Thu Jun 25 15:21:56 2026 -0700

    IMPALA-15219: Enforce strict validation for AI endpoint hostnames
    
    The AiFunctions currently validates allowed AI endpoints using
    a simple substring match (gstrncasestr). This can allow
    improperly formatted or malicious URLs to pass validation if
    they contain the target domain string in the path.
    
    This commit makes the validation stricter and more robust by
    enforcing structural hostname matching:
    1. Added a helper ExtractHost() to parse the URL and isolate
    the canonical hostname.
    2. Added a helper IsHostMatch() to perform exact domain and
    strict subdomain matching against the allowed endpoint list.
    3. Updated GetAiPlatformFromEndpoint() to evaluate the exact
    host using the new helper functions, ensuring malicious URLs
    are rejected as unsupported.
    
    Testing:
    Added test cases to AiFunctionsTest and
    AiFunctionsTestAdditionalSites to ensure that URLs with
    target domains hidden in paths are correctly rejected.
    
    Change-Id: I5481dec7d798dc6ca7273b0028b9317ab69c39f5
    Reviewed-on: http://gerrit.cloudera.org:8080/24775
    Reviewed-by: Joe McDonnell <jo...@cloudera.com>
    Tested-by: Michael Smith <mi...@cloudera.com>
---
 be/src/exprs/ai-functions-ir.cc | 70 +++++++++++++++++++++++++++++++----------
 be/src/exprs/expr-test.cc       | 21 +++++++++++++
 2 files changed, 74 insertions(+), 17 deletions(-)

diff --git a/be/src/exprs/ai-functions-ir.cc b/be/src/exprs/ai-functions-ir.cc
index d5906588a..0fb43affa 100644
--- a/be/src/exprs/ai-functions-ir.cc
+++ b/be/src/exprs/ai-functions-ir.cc
@@ -18,6 +18,9 @@
 // The functions in this file are specifically not cross-compiled to IR because there
 // is no signifcant performance benefit to be gained.
 
+#include <set>
+
+#include <boost/algorithm/string/predicate.hpp>
 #include <boost/algorithm/string/trim.hpp>
 
 #include "exprs/ai-functions.inline.h"
@@ -84,6 +87,39 @@ static const char* OPEN_AI_RESPONSE_FIELD_CHOICES = "choices";
 static const char* OPEN_AI_RESPONSE_FIELD_MESSAGE = "message";
 static const char* OPEN_AI_RESPONSE_FIELD_CONTENT = "content";
 
+// Helper to extract the hostname from a URL. Returns empty string if invalid.
+static string_view ExtractHost(const string_view& endpoint) {
+  string_view prefix = AI_API_ENDPOINT_PREFIX;
+  if (endpoint.length() < prefix.length() ||
+      strncasecmp(endpoint.data(), prefix.data(), prefix.length()) != 0) {
+    return {};
+  }
+  string_view without_prefix = endpoint.substr(prefix.length());
+  // Truncate the string at the first standard URL delimiter ('/', '?', '#',
+  // or ':'). This extracts the raw host name, ensuring that any trailing
+  // paths are ignored.
+  size_t host_end = without_prefix.find_first_of("/?:#");
+  if (host_end == string_view::npos) {
+    return without_prefix;
+  }
+  return without_prefix.substr(0, host_end);
+}
+
+// Helper to strictly validate a host against a target domain.
+bool IsHostMatch(const string_view& host, const string_view& target_domain) {
+  if (host.length() == target_domain.length()) {
+    return strncasecmp(host.data(), target_domain.data(), target_domain.length()) == 0;
+  } else if (host.length() > target_domain.length() + 1) {
+    // Allow subdomains by checking for a preceding dot.
+    // Length +1 ensures there is at least one character before the dot
+    // (e.g., 'a.domain.com').
+    if (host[host.length() - target_domain.length() - 1] == '.') {
+      return boost::algorithm::iends_with(host, target_domain);
+    }
+  }
+  return false;
+}
+
 /**
  * Singleton class for managing the additional AI platforms endpoints.
  * The additional platforms are loaded and parsed once to optimize for efficiency.
@@ -100,11 +136,11 @@ class AIAdditionalPlatforms {
   AIAdditionalPlatforms(const AIAdditionalPlatforms&) = delete;
   AIAdditionalPlatforms& operator=(const AIAdditionalPlatforms&) = delete;
 
-  // Check if the endpoint matches any of the additional platforms.
-  bool IsGeneralSite(const string_view& endpoint) const {
+  // Check if the extracted host matches any of the additional platforms.
+  bool IsGeneralSite(const string_view& host) const {
     return any_of(additional_platforms.begin(), additional_platforms.end(),
-        [&endpoint](const string& site) {
-          return gstrncasestr(endpoint.data(), site.c_str(), endpoint.size()) != nullptr;
+        [&host](const string& site) {
+          return IsHostMatch(host, site);
         });
   }
 
@@ -144,34 +180,33 @@ class AIAdditionalPlatforms {
 bool AiFunctions::is_api_endpoint_valid(const string_view& endpoint) {
   // Simple validation for endpoint. It should start with https://
   return (strncaseprefix(endpoint.data(), endpoint.size(), AI_API_ENDPOINT_PREFIX,
-              sizeof(AI_API_ENDPOINT_PREFIX))
+              strlen(AI_API_ENDPOINT_PREFIX))
       != nullptr);
 }
 
 bool AiFunctions::is_api_endpoint_supported(const string_view& endpoint) {
-  // Only OpenAI or configured general endpoints are supported.
-  return (
-      gstrncasestr(endpoint.data(), OPEN_AI_AZURE_ENDPOINT, endpoint.size()) != nullptr
-      || gstrncasestr(endpoint.data(), OPEN_AI_PUBLIC_ENDPOINT, endpoint.size())
-          != nullptr
-      || AIAdditionalPlatforms::GetInstance().IsGeneralSite(endpoint));
+  return GetAiPlatformFromEndpoint(endpoint, false)
+      != AiFunctions::AI_PLATFORM::UNSUPPORTED;
 }
 
 AiFunctions::AI_PLATFORM AiFunctions::GetAiPlatformFromEndpoint(
     const string_view& endpoint, bool dry_run) {
   if (UNLIKELY(dry_run)) AIAdditionalPlatforms::GetInstance().Reset();
 
-  // Only OpenAI or configured general endpoints are supported.
-  if (gstrncasestr(endpoint.data(), OPEN_AI_PUBLIC_ENDPOINT, endpoint.size())
-      != nullptr) {
+  // Validate the canonical host.
+  string_view host = ExtractHost(endpoint);
+  if (host.empty()) return AiFunctions::AI_PLATFORM::UNSUPPORTED;
+
+  if (IsHostMatch(host, OPEN_AI_PUBLIC_ENDPOINT)) {
     return AiFunctions::AI_PLATFORM::OPEN_AI;
   }
-  if (gstrncasestr(endpoint.data(), OPEN_AI_AZURE_ENDPOINT, endpoint.size()) != nullptr) {
+  if (IsHostMatch(host, OPEN_AI_AZURE_ENDPOINT)) {
     return AiFunctions::AI_PLATFORM::AZURE_OPEN_AI;
   }
-  if (AIAdditionalPlatforms::GetInstance().IsGeneralSite(endpoint)) {
+  if (AIAdditionalPlatforms::GetInstance().IsGeneralSite(host)) {
     return AI_PLATFORM::GENERAL;
   }
+
   return AiFunctions::AI_PLATFORM::UNSUPPORTED;
 }
 
@@ -233,6 +268,7 @@ StringVal AiFunctions::AiGenerateTextHelper(FunctionContext* ctx,
       return StringVal(AI_GENERATE_TXT_INVALID_PROTOCOL_ERROR.c_str());
     }
   }
+
   AI_PLATFORM platform = GetAiPlatformFromEndpoint(endpoint_sv);
   switch(platform) {
     case AI_PLATFORM::OPEN_AI:
@@ -247,7 +283,7 @@ StringVal AiFunctions::AiGenerateTextHelper(FunctionContext* ctx,
           prompt, model, auth_credential, platform_params, impala_options, false);
     default:
       if (fastpath) {
-        DCHECK(false) << "Default endpoint " << FLAGS_ai_endpoint << "must be supported";
+        DCHECK(false) << "Default endpoint " << FLAGS_ai_endpoint << " must be supported";
       }
       LOG(ERROR) << "AI Generate Text: \nunsupported endpoint: " << endpoint_sv;
       return StringVal(AI_GENERATE_TXT_UNSUPPORTED_ENDPOINT_ERROR.c_str());
diff --git a/be/src/exprs/expr-test.cc b/be/src/exprs/expr-test.cc
index d3455ecb3..47e459a6f 100644
--- a/be/src/exprs/expr-test.cc
+++ b/be/src/exprs/expr-test.cc
@@ -11313,6 +11313,20 @@ TEST_P(ExprTest, AiFunctionsTest) {
   EXPECT_EQ(AiFunctions::AI_PLATFORM::UNSUPPORTED,
       AiFunctions::GetAiPlatformFromEndpoint("https://qwerty.com"));
 
+  // Test Malicious Endpoint Rejection
+  // Keyword hidden in the path
+  EXPECT_EQ(AiFunctions::AI_PLATFORM::UNSUPPORTED,
+      AiFunctions::GetAiPlatformFromEndpoint("https://fake.com/api.openai.com"));
+  // Keyword hidden in the query parameters
+  EXPECT_EQ(AiFunctions::AI_PLATFORM::UNSUPPORTED,
+      AiFunctions::GetAiPlatformFromEndpoint("https://fake.com/?url=openai.azure.com"));
+  // Spoofed Top-Level Domain (TLD) / suffix extension
+  EXPECT_EQ(AiFunctions::AI_PLATFORM::UNSUPPORTED,
+      AiFunctions::GetAiPlatformFromEndpoint("https://api.openai.com.fake.com/v1/chat"));
+  // Prefix extension bypassing exact dot notation (e.g., fake-api instead of .api)
+  EXPECT_EQ(AiFunctions::AI_PLATFORM::UNSUPPORTED,
+      AiFunctions::GetAiPlatformFromEndpoint("https://fake-api.openai.com/v1/chat"));
+
   // Test fastpath
   StringVal result =
       AiFunctions::AiGenerateTextInternal<true, AiFunctions::AI_PLATFORM::OPEN_AI>(ctx,
@@ -11591,6 +11605,13 @@ TEST_P(ExprTest, AiFunctionsTestAdditionalSites) {
   EXPECT_EQ(
       AiFunctions::GetAiPlatformFromEndpoint("https://AI-API.COM/v1/generate", true),
       AiFunctions::AI_PLATFORM::GENERAL);
+
+  // Test for supported/unsupported endpoints.
+  EXPECT_EQ(
+      AiFunctions::is_api_endpoint_supported("https://fake.com/ai-api.com"), false);
+  EXPECT_EQ(
+      AiFunctions::is_api_endpoint_supported("https://another-ai.org.fake.com"), false);
+
 }
 
 } // namespace impala