feat: Support extracting model argument from body in multipart/form-data format (#1940)

plugins/wasm-cpp/extensions/model_router/plugin.cc

@@ -16,6 +16,7 @@
 
 #include <array>
 #include <limits>
+#include <regex>
 
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_split.h"
@@ -123,6 +124,7 @@ bool PluginRootContext::configure(size_t configuration_size) {
 }
 
 FilterHeadersStatus PluginRootContext::onHeader(
+    PluginContext& ctx,
     const ModelRouterConfigRule& rule) {
   if (!Wasm::Common::Http::hasRequestBody()) {
     return FilterHeadersStatus::Continue;
@@ -150,19 +152,49 @@ FilterHeadersStatus PluginRootContext::onHeader(
   if (!enable) {
     return FilterHeadersStatus::Continue;
   }
-  auto content_type_value =
+  auto content_type_ptr =
       getRequestHeader(Wasm::Common::Http::Header::ContentType);
-  if (!absl::StrContains(content_type_value->view(),
+  auto content_type_value = content_type_ptr->view();
+  LOG_DEBUG(absl::StrCat("Content-Type: ", content_type_value));
+  if (absl::StrContains(content_type_value,
                          Wasm::Common::Http::ContentTypeValues::Json)) {
-    return FilterHeadersStatus::Continue;
+    ctx.mode_ = MODE_JSON;
+    LOG_DEBUG("Enable JSON mode.");
+    removeRequestHeader(Wasm::Common::Http::Header::ContentLength);
+    setFilterState(SetDecoderBufferLimitKey, DefaultMaxBodyBytes);
+    LOG_INFO(absl::StrCat("SetRequestBodyBufferLimit: ", DefaultMaxBodyBytes));
+    return FilterHeadersStatus::StopIteration;
   }
-  removeRequestHeader(Wasm::Common::Http::Header::ContentLength);
-  setFilterState(SetDecoderBufferLimitKey, DefaultMaxBodyBytes);
-  LOG_INFO(absl::StrCat("SetRequestBodyBufferLimit: ", DefaultMaxBodyBytes));
-  return FilterHeadersStatus::StopIteration;
+  if (absl::StrContains(content_type_value,
+                         Wasm::Common::Http::ContentTypeValues::MultipartFormData)) {
+    // Get the boundary from the content type
+    auto boundary_start = content_type_value.find("boundary=");
+    if (boundary_start == std::string::npos) {
+      LOG_WARN(absl::StrCat("No boundary found in a multipart/form-data content-type: ", content_type_value));
+      return FilterHeadersStatus::Continue;
+    }
+    boundary_start += 9;
+    auto boundary_end = content_type_value.find(';', boundary_start);
+    if (boundary_end == std::string::npos) {
+      boundary_end = content_type_value.size();
+    }
+    auto boundary_length = boundary_end - boundary_start;
+    if (boundary_length < 1 || boundary_length > 70) {
+      // See https://www.w3.org/Protocols/rfc1341/7_2_Multipart.html
+      LOG_WARN(absl::StrCat("Invalid boundary value in a multipart/form-data content-type: ", content_type_value));
+      return FilterHeadersStatus::Continue;
+    }
+    auto boundary_value = content_type_value.substr(boundary_start, boundary_end - boundary_start);
+    ctx.mode_ = MODE_MULTIPART;
+    ctx.boundary_ = boundary_value;
+    LOG_DEBUG(absl::StrCat("Enable multipart/form-data mode. Boundary=", boundary_value));
+    removeRequestHeader(Wasm::Common::Http::Header::ContentLength);
+    return FilterHeadersStatus::StopIteration;
+  }
+  return FilterHeadersStatus::Continue;
 }
 
-FilterDataStatus PluginRootContext::onBody(const ModelRouterConfigRule& rule,
+FilterDataStatus PluginRootContext::onJsonBody(const ModelRouterConfigRule& rule,
                                            std::string_view body) {
   const auto& model_key = rule.model_key_;
   const auto& add_provider_header = rule.add_provider_header_;
@@ -198,10 +230,85 @@ FilterDataStatus PluginRootContext::onBody(const ModelRouterConfigRule& rule,
   return FilterDataStatus::Continue;
 }
 
+FilterDataStatus PluginRootContext::onMultipartBody(
+    PluginContext& ctx,
+    const ModelRouterConfigRule& rule,
+    WasmDataPtr& body,
+    bool end_stream) {
+  const auto& add_provider_header = rule.add_provider_header_;
+  const auto& model_to_header = rule.model_to_header_;
+
+  const auto boundary = ctx.boundary_;
+  const auto body_view = body->view();
+  const auto model_param_header = absl::StrCat("Content-Disposition: form-data; name=\"", rule.model_key_, "\"");
+
+  for (size_t pos = 0; (pos = body_view.find(boundary, pos)) != std::string_view::npos;) {
+    LOG_DEBUG(absl::StrCat("Found boundary at ", pos));
+    pos += boundary.length();
+    size_t end_pos = body_view.find(boundary, pos);
+    if (end_pos == std::string_view::npos) {
+      end_pos = body_view.length();
+    }
+    std::string_view part = body_view.substr(pos, end_pos - pos);
+    LOG_DEBUG(absl::StrCat("Part: ", part));
+    auto part_pos = pos;
+    pos = end_pos;
+
+    // Check if this part contains the model parameter
+    if (!absl::StrContains(part, model_param_header)) {
+      LOG_DEBUG("Part does not contain model parameter");
+      continue;
+    }
+    size_t value_start = part.find(CRLF_CRLF);
+    if (value_start == std::string_view::npos) {
+      LOG_DEBUG("No value start found in part");
+      break;
+    }
+    value_start += 4; // Skip the "\r\n\r\n"
+    // model parameter should be only one line
+    size_t value_end = part.find(CRLF, value_start);
+    if (value_end == std::string_view::npos) {
+      LOG_DEBUG("No value end found in part");
+      break;
+    }
+    auto model_value = part.substr(value_start, value_end - value_start);
+    LOG_DEBUG(absl::StrCat("Model value: ", model_value));
+    if (!model_to_header.empty()) {
+      replaceRequestHeader(model_to_header, model_value);
+    }
+    if (!add_provider_header.empty()) {
+      auto pos = model_value.find('/');
+      if (pos != std::string::npos) {
+        const auto& provider = model_value.substr(0, pos);
+        const auto& model = model_value.substr(pos + 1);
+        replaceRequestHeader(add_provider_header, provider);
+        size_t new_size = 0;
+        auto new_buffer_data = absl::StrCat(body_view.substr(0, part_pos + value_start), model, body_view.substr(part_pos + value_end));
+        auto result = setBuffer(WasmBufferType::HttpRequestBody, 0, std::numeric_limits<size_t>::max(), new_buffer_data, &new_size);
+        LOG_DEBUG(absl::StrCat("model route to provider:", provider,
+                               ", model:", model));
+        LOG_DEBUG(absl::StrCat("result=", result, "  new_size=", new_size));
+      } else {
+        LOG_DEBUG(absl::StrCat("model route to provider not work, model:",
+                               model_value));
+      }
+    }
+    // We are done now. We can stop processing the body.
+    LOG_DEBUG(absl::StrCat("Done processing multipart body after caching ", body_view.length() , " bytes."));
+    ctx.mode_ = MODE_BYPASS;
+    return FilterDataStatus::Continue;
+  }
+  if (end_stream) {
+    LOG_DEBUG("No model parameter found in the body");
+    return FilterDataStatus::Continue;
+  }
+  return FilterDataStatus::StopIterationAndBuffer;
+}
+
 FilterHeadersStatus PluginContext::onRequestHeaders(uint32_t, bool) {
   auto* rootCtx = rootContext();
   return rootCtx->onHeaders([rootCtx, this](const auto& config) {
-    auto ret = rootCtx->onHeader(config);
+    auto ret = rootCtx->onHeader(*this, config);
     if (ret == FilterHeadersStatus::StopIteration) {
       this->config_ = &config;
     }
@@ -214,14 +321,28 @@ FilterDataStatus PluginContext::onRequestBody(size_t body_size,
   if (config_ == nullptr) {
     return FilterDataStatus::Continue;
   }
-  body_total_size_ += body_size;
-  if (!end_stream) {
-    return FilterDataStatus::StopIterationAndBuffer;
-  }
-  auto body =
-      getBufferBytes(WasmBufferType::HttpRequestBody, 0, body_total_size_);
   auto* rootCtx = rootContext();
-  return rootCtx->onBody(*config_, body->view());
+  body_total_size_ += body_size;
+  switch (mode_) {
+    case MODE_JSON:
+    {
+      if (!end_stream) {
+        return FilterDataStatus::StopIterationAndBuffer;
+      }
+      auto body =
+          getBufferBytes(WasmBufferType::HttpRequestBody, 0, body_total_size_);
+      return rootCtx->onJsonBody(*config_, body->view());
+    }
+    case MODE_MULTIPART:
+    {
+      auto body =
+          getBufferBytes(WasmBufferType::HttpRequestBody, 0, body_total_size_);
+      return rootCtx->onMultipartBody(*this, *config_, body, end_stream);
+    }
+    case MODE_BYPASS:
+    default:
+      return FilterDataStatus::Continue;
+  }
 }
 
 #ifdef NULL_PLUGIN