From 5630efe6e0d8b74b72a793678722f89112240f76 Mon Sep 17 00:00:00 2001
From: Yingge He <157551214+yinggeh@users.noreply.github.com>
Date: Tue, 27 Feb 2024 12:00:36 -0800
Subject: [PATCH] Match forward headers case insensitively. (#6889)

* Modify "header_forward_pattern" to match headers case-insensitively. Add unit tests.

* fix indentation

* fix pre-comiit errors

* Update doc

* Update copyright

* Add test case for "(?-i)", which disables regex case-insensitive mode.

* fix pre-commit

* Name each test. Remove support of disabling --http-header-forward-pattern case-insensitive mode on http python client.

* Update .md file.

* fix typo

* Reformat args.

* Fix pre-commit

* Fix test name issue.

* Fix pre-commit.

* Update md file and copyright.
---
 docs/protocol/extension_parameters.md |  8 +++++++-
 qa/L0_parameters/parameters_test.py   | 18 ++++++++++-------
 qa/L0_parameters/test.sh              | 29 +++++++++++++++++++--------
 src/command_line_parser.cc            | 11 ++++++++--
 4 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/docs/protocol/extension_parameters.md b/docs/protocol/extension_parameters.md
index 4cdb60cf38..14ed4d1dc5 100644
--- a/docs/protocol/extension_parameters.md
+++ b/docs/protocol/extension_parameters.md
@@ -1,5 +1,5 @@
 <!--
-# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -98,6 +98,12 @@ value. For example to forward all the headers that start with 'PREFIX_' from
 both HTTP and GRPC, you should add `--http-header-forward-pattern PREFIX_.*
 --grpc-header-forward-pattern PREFIX_.*` to your `tritonserver` command.
 
+By default, the regular expression pattern matches headers with case-insensitive
+mode according to the HTTP protocol. If you want to enforce case-sensitive mode,
+simplying adding the `(?-i)` prefix which turns off case-insensitive mode, e.g.
+`--http-header-forward-pattern (?-i)PREFIX_.*`. Note, headers sent through the
+Python HTTP client may be automatically lower-cased by internal client libraries.
+
 The forwarded headers can be accessed using the
 [Python](https://github.com/triton-inference-server/python_backend#inference-request-parameters)
 or C Backend APIs as inference request parameters.
diff --git a/qa/L0_parameters/parameters_test.py b/qa/L0_parameters/parameters_test.py
index 190326fc2c..a20d13c1eb 100755
--- a/qa/L0_parameters/parameters_test.py
+++ b/qa/L0_parameters/parameters_test.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -60,7 +60,8 @@ async def asyncSetUp(self):
         self.parameter_list.append({"key1": True, "key2": "value2"})
         self.parameter_list.append({"triton_": True, "key2": "value2"})
 
-        if TEST_HEADER == "1":
+        # Only "test_params" tests parameters without headers.
+        if TEST_HEADER != "test_params":
             self.headers = {
                 "header_1": "value_1",
                 "header_2": "value_2",
@@ -70,11 +71,14 @@ async def asyncSetUp(self):
             }
 
             # only these headers should be forwarded to the model.
-            self.expected_headers = {
-                "my_header_1": "my_value_1",
-                "my_header_2": "my_value_2",
-                "my_header_3": 'This is a "quoted" string with a backslash\ ',
-            }
+            if TEST_HEADER == "test_grpc_header_forward_pattern_case_sensitive":
+                self.expected_headers = {}
+            else:
+                self.expected_headers = {
+                    "my_header_1": "my_value_1",
+                    "my_header_2": "my_value_2",
+                    "my_header_3": 'This is a "quoted" string with a backslash\ ',
+                }
         else:
             self.headers = {}
             self.expected_headers = {}
diff --git a/qa/L0_parameters/test.sh b/qa/L0_parameters/test.sh
index 967ead15c7..c53b02d4b7 100755
--- a/qa/L0_parameters/test.sh
+++ b/qa/L0_parameters/test.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright 2023-2024, NVIDIA CORPORATION. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -53,15 +53,28 @@ mkdir -p "${MODELDIR}/ensemble/1"
 # TODO: Add support and testing for C++ client parameters:
 # https://jirasw.nvidia.com/browse/DLIS-4673
 
-RET=0
-for i in {0..1}; do
+all_tests=("test_params"
+           "test_headers"
+           "test_header_forward_pattern_case_insensitive"
+           "test_grpc_header_forward_pattern_case_sensitive")
 
+RET=0
+for i in "${all_tests[@]}"; do
   # TEST_HEADER is a parameter used by `parameters_test.py` that controls
   # whether the script will test for inclusion of headers in parameters or not.
-  if [ $i == 1 ]; then
-    SERVER_ARGS="--model-repository=${MODELDIR} --exit-timeout-secs=120 --grpc-header-forward-pattern my_header.* --http-header-forward-pattern my_header.*"
-  else
-    SERVER_ARGS="--model-repository=${MODELDIR} --exit-timeout-secs=120"
+  SERVER_ARGS="--model-repository=${MODELDIR} --exit-timeout-secs=120"
+  if [ "$i" == "test_headers" ]; then
+    SERVER_ARGS+=" --grpc-header-forward-pattern my_header.*"
+    SERVER_ARGS+=" --http-header-forward-pattern my_header.*"
+  elif [ "$i" == "test_header_forward_pattern_case_insensitive" ]; then
+    SERVER_ARGS+=" --grpc-header-forward-pattern MY_HEADER.*"
+    SERVER_ARGS+=" --http-header-forward-pattern MY_HEADER.*"
+  # NOTE: headers sent through the python HTTP client may be automatically
+  # lowercased by internal libraries like geventhttpclient, so we only test
+  # GRPC client for case-sensitivity here:
+  # https://github.com/geventhttpclient/geventhttpclient/blob/d1e14356c3b02099c879cf9b3bdb684a0cbd8bf5/src/geventhttpclient/header.py#L62-L63
+  elif [ "$i" == "test_grpc_header_forward_pattern_case_sensitive" ]; then
+    SERVER_ARGS+=" --grpc-header-forward-pattern (?-i)MY_HEADER.*"
   fi
   run_server
   if [ "$SERVER_PID" == "0" ]; then
@@ -71,7 +84,7 @@ for i in {0..1}; do
   fi
 
   set +e
-  TEST_HEADER=$i python3 $TEST_SCRIPT_PY >$CLIENT_LOG 2>&1
+  TEST_HEADER="$i" python3 $TEST_SCRIPT_PY >$CLIENT_LOG 2>&1
   if [ $? -ne 0 ]; then
       cat $CLIENT_LOG
       echo -e "\n***\n*** Test Failed\n***"
diff --git a/src/command_line_parser.cc b/src/command_line_parser.cc
index 5b490a45f6..1ba9fe2175 100644
--- a/src/command_line_parser.cc
+++ b/src/command_line_parser.cc
@@ -1235,6 +1235,11 @@ TritonParser::Parse(int argc, char** argv)
   triton::server::grpc::Options& lgrpc_options = lparams.grpc_options_;
 #endif  // TRITON_ENABLE_GRPC
 
+#if defined TRITON_ENABLE_HTTP || defined TRITON_ENABLE_GRPC
+  // According to HTTP specification header names are case-insensitive.
+  const std::string case_insensitive_prefix{"(?i)"};
+#endif  // TRITON_ENABLE_HTTP || TRITON_ENABLE_GRPC
+
 #ifdef TRITON_ENABLE_VERTEX_AI
   // Set different default value if specific flag is set
   {
@@ -1345,7 +1350,8 @@ TritonParser::Parse(int argc, char** argv)
           lparams.http_address_ = optarg;
           break;
         case OPTION_HTTP_HEADER_FORWARD_PATTERN:
-          lparams.http_forward_header_pattern_ = optarg;
+          lparams.http_forward_header_pattern_ =
+              std::move(case_insensitive_prefix + optarg);
           break;
           break;
         case OPTION_HTTP_THREAD_COUNT:
@@ -1484,7 +1490,8 @@ TritonParser::Parse(int argc, char** argv)
           break;
         }
         case OPTION_GRPC_HEADER_FORWARD_PATTERN:
-          lgrpc_options.forward_header_pattern_ = optarg;
+          lgrpc_options.forward_header_pattern_ =
+              std::move(case_insensitive_prefix + optarg);
           break;
 #endif  // TRITON_ENABLE_GRPC