grpc 1.81.1 → 1.82.0.pre1

data/src/core/ext/transport/chttp2/transport/header_assembler.h

@@ -23,6 +23,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <optional>
 #include <string>
 #include <utility>
 
@@ -41,7 +42,7 @@
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 
-// TODO(tjagtap) TODO(akshitpatel): [PH2][P3] : Write micro benchmarks for
+// TODO(akshitpatel): [PH2][P4] : Write micro benchmarks for
 // assembler and disassembler code
 
 namespace grpc_core {
@@ -100,6 +101,8 @@ class HeaderAssembler {
 
     // Start header workflow
     header_in_progress_ = true;
+    GRPC_DCHECK_GT(stream_id_, 0u);
+    GRPC_DCHECK_EQ(stream_id_, frame.stream_id);
 
     // Manage payload
     frame.payload.MoveFirstNBytesIntoSliceBuffer(current_len, buffer_);
@@ -121,6 +124,8 @@ class HeaderAssembler {
   Http2Status AppendFrame(Http2ContinuationFrame& frame) {
     // Validate Assembler state
     GRPC_DCHECK(header_in_progress_);
+    GRPC_DCHECK_GT(stream_id_, 0u);
+    GRPC_DCHECK_EQ(stream_id_, frame.stream_id);
 
     // Manage payload
     const size_t current_len = frame.payload.Length();
@@ -195,18 +200,27 @@ class HeaderAssembler {
                   << max_header_list_size_soft_limit
                   << "max_header_list_size_hard_limit: "
                   << max_header_list_size_hard_limit;
-    Http2Status status = ParseHeader(
-        parser, std::move(buffer_), /*grpc_metadata_batch=*/nullptr,
-        ParseHeaderArgs{
-            /*is_initial_metadata=*/is_initial_metadata,
-            /*is_end_headers=*/is_ready_,
-            /*is_client=*/is_client_,
-            /*max_header_list_size_soft_limit=*/max_header_list_size_soft_limit,
-            /*max_header_list_size_hard_limit=*/max_header_list_size_hard_limit,
-            /*stream_id=*/stream_id_,
-        });
-    Cleanup();
-    return status;
+    if (buffer_.Length() > 0) {
+      Http2Status status = ParseHeader(
+          parser, std::move(buffer_), /*grpc_metadata_batch=*/nullptr,
+          ParseHeaderArgs{
+              /*is_initial_metadata=*/is_initial_metadata,
+              /*is_end_headers=*/is_ready_,
+              /*is_client=*/is_client_,
+              /*max_header_list_size_soft_limit=*/
+              max_header_list_size_soft_limit,
+              /*max_header_list_size_hard_limit=*/
+              max_header_list_size_hard_limit,
+              /*stream_id=*/stream_id_,
+          });
+      Cleanup();
+      return status;
+    }
+
+    GRPC_DCHECK(!is_ready_);
+    GRPC_DCHECK(buffer_.Length() == 0);
+    GRPC_DCHECK(!header_in_progress_);
+    return Http2Status::Ok();
   }
 
   size_t GetBufferedHeadersLength() const { return buffer_.Length(); }
@@ -217,7 +231,6 @@ class HeaderAssembler {
   explicit HeaderAssembler(const bool is_client)
       : header_in_progress_(false),
         is_ready_(false),
-        allow_true_binary_metadata_acked_(true),
         is_client_(is_client),
         stream_id_(0) {}
 
@@ -228,12 +241,24 @@ class HeaderAssembler {
   HeaderAssembler(const HeaderAssembler&) = delete;
   HeaderAssembler& operator=(const HeaderAssembler&) = delete;
 
-  void InitializeStream(const uint32_t stream_id,
-                        const bool allow_true_binary_metadata_acked) {
-    GRPC_DCHECK_EQ(stream_id_, 0u);
+  // This is set only once in the lifetime of a transport.
+  // TODO(tjagtap) [PH2][P3] : This value in settings can only be set once as
+  // the allow_true_binary_metadata setting is sent only once. Re-think this
+  // once we add actual use of this value.
+  void MaybeSetAllowTrueBinaryMetadataAcked(
+      const bool allow_true_binary_metadata_acked) {
+    if (allow_true_binary_metadata_acked_.has_value()) {
+      return;
+    }
+    allow_true_binary_metadata_acked_ = allow_true_binary_metadata_acked;
+  }
+
+  void SetStreamId(uint32_t stream_id) {
     GRPC_DCHECK_NE(stream_id, 0u);
+    GRPC_DCHECK(!header_in_progress_);
+    GRPC_DCHECK_EQ(is_ready_, false);
+    GRPC_DCHECK(buffer_.Length() == 0);
     stream_id_ = stream_id;
-    allow_true_binary_metadata_acked_ = allow_true_binary_metadata_acked;
   }
 
   // HPACK parser helpers
@@ -250,14 +275,15 @@ class HeaderAssembler {
           "is_initial_metadata: ", is_initial_metadata,
           " is_end_headers: ", is_end_headers, " is_client: ", is_client,
           " max_header_list_size_soft_limit: ", max_header_list_size_soft_limit,
-          " max_header_list_size_hard_limit: ",
-          max_header_list_size_hard_limit);
+          " max_header_list_size_hard_limit: ", max_header_list_size_hard_limit,
+          " stream_id:", stream_id);
     }
   };
 
   static Http2Status ParseHeader(HPackParser& parser, SliceBuffer&& buffer,
                                  grpc_metadata_batch* grpc_metadata_batch,
                                  const ParseHeaderArgs args) {
+    // TODO(alishananda): Propagate mitigation engine here.
     parser.BeginFrame(
         /*grpc_metadata_batch*/ grpc_metadata_batch,
         args.max_header_list_size_soft_limit,
@@ -268,7 +294,8 @@ class HeaderAssembler {
                              args.is_initial_metadata
                                  ? HPackParser::LogInfo::Type::kHeaders
                                  : HPackParser::LogInfo::Type::kTrailers,
-                             args.is_client});
+                             args.is_client},
+        /*mitigation_engine=*/nullptr);
     // TODO(tjagtap) [PH2][P5] Bug fix : Check if the received metadata honours
     // allow_true_binary_metadata or not. Will need changes to HPack code.
     absl::Status stream_error = absl::OkStatus();
@@ -314,16 +341,21 @@ class HeaderAssembler {
     return boundary;
   }
 
+  std::optional<bool> TestOnlyAllowTrueBinaryMetadataAcked() const {
+    return allow_true_binary_metadata_acked_;
+  }
+
  private:
   void Cleanup() {
     buffer_.Clear();
     header_in_progress_ = false;
     is_ready_ = false;
+    stream_id_ = 0;
   }
 
   bool header_in_progress_;
   bool is_ready_;
-  bool allow_true_binary_metadata_acked_;
+  std::optional<bool> allow_true_binary_metadata_acked_;
   const bool is_client_;
   uint32_t stream_id_;
   SliceBuffer buffer_;

data/src/core/ext/transport/chttp2/transport/hpack_encoder.cc

@@ -46,6 +46,21 @@ namespace {
 
 constexpr size_t kHeadersFrameHeaderSize = 9;
 
+std::optional<size_t> CheckIndexSize(size_t key_len, size_t value_len) {
+  constexpr size_t max_size = HPackEncoderTable::MaxEntrySize();
+
+  if (key_len > max_size - hpack_constants::kEntryOverhead) {
+    // Too large to index or overflow.
+    return std::nullopt;
+  }
+  size_t sum = key_len + hpack_constants::kEntryOverhead;
+  if (value_len > max_size - sum) {
+    // Too large to index or overflow.
+    return std::nullopt;
+  }
+  return sum + value_len;
+}
+
 }  // namespace
 
 // fills p (which is expected to be kHeadersFrameHeaderSize bytes long)
@@ -254,6 +269,21 @@ uint32_t HPackWriter::EmitLitHdrWithNonBinaryStringKeyIncIdx(
     HPackEncoderTable& table) {
   auto key_len = key_slice.length();
   auto value_len = value_slice.length();
+
+  size_t size_to_allocate;
+  if (IsOptimization03Enabled()) {
+    auto size = CheckIndexSize(key_len, value_len);
+    if (!size.has_value()) {
+      // Fall back to not indexing.
+      EmitLitHdrWithNonBinaryStringKeyNotIdx(std::move(key_slice),
+                                             std::move(value_slice), output);
+      return 0;
+    }
+    size_to_allocate = *size;
+  } else {
+    size_to_allocate = key_len + value_len + hpack_constants::kEntryOverhead;
+  }
+
   StringKey key(std::move(key_slice));
   key.WritePrefix(0x40, output.AddTiny(key.prefix_length()));
   output.Append(key.key());
@@ -261,8 +291,7 @@ uint32_t HPackWriter::EmitLitHdrWithNonBinaryStringKeyIncIdx(
   emit.WritePrefix(output.AddTiny(emit.prefix_length()));
   // Allocate an index in the hpack table for this newly emitted entry.
   // (we do so here because we know the length of the key and value)
-  uint32_t index = table.AllocateIndex(key_len + value_len +
-                                       hpack_constants::kEntryOverhead);
+  uint32_t index = table.AllocateIndex(size_to_allocate);
   output.Append(emit.data());
   return index;
 }
@@ -282,15 +311,34 @@ uint32_t HPackWriter::EmitLitHdrWithBinaryStringKeyIncIdx(
     Slice key_slice, Slice value_slice, SliceBuffer& output,
     HPackEncoderTable& table, bool use_true_binary_metadata) {
   auto key_len = key_slice.length();
+  BinaryStringValue emit(std::move(value_slice), use_true_binary_metadata);
+
+  size_t size_to_allocate;
+
+  if (IsOptimization03Enabled()) {
+    auto size = CheckIndexSize(key_len, emit.hpack_length());
+    if (!size.has_value()) {
+      // Fall back to not indexing.
+      StringKey key(std::move(key_slice));
+      key.WritePrefix(0x00, output.AddTiny(key.prefix_length()));
+      output.Append(key.key());
+      emit.WritePrefix(output.AddTiny(emit.prefix_length()));
+      output.Append(emit.data());
+      return 0;
+    }
+    size_to_allocate = *size;
+  } else {
+    size_to_allocate =
+        key_len + emit.hpack_length() + hpack_constants::kEntryOverhead;
+  }
+
   StringKey key(std::move(key_slice));
   key.WritePrefix(0x40, output.AddTiny(key.prefix_length()));
   output.Append(key.key());
-  BinaryStringValue emit(std::move(value_slice), use_true_binary_metadata);
   emit.WritePrefix(output.AddTiny(emit.prefix_length()));
   // Allocate an index in the hpack table for this newly emitted entry.
   // (we do so here because we know the length of the key and value)
-  uint32_t index = table.AllocateIndex(key_len + emit.hpack_length() +
-                                       hpack_constants::kEntryOverhead);
+  uint32_t index = table.AllocateIndex(size_to_allocate);
   output.Append(emit.data());
   return index;
 }

data/src/core/ext/transport/chttp2/transport/hpack_parse_result.cc

@@ -176,6 +176,28 @@ absl::Status HpackParseResult::BuildMaterialized() const {
           "Invalid HPACK index received (%d)", state_->invalid_hpack_index));
     case HpackParseStatus::kIllegalHpackOpCode:
       return absl::InternalError("Illegal hpack op code");
+    case HpackParseStatus::kMitigationEngineConnectionError:
+      if (state_->key.empty()) {
+        return absl::InternalError(
+            absl::StrCat("Mitigation engine triggered action ",
+                         MitigationEngine::ActionToString(state_->action)));
+      } else {
+        return absl::InternalError(
+            absl::StrCat("Mitigation engine triggered action ",
+                         MitigationEngine::ActionToString(state_->action),
+                         " for key: ", state_->key));
+      }
+    case HpackParseStatus::kMitigationEngineStreamError:
+      if (state_->key.empty()) {
+        return MakeStreamError(absl::InternalError(
+            absl::StrCat("Mitigation engine triggered action ",
+                         MitigationEngine::ActionToString(state_->action))));
+      } else {
+        return MakeStreamError(absl::InternalError(
+            absl::StrCat("Mitigation engine triggered action ",
+                         MitigationEngine::ActionToString(state_->action),
+                         " for key: ", state_->key)));
+      }
   }
   GPR_UNREACHABLE_CODE(return absl::UnknownError("Should never reach here"));
 }

data/src/core/ext/transport/chttp2/transport/hpack_parse_result.h

@@ -25,6 +25,7 @@
 
 #include "src/core/call/metadata_batch.h"
 #include "src/core/lib/surface/validate_metadata.h"
+#include "src/core/mitigation_engine/mitigation_engine.h"
 #include "src/core/util/crash.h"
 #include "src/core/util/grpc_check.h"
 #include "src/core/util/ref_counted.h"
@@ -69,6 +70,8 @@ enum class HpackParseStatus : uint8_t {
   kMetadataParseError,
   // Parse failed due to a base64 decode error
   kUnbase64Failed,
+  // Error triggered by the mitigation engine (stream error)
+  kMitigationEngineStreamError,
 
   ///////////////////////////////////////////////////////////////////
   // Connection Errors - result in the tcp connection closing
@@ -102,6 +105,8 @@ enum class HpackParseStatus : uint8_t {
   kMaliciousVarintEncoding,
   // Illegal hpack op code
   kIllegalHpackOpCode,
+  // Error triggered by the mitigation engine
+  kMitigationEngineConnectionError,
 };
 
 inline bool IsStreamError(HpackParseStatus status) {
@@ -260,6 +265,18 @@ class HpackParseResult {
     return p;
   }
 
+  static HpackParseResult MitigationEngineError(
+      absl::string_view key, MitigationEngine::Action action) {
+    HpackParseStatus status =
+        action == MitigationEngine::Action::kCloseConnection
+            ? HpackParseStatus::kMitigationEngineConnectionError
+            : HpackParseStatus::kMitigationEngineStreamError;
+    HpackParseResult p{status};
+    p.state_->key = std::string(key);
+    p.state_->action = action;
+    return p;
+  }
+
   // Compute the absl::Status that goes along with this HpackParseResult.
   // (may be cached, so this is not thread safe)
   absl::Status Materialize() const;
@@ -335,6 +352,9 @@ class HpackParseResult {
       MetadataLimitExceededByAtom metadata_limit_exceeded_by_atom;
       // Set if status == kIllegalTableSizeChange
       IllegalTableSizeChange illegal_table_size_change;
+      // Set if status == kMitigationEngineConnectionError or
+      // kMitigationEngineStreamError
+      MitigationEngine::Action action;
     };
     std::string key;
     mutable std::optional<absl::Status> materialized_status;

data/src/core/ext/transport/chttp2/transport/hpack_parser.cc

@@ -40,6 +40,7 @@
 #include "src/core/lib/slice/slice.h"
 #include "src/core/lib/slice/slice_refcount.h"
 #include "src/core/lib/surface/validate_metadata.h"
+#include "src/core/mitigation_engine/mitigation_engine.h"
 #include "src/core/telemetry/call_tracer.h"
 #include "src/core/telemetry/stats.h"
 #include "src/core/telemetry/stats_data.h"
@@ -997,6 +998,21 @@ class HPackParser::Parser {
     auto value_slice = value.value.Take();
     const auto transport_size =
         key_string.size() + value.wire_size + hpack_constants::kEntryOverhead;
+    if (state_.mitigation_engine != nullptr) {
+      auto action = state_.mitigation_engine->EvaluateIncomingMetadata(
+          key_string, value_slice.as_string_view());
+      if (action.has_value()) {
+        if (action.value() == MitigationEngine::Action::kCloseConnection) {
+          input_->SetErrorAndStopParsing(
+              HpackParseResult::MitigationEngineError(
+                  key_string, MitigationEngine::Action::kCloseConnection));
+        } else if (action.value() == MitigationEngine::Action::kRejectRpc) {
+          input_->SetErrorAndContinueParsing(
+              HpackParseResult::MitigationEngineError(
+                  key_string, MitigationEngine::Action::kRejectRpc));
+        }
+      }
+    }
     auto md = grpc_metadata_batch::Parse(
         key_string, std::move(value_slice), state_.add_to_table, transport_size,
         [key_string, this](absl::string_view message, const Slice&) {
@@ -1092,13 +1108,15 @@ void HPackParser::BeginFrame(grpc_metadata_batch* metadata_buffer,
                              uint32_t metadata_size_soft_limit,
                              uint32_t metadata_size_hard_limit,
                              Boundary boundary, Priority priority,
-                             LogInfo log_info) {
+                             LogInfo log_info,
+                             MitigationEngine* mitigation_engine) {
   metadata_buffer_ = metadata_buffer;
   if (metadata_buffer != nullptr) {
     metadata_buffer->Set(GrpcStatusFromWire(), true);
   }
   boundary_ = boundary;
   priority_ = priority;
+  state_.mitigation_engine = mitigation_engine;
   state_.dynamic_table_updates_allowed = 2;
   state_.metadata_early_detection.SetLimits(
       /*soft_limit=*/metadata_size_soft_limit,
@@ -1133,6 +1151,20 @@ grpc_error_handle HPackParser::ParseInput(Input input, bool is_last,
                                           CallSpan* call_tracer) {
   ParseInputInner(&input);
   if (is_last && is_boundary()) {
+    if (state_.mitigation_engine != nullptr && metadata_buffer_ != nullptr) {
+      auto action = state_.mitigation_engine->EvaluateAllIncomingMetadata(
+          *metadata_buffer_);
+      if (action.has_value()) {
+        if (action.value() == MitigationEngine::Action::kCloseConnection) {
+          input.SetErrorAndStopParsing(HpackParseResult::MitigationEngineError(
+              /*key=*/"", MitigationEngine::Action::kCloseConnection));
+        } else if (action.value() == MitigationEngine::Action::kRejectRpc) {
+          input.SetErrorAndContinueParsing(
+              HpackParseResult::MitigationEngineError(
+                  /*key=*/"", MitigationEngine::Action::kRejectRpc));
+        }
+      }
+    }
     if (state_.metadata_early_detection.Reject(state_.frame_length,
                                                input.bitsrc())) {
       HandleMetadataSoftSizeLimitExceeded(&input);

data/src/core/ext/transport/chttp2/transport/hpack_parser.h

@@ -37,6 +37,7 @@
 #include "src/core/lib/iomgr/error.h"
 #include "src/core/lib/slice/slice.h"
 #include "src/core/lib/slice/slice_refcount.h"
+#include "src/core/mitigation_engine/mitigation_engine.h"
 #include "src/core/telemetry/call_tracer.h"
 #include "src/core/util/random_early_detection.h"
 #include "absl/random/bit_gen_ref.h"
@@ -99,7 +100,8 @@ class HPackParser {
   void BeginFrame(grpc_metadata_batch* metadata_buffer,
                   uint32_t metadata_size_soft_limit,
                   uint32_t metadata_size_hard_limit, Boundary boundary,
-                  Priority priority, LogInfo log_info);
+                  Priority priority, LogInfo log_info,
+                  MitigationEngine* mitigation_engine);
   // Start throwing away any received headers after parsing them.
   void StopBufferingFrame() { metadata_buffer_ = nullptr; }
   // Parse one slice worth of data
@@ -252,6 +254,7 @@ class HPackParser {
     uint8_t dynamic_table_updates_allowed;
     // Current parse state
     ParseState parse_state = ParseState::kTop;
+    MitigationEngine* mitigation_engine = nullptr;
     std::variant<const HPackTable::Memento*, Slice> key;
   };