From 980bd76b0466d112cc1c5cdd918618993d53b061 Mon Sep 17 00:00:00 2001
From: GuanLuo <41310872+GuanLuo@users.noreply.github.com>
Date: Tue, 21 Mar 2023 10:27:29 -0700
Subject: [PATCH] Add GRPC option for restricted protocol access (#5397)

* [WIP] simple example

* WIP

* Implement GRPC restricted protocol option

* Add GRPC restricted protocol test. Fix bug

* Add documentation for endpoint configuration

* Add missing protocol

* Change option documentation based on change in detail

* Update command line option to restrict protocol access

* minor fix

* Update doc

* Address comment

* fix up

* Address comment
---
 .../inference_protocols.md                    |  52 +-
 docs/protocol/README.md                       |  12 +-
 qa/L0_grpc/python_unit_test.py                | 154 ++++++
 qa/L0_grpc/test.sh                            |  39 ++
 src/command_line_parser.cc                    |  95 +++-
 src/command_line_parser.h                     |   7 +
 src/grpc_server.cc                            | 473 ++++++++++++------
 src/grpc_server.h                             |  12 +
 8 files changed, 688 insertions(+), 156 deletions(-)
 create mode 100644 qa/L0_grpc/python_unit_test.py

diff --git a/docs/customization_guide/inference_protocols.md b/docs/customization_guide/inference_protocols.md
index 03b9d31e68..3c67371d1f 100644
--- a/docs/customization_guide/inference_protocols.md
+++ b/docs/customization_guide/inference_protocols.md
@@ -1,5 +1,5 @@
 <!--
-# Copyright 2018-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2018-2023, 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
@@ -114,6 +114,56 @@ These options can be used to configure the KeepAlive settings:
 
 For client-side documentation, see [Client-Side GRPC KeepAlive](https://github.com/triton-inference-server/client/blob/main/README.md#grpc-keepalive).
 
+#### Limit Endpoint Access (BETA)
+
+In some use cases, Triton users may want to restrict the access of the protocols on a given endpoint.
+For example, there can be need for two separate protocol groups that one exposes standard inference
+protocols for user access, while the other one exposes other extension protocols for administration
+usage and should not be accessible by non-admin user.
+
+The following option can be specified to declare an restricted protocol group:
+
+```
+--grpc-restricted-protocol=<protocol_1>,<protocol_2>,...:<restricted-key>=<restricted-value>
+```
+
+The option can be specified multiple times to specifies multiple groups of
+protocols with different restriction settings.
+
+* `protocols` : A comma-separated list of protocols to be included in this
+group. Note that currently a given protocol is not allowed to be included in
+multiple groups. The following protocols are currently recognized by all network
+protocol types mentioned above:
+
+  * `health` : Health endpoint defined for [HTTP/REST](https://github.com/kserve/kserve/blob/master/docs/predict-api/v2/required_api.md#health) and [GRPC](https://github.com/kserve/kserve/blob/master/docs/predict-api/v2/required_api.md#health-1). For GRPC endpoint, this value also exposes [GRPC health check protocol](https://github.com/triton-inference-server/common/blob/main/protobuf/health.proto).
+  * `metadata` : Server / model metadata endpoints defined for [HTTP/REST](https://github.com/kserve/kserve/blob/master/docs/predict-api/v2/required_api.md#server-metadata) and [GRPC](https://github.com/kserve/kserve/blob/master/docs/predict-api/v2/required_api.md#server-metadata-1).
+  * `inference` : Inference endpoints defined for [HTTP/REST](https://github.com/kserve/kserve/blob/master/docs/predict-api/v2/required_api.md#inference) and [GRPC](https://github.com/kserve/kserve/blob/master/docs/predict-api/v2/required_api.md#inference-1).
+  * `shared-memory` : [Shared-memory endpoint](https://github.com/triton-inference-server/server/blob/main/docs/protocol/extension_shared_memory.md).
+  * `model-config` : [Model configuration endpoint](https://github.com/triton-inference-server/server/blob/main/docs/protocol/extension_model_configuration.md).
+  * `model-repository` : [Model repository endpoint](https://github.com/triton-inference-server/server/blob/main/docs/protocol/extension_model_repository.md).
+  * `statistics` : [statistics endpoint](https://github.com/triton-inference-server/server/blob/main/docs/protocol/extension_statistics.md).
+  * `trace` : [trace endpoint](https://github.com/triton-inference-server/server/blob/main/docs/protocol/extension_trace.md).
+  * `logging` : [logging endpoint](https://github.com/triton-inference-server/server/blob/main/docs/protocol/extension_logging.md).
+
+* `restricted-key` : Key to determine the GRPC request header to be checked when a
+request to the protocol is received. The completed header will be in the form of
+`triton-grpc-protocol-<restricted-key>`
+
+* `restricted-value` : The value of the header to be matched in order to preceed
+the process of the specified protocols.
+
+#### Example
+
+To start server with a subset of protocols to be restricted in use case
+described above, the following command line arguments can be set to accept
+"standard inference" request without additional header and the rest of the
+protocols with `triton-grpc-protocol-<admin-key>=<admin-value>` specified in header:
+
+```
+tritonserver --grpc-restricted-protocol=shared-memory,model-config,model-repository,statistics,trace:<admin-key>=<admin-value> ...
+```
+
+
 ## In-Process Triton Server API
 
 The Triton Inference Server provides a backwards-compatible C API that
diff --git a/docs/protocol/README.md b/docs/protocol/README.md
index 808751957c..35e62e837e 100644
--- a/docs/protocol/README.md
+++ b/docs/protocol/README.md
@@ -36,17 +36,25 @@ plus several extensions that are defined in the following documents:
 
 - [Binary tensor data extension](./extension_binary_data.md)
 - [Classification extension](./extension_classification.md)
-- [Model configuration extension](./extension_model_configuration.md)
-- [Model repository extension](./extension_model_repository.md)
 - [Schedule policy extension](./extension_schedule_policy.md)
 - [Sequence extension](./extension_sequence.md)
 - [Shared-memory extension](./extension_shared_memory.md)
+- [Model configuration extension](./extension_model_configuration.md)
+- [Model repository extension](./extension_model_repository.md)
 - [Statistics extension](./extension_statistics.md)
 - [Trace extension](./extension_trace.md)
 - [Logging extension](./extension_logging.md)
 - [Parameters extension](./extension_parameters.md)
 
+Note that some extensions introduce new fields onto the inference protocols,
+and the other extensions define new protocols that Triton follows, please refer
+to the extension documents for detail.
+
 For the GRPC protocol, the [protobuf
 specification](https://github.com/triton-inference-server/common/blob/main/protobuf/grpc_service.proto)
 is also available. In addition, you can find the GRPC health checking protocol protobuf
 specification [here](https://github.com/triton-inference-server/common/blob/main/protobuf/health.proto).
+
+You can configure the Triton endpoints, which implement the protocols, to
+restrict access to some protocols and to control network settings, please refer
+to [protocol customization guide](https://github.com/triton-inference-server/server/blob/main/docs/customization_guide/inference_protocols.md#httprest-and-grpc-protocols) for detail.
diff --git a/qa/L0_grpc/python_unit_test.py b/qa/L0_grpc/python_unit_test.py
new file mode 100644
index 0000000000..db2a63f0a5
--- /dev/null
+++ b/qa/L0_grpc/python_unit_test.py
@@ -0,0 +1,154 @@
+#!/usr/bin/env python
+# Copyright 2023, 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
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import unittest
+import numpy as np
+import time
+
+import tritonclient.grpc as grpcclient
+from tritonclient.utils import InferenceServerException
+
+# For stream infer test
+from functools import partial
+import queue
+
+
+class UserData:
+
+    def __init__(self):
+        self._completed_requests = queue.Queue()
+
+
+def callback(user_data, result, error):
+    if error:
+        user_data._completed_requests.put(error)
+    else:
+        user_data._completed_requests.put(result)
+
+
+class RestrictedProtocolTest(unittest.TestCase):
+
+    def setUp(self):
+        self.client_ = grpcclient.InferenceServerClient(url="localhost:8001")
+        self.model_name_ = "simple"
+        self.prefix_ = "triton-grpc-protocol-"
+
+    # Other unspecified protocols should not be restricted
+    def test_sanity(self):
+        self.client_.get_inference_statistics("simple")
+        self.client_.get_inference_statistics(
+            "simple", headers={self.prefix_ + "infer-key": "infer-value"})
+
+    # health, infer, model repository protocols are restricted.
+    # health and infer expects "triton-grpc-restricted-infer-key : infer-value" header,
+    # model repository expected "triton-grpc-restricted-admin-key : admin-value".
+    def test_model_repository(self):
+        with self.assertRaisesRegex(InferenceServerException,
+                                    "This protocol is restricted"):
+            self.client_.unload_model(
+                self.model_name_,
+                headers={self.prefix_ + "infer-key": "infer-value"})
+        # Request go through and get actual transaction error
+        with self.assertRaisesRegex(
+                InferenceServerException,
+                "explicit model load / unload is not allowed"):
+            self.client_.unload_model(
+                self.model_name_,
+                headers={self.prefix_ + "admin-key": "admin-value"})
+
+    def test_health(self):
+        with self.assertRaisesRegex(InferenceServerException,
+                                    "This protocol is restricted"):
+            self.client_.is_server_live()
+        self.client_.is_server_live({self.prefix_ + "infer-key": "infer-value"})
+
+    def test_infer(self):
+        # setup
+        inputs = [
+            grpcclient.InferInput('INPUT0', [1, 16], "INT32"),
+            grpcclient.InferInput('INPUT1', [1, 16], "INT32")
+        ]
+        inputs[0].set_data_from_numpy(np.ones(shape=(1, 16), dtype=np.int32))
+        inputs[1].set_data_from_numpy(np.ones(shape=(1, 16), dtype=np.int32))
+
+        # This test only care if the request goes through
+        with self.assertRaisesRegex(InferenceServerException,
+                                    "This protocol is restricted"):
+            results = self.client_.infer(model_name=self.model_name_,
+                                         inputs=inputs,
+                                         headers={'test': '1'})
+        self.client_.infer(model_name=self.model_name_,
+                           inputs=inputs,
+                           headers={self.prefix_ + "infer-key": "infer-value"})
+
+    def test_stream_infer(self):
+        # setup
+        inputs = [
+            grpcclient.InferInput('INPUT0', [1, 16], "INT32"),
+            grpcclient.InferInput('INPUT1', [1, 16], "INT32")
+        ]
+        inputs[0].set_data_from_numpy(np.ones(shape=(1, 16), dtype=np.int32))
+        inputs[1].set_data_from_numpy(np.ones(shape=(1, 16), dtype=np.int32))
+        user_data = UserData()
+        # The server can't interfere with whether GRPC should create the stream,
+        # server will be notified after the stream is established and only
+        # until then be able to access metadata to decide whether to continue
+        # the stream.
+        # So on client side, it will always perceive that the stream is
+        # successfully created and can only check its health at a later time.
+        self.client_.start_stream(partial(callback, user_data),
+                                  headers={'test': '1'})
+        # wait for sufficient round-trip time
+        time.sleep(1)
+        with self.assertRaisesRegex(InferenceServerException,
+                                    "The stream is no longer in valid state"):
+            self.client_.async_stream_infer(model_name=self.model_name_,
+                                            inputs=inputs)
+        # callback should record error detail
+        self.assertFalse(user_data._completed_requests.empty())
+        with self.assertRaisesRegex(InferenceServerException,
+                                    "This protocol is restricted"):
+            raise user_data._completed_requests.get()
+
+        self.assertTrue(user_data._completed_requests.empty())
+
+        # Stop and start new stream with proper header
+        self.client_.stop_stream()
+        self.client_.start_stream(
+            partial(callback, user_data),
+            headers={self.prefix_ + "infer-key": "infer-value"})
+        self.client_.async_stream_infer(model_name=self.model_name_,
+                                        inputs=inputs)
+        # wait for response
+        time.sleep(1)
+        self.assertFalse(user_data._completed_requests.empty())
+        self.assertNotEqual(type(user_data._completed_requests.get()),
+                            InferenceServerException)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/qa/L0_grpc/test.sh b/qa/L0_grpc/test.sh
index 1788785560..e3286d847d 100644
--- a/qa/L0_grpc/test.sh
+++ b/qa/L0_grpc/test.sh
@@ -137,6 +137,7 @@ else
     SIMPLE_CUSTOM_ARGS_CLIENT=../clients/simple_grpc_custom_args_client
     CC_UNIT_TEST=../clients/cc_client_test
 fi
+PYTHON_UNIT_TEST=python_unit_test.py
 
 # Add string_dyna_sequence model to repo
 cp -r ${MODELDIR}/simple_dyna_sequence ${MODELDIR}/simple_string_dyna_sequence
@@ -571,6 +572,44 @@ fi
 kill $SERVER_PID
 wait $SERVER_PID
 
+# Repeated protocol, not allowed
+SERVER_ARGS="--model-repository=${MODELDIR} \
+             --grpc-restricted-protocol=model-repository,health:k1=v1 \
+             --grpc-restricted-protocol=metadata,health:k2=v2"
+run_server
+EXPECTED_MSG="protocol 'health' can not be specified in multiple config group"
+if [ "$SERVER_PID" != "0" ]; then
+    echo -e "\n***\n*** Expect fail to start $SERVER\n***"
+    kill $SERVER_PID
+    wait $SERVER_PID
+    RET=1
+elif [ `grep -c "${EXPECTED_MSG}" ${SERVER_LOG}` != "1" ]; then
+    echo -e "\n***\n*** Failed. Expected ${EXPECTED_MSG} to be found in log\n***"
+    cat $SERVER_LOG
+    RET=1
+fi
+
+# Test restricted protocols
+SERVER_ARGS="--model-repository=${MODELDIR} \
+             --grpc-restricted-protocol=model-repository:admin-key=admin-value \
+             --grpc-restricted-protocol=inference,health:infer-key=infer-value"
+run_server
+if [ "$SERVER_PID" == "0" ]; then
+    echo -e "\n***\n*** Failed to start $SERVER\n***"
+    cat $SERVER_LOG
+    exit 1
+fi
+set +e
+python $PYTHON_UNIT_TEST RestrictedProtocolTest > $CLIENT_LOG 2>&1
+if [ $? -ne 0 ]; then
+    cat $CLIENT_LOG
+    echo -e "\n***\n*** Python GRPC Restricted Protocol Test Failed\n***"
+    RET=1
+fi
+set -e
+kill $SERVER_PID
+wait $SERVER_PID
+
 if [ $RET -eq 0 ]; then
     echo -e "\n***\n*** Test Passed\n***"
 else
diff --git a/src/command_line_parser.cc b/src/command_line_parser.cc
index 8956decee4..725dd450d7 100644
--- a/src/command_line_parser.cc
+++ b/src/command_line_parser.cc
@@ -213,6 +213,25 @@ ParsePairOption(const std::string& arg, const std::string& delim_str)
   return {ParseOption<T1>(first_string), ParseOption<T2>(second_string)};
 }
 
+#ifdef TRITON_ENABLE_GRPC
+// Split 'options' by 'delim_str' and place splitted strings into a vector
+std::vector<std::string>
+SplitOptions(std::string options, const std::string& delim_str)
+{
+  std::vector<std::string> res;
+
+  int delim = options.find(delim_str);
+  while ((delim >= 0)) {
+    res.emplace_back(options.substr(0, delim));
+    options = options.substr(delim + delim_str.length());
+    delim = options.find(delim_str);
+  }
+  // include last element
+  res.emplace_back(options);
+  return res;
+}
+#endif  // TRITON_ENABLE_GRPC
+
 }  // namespace
 
 enum TritonOptionId {
@@ -256,6 +275,7 @@ enum TritonOptionId {
   OPTION_GRPC_ARG_HTTP2_MAX_PINGS_WITHOUT_DATA,
   OPTION_GRPC_ARG_HTTP2_MIN_RECV_PING_INTERVAL_WITHOUT_DATA_MS,
   OPTION_GRPC_ARG_HTTP2_MAX_PING_STRIKES,
+  OPTION_GRPC_RESTRICTED_PROTOCOL,
 #endif  // TRITON_ENABLE_GRPC
 #if defined(TRITON_ENABLE_SAGEMAKER)
   OPTION_ALLOW_SAGEMAKER,
@@ -439,6 +459,13 @@ std::vector<Option> TritonParser::recognized_options_
        "Maximum number of bad pings that the server will tolerate before "
        "sending an HTTP2 GOAWAY frame and closing the transport. Setting it to "
        "0 allows the server to accept any number of bad pings. Default is 2."},
+      {OPTION_GRPC_RESTRICTED_PROTOCOL, "grpc-restricted-protocol",
+       "<string>:<string>=<string>",
+       "Specify restricted GRPC protocol setting. The format of this "
+       "flag is --grpc-restricted-protocol=<protocols>,<key>=<value>. Where "
+       "<protocol> is a comma-separated list of protocols to be restricted. "
+       "<key> will be additional header key to be checked when a GRPC request "
+       "is received, and <value> is the value expected to be matched."},
 #endif  // TRITON_ENABLE_GRPC
 #if defined(TRITON_ENABLE_SAGEMAKER)
       {OPTION_ALLOW_SAGEMAKER, "allow-sagemaker", Option::ArgBool,
@@ -1147,6 +1174,19 @@ TritonParser::Parse(int argc, char** argv)
         lgrpc_options.keep_alive_.http2_max_ping_strikes_ =
             ParseOption<int>(optarg);
         break;
+      case OPTION_GRPC_RESTRICTED_PROTOCOL: {
+        const auto& parsed_tuple = ParseGrpcRestrictedProtocolOption(optarg);
+        const auto& protocols = SplitOptions(std::get<0>(parsed_tuple), ",");
+        const auto& key = std::get<1>(parsed_tuple);
+        const auto& value = std::get<2>(parsed_tuple);
+        grpc::ProtocolGroup pg;
+        for (const auto& p : protocols) {
+          pg.protocols_.emplace(p);
+        }
+        pg.restricted_key_ = std::make_pair(key, value);
+        lgrpc_options.protocol_groups_.emplace_back(pg);
+        break;
+      }
 #endif  // TRITON_ENABLE_GRPC
 
 #ifdef TRITON_ENABLE_METRICS
@@ -1320,7 +1360,6 @@ TritonParser::Parse(int argc, char** argv)
   // others which are not determined until after parsing.
   //
 
-  // [FIXME] check at parsing?
   if (lparams.control_mode_ != TRITONSERVER_MODEL_CONTROL_POLL) {
     lparams.repository_poll_secs_ = 0;
   }
@@ -1527,20 +1566,57 @@ TritonParser::ParseBackendConfigOption(const std::string& arg)
 }
 
 std::tuple<std::string, std::string, std::string>
-TritonParser::ParseHostPolicyOption(const std::string& arg)
+TritonParser::ParseGrpcRestrictedProtocolOption(const std::string& arg)
 {
-  // Format is "<policy_name>,<setting>=<value>"
-  int delim_name = arg.find(",");
-  int delim_setting = arg.find("=", delim_name + 1);
+  try {
+    return ParseGenericConfigOption(arg, ":", "=");
+  }
+  catch (const ParseException& pe) {
+    // catch and throw exception with option specific message
+    std::stringstream ss;
+    ss << "--grpc-restricted-protocol option format is '<config "
+          "name>:<setting>=<value>'. Got "
+       << arg << std::endl;
+    throw ParseException(ss.str());
+  }
+  // Should not reach here
+  return {};
+}
 
-  // Check for 2 semicolons
-  if ((delim_name < 0) || (delim_setting < 0)) {
+std::tuple<std::string, std::string, std::string>
+TritonParser::ParseHostPolicyOption(const std::string& arg)
+{
+  try {
+    return ParseGenericConfigOption(arg, ",", "=");
+  }
+  catch (const ParseException& pe) {
+    // catch and throw exception with option specific message
     std::stringstream ss;
     ss << "--host-policy option format is '<policy "
           "name>,<setting>=<value>'. Got "
        << arg << std::endl;
     throw ParseException(ss.str());
   }
+  // Should not reach here
+  return {};
+}
+
+std::tuple<std::string, std::string, std::string>
+TritonParser::ParseGenericConfigOption(
+    const std::string& arg, const std::string& first_delim,
+    const std::string& second_delim)
+{
+  // Format is "<string>,<string>=<string>"
+  int delim_name = arg.find(first_delim);
+  int delim_setting = arg.find(second_delim, delim_name + 1);
+
+  // Check for 2 semicolons
+  if ((delim_name < 0) || (delim_setting < 0)) {
+    std::stringstream ss;
+    ss << "option format is '<string>" << first_delim << "<string>"
+       << second_delim << "<string>'. Got " << arg << std::endl;
+    throw ParseException(ss.str());
+  }
 
   std::string name_string = arg.substr(0, delim_name);
   std::string setting_string =
@@ -1549,9 +1625,8 @@ TritonParser::ParseHostPolicyOption(const std::string& arg)
 
   if (name_string.empty() || setting_string.empty() || value_string.empty()) {
     std::stringstream ss;
-    ss << "--host-policy option format is '<policy "
-          "name>,<setting>=<value>'. Got "
-       << arg << std::endl;
+    ss << "option format is '<string>" << first_delim << "<string>"
+       << second_delim << "<string>'. Got " << arg << std::endl;
     throw ParseException(ss.str());
   }
 
diff --git a/src/command_line_parser.h b/src/command_line_parser.h
index c54df60b90..89a1de753b 100644
--- a/src/command_line_parser.h
+++ b/src/command_line_parser.h
@@ -272,9 +272,16 @@ class TritonParser {
       const std::string& arg);
   std::tuple<std::string, std::string, std::string> ParseHostPolicyOption(
       const std::string& arg);
+  std::tuple<std::string, std::string, std::string>
+  ParseGrpcRestrictedProtocolOption(const std::string& arg);
 #ifdef TRITON_ENABLE_TRACING
   TRITONSERVER_InferenceTraceLevel ParseTraceLevelOption(std::string arg);
 #endif  // TRITON_ENABLE_TRACING
+  // Helper function to parse option in
+  // "<string>[1st_delim]<string>[2nd_delim]<string>" format
+  std::tuple<std::string, std::string, std::string> ParseGenericConfigOption(
+      const std::string& arg, const std::string& first_delim,
+      const std::string& second_delim);
 
   static std::vector<Option> recognized_options_;
 };
diff --git a/src/grpc_server.cc b/src/grpc_server.cc
index 9d4539726f..e82e23be25 100644
--- a/src/grpc_server.cc
+++ b/src/grpc_server.cc
@@ -219,9 +219,11 @@ class CommonCallData : public Server::ICallData {
       const std::string& name, const uint64_t id,
       const StandardRegisterFunc OnRegister,
       const StandardCallbackFunc OnExecute, const bool async,
-      ::grpc::ServerCompletionQueue* cq)
+      ::grpc::ServerCompletionQueue* cq,
+      const std::pair<std::string, std::string>& restricted_kv)
       : name_(name), id_(id), OnRegister_(OnRegister), OnExecute_(OnExecute),
-        async_(async), cq_(cq), responder_(&ctx_), step_(Steps::START)
+        async_(async), cq_(cq), responder_(&ctx_), step_(Steps::START),
+        restricted_kv_(restricted_kv)
   {
     OnRegister_(&ctx_, &request_, &responder_, this);
     LOG_VERBOSE(1) << "Ready for RPC '" << name_ << "', " << id_;
@@ -244,6 +246,7 @@ class CommonCallData : public Server::ICallData {
   void Execute();
   void AddToCompletionQueue();
   void WriteResponse();
+  bool ExecutePrecondition();
 
   const std::string name_;
   const uint64_t id_;
@@ -263,6 +266,8 @@ class CommonCallData : public Server::ICallData {
   std::thread async_thread_;
 
   Steps step_;
+
+  std::pair<std::string, std::string> restricted_kv_{"", ""};
 };
 
 template <typename ResponderType, typename RequestType, typename ResponseType>
@@ -288,7 +293,7 @@ CommonCallData<ResponderType, RequestType, ResponseType>::Process(bool rpc_ok)
     // Start a new request to replace this one...
     if (!shutdown) {
       new CommonCallData<ResponderType, RequestType, ResponseType>(
-          name_, id_ + 1, OnRegister_, OnExecute_, async_, cq_);
+          name_, id_ + 1, OnRegister_, OnExecute_, async_, cq_, restricted_kv_);
     }
 
     if (!async_) {
@@ -316,7 +321,14 @@ template <typename ResponderType, typename RequestType, typename ResponseType>
 void
 CommonCallData<ResponderType, RequestType, ResponseType>::Execute()
 {
-  OnExecute_(request_, &response_, &status_);
+  if (ExecutePrecondition()) {
+    OnExecute_(request_, &response_, &status_);
+  } else {
+    status_ = ::grpc::Status(
+        ::grpc::StatusCode::UNAVAILABLE,
+        std::string("This protocol is restricted, expecting header '") +
+            restricted_kv_.first + "'");
+  }
   step_ = Steps::WRITEREADY;
 
   if (async_) {
@@ -327,6 +339,18 @@ CommonCallData<ResponderType, RequestType, ResponseType>::Execute()
   }
 }
 
+template <typename ResponderType, typename RequestType, typename ResponseType>
+bool
+CommonCallData<ResponderType, RequestType, ResponseType>::ExecutePrecondition()
+{
+  if (!restricted_kv_.first.empty()) {
+    const auto& metadata = ctx_.client_metadata();
+    const auto it = metadata.find(restricted_kv_.first);
+    return (it != metadata.end()) && (it->second == restricted_kv_.second);
+  }
+  return true;
+}
+
 template <typename ResponderType, typename RequestType, typename ResponseType>
 void
 CommonCallData<ResponderType, RequestType, ResponseType>::AddToCompletionQueue()
@@ -356,7 +380,9 @@ class CommonHandler : public Server::HandlerBase {
       TraceManager* trace_manager,
       inference::GRPCInferenceService::AsyncService* service,
       ::grpc::health::v1::Health::AsyncService* health_service,
-      ::grpc::ServerCompletionQueue* cq);
+      ::grpc::ServerCompletionQueue* cq,
+      std::map<std::string, std::pair<std::string, std::string>>
+          restricted_keys);
 
   // Descriptive name of of the handler.
   const std::string& Name() const { return name_; }
@@ -401,8 +427,13 @@ class CommonHandler : public Server::HandlerBase {
   ::grpc::health::v1::Health::AsyncService* health_service_;
   ::grpc::ServerCompletionQueue* cq_;
   std::unique_ptr<std::thread> thread_;
+  std::map<std::string, std::pair<std::string, std::string>> restricted_keys_;
+  static std::pair<std::string, std::string> empty_restricted_key_;
 };
 
+std::pair<std::string, std::string> CommonHandler::empty_restricted_key_{"",
+                                                                         ""};
+
 CommonHandler::CommonHandler(
     const std::string& name,
     const std::shared_ptr<TRITONSERVER_Server>& tritonserver,
@@ -410,10 +441,12 @@ CommonHandler::CommonHandler(
     TraceManager* trace_manager,
     inference::GRPCInferenceService::AsyncService* service,
     ::grpc::health::v1::Health::AsyncService* health_service,
-    ::grpc::ServerCompletionQueue* cq)
+    ::grpc::ServerCompletionQueue* cq,
+    std::map<std::string, std::pair<std::string, std::string>> restricted_keys)
     : name_(name), tritonserver_(tritonserver), shm_manager_(shm_manager),
       trace_manager_(trace_manager), service_(service),
-      health_service_(health_service), cq_(cq)
+      health_service_(health_service), cq_(cq),
+      restricted_keys_(restricted_keys)
 {
 }
 
@@ -534,11 +567,14 @@ CommonHandler::RegisterServerLive()
     TRITONSERVER_ErrorDelete(err);
   };
 
+  const auto it = restricted_keys_.find("health");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<inference::ServerLiveResponse>,
       inference::ServerLiveRequest, inference::ServerLiveResponse>(
       "ServerLive", 0, OnRegisterServerLive, OnExecuteServerLive,
-      false /* async */, cq_);
+      false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -568,11 +604,14 @@ CommonHandler::RegisterServerReady()
     TRITONSERVER_ErrorDelete(err);
   };
 
+  const auto it = restricted_keys_.find("health");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<inference::ServerReadyResponse>,
       inference::ServerReadyRequest, inference::ServerReadyResponse>(
       "ServerReady", 0, OnRegisterServerReady, OnExecuteServerReady,
-      false /* async */, cq_);
+      false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -613,13 +652,16 @@ CommonHandler::RegisterHealthCheck()
     TRITONSERVER_ErrorDelete(err);
   };
 
+  const auto it = restricted_keys_.find("health");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<
           ::grpc::health::v1::HealthCheckResponse>,
       ::grpc::health::v1::HealthCheckRequest,
       ::grpc::health::v1::HealthCheckResponse>(
       "Check", 0, OnRegisterHealthCheck, OnExecuteHealthCheck,
-      false /* async */, cq_);
+      false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -655,11 +697,14 @@ CommonHandler::RegisterModelReady()
     TRITONSERVER_ErrorDelete(err);
   };
 
+  const auto it = restricted_keys_.find("health");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<inference::ModelReadyResponse>,
       inference::ModelReadyRequest, inference::ModelReadyResponse>(
       "ModelReady", 0, OnRegisterModelReady, OnExecuteModelReady,
-      false /* async */, cq_);
+      false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -731,11 +776,14 @@ CommonHandler::RegisterServerMetadata()
         TRITONSERVER_ErrorDelete(err);
       };
 
+  const auto it = restricted_keys_.find("metadata");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<inference::ServerMetadataResponse>,
       inference::ServerMetadataRequest, inference::ServerMetadataResponse>(
       "ServerMetadata", 0, OnRegisterServerMetadata, OnExecuteServerMetadata,
-      false /* async */, cq_);
+      false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -897,11 +945,14 @@ CommonHandler::RegisterModelMetadata()
     TRITONSERVER_ErrorDelete(err);
   };
 
+  const auto it = restricted_keys_.find("metadata");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<inference::ModelMetadataResponse>,
       inference::ModelMetadataRequest, inference::ModelMetadataResponse>(
       "ModelMetadata", 0, OnRegisterModelMetadata, OnExecuteModelMetadata,
-      false /* async */, cq_);
+      false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -948,11 +999,14 @@ CommonHandler::RegisterModelConfig()
     TRITONSERVER_ErrorDelete(err);
   };
 
+  const auto it = restricted_keys_.find("model-config");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<inference::ModelConfigResponse>,
       inference::ModelConfigRequest, inference::ModelConfigResponse>(
       "ModelConfig", 0, OnRegisterModelConfig, OnExecuteModelConfig,
-      false /* async */, cq_);
+      false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -1245,11 +1299,14 @@ CommonHandler::RegisterModelStatistics()
 #endif
   };
 
+  const auto it = restricted_keys_.find("statistics");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<inference::ModelStatisticsResponse>,
       inference::ModelStatisticsRequest, inference::ModelStatisticsResponse>(
       "ModelStatistics", 0, OnRegisterModelStatistics, OnExecuteModelStatistics,
-      false /* async */, cq_);
+      false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -1456,10 +1513,14 @@ CommonHandler::RegisterTrace()
 #endif
   };
 
+  const auto it = restricted_keys_.find("trace");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<inference::TraceSettingResponse>,
       inference::TraceSettingRequest, inference::TraceSettingResponse>(
-      "Trace", 0, OnRegisterTrace, OnExecuteTrace, false /* async */, cq_);
+      "Trace", 0, OnRegisterTrace, OnExecuteTrace, false /* async */, cq_,
+      restricted_kv);
 }
 
 void
@@ -1662,11 +1723,14 @@ CommonHandler::RegisterLogging()
 #endif
   };
 
+  const auto it = restricted_keys_.find("logging");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<inference::LogSettingsResponse>,
       inference::LogSettingsRequest, inference::LogSettingsResponse>(
-      "Logging", 0, OnRegisterLogging, OnExecuteLogging, false /* async */,
-      cq_);
+      "Logging", 0, OnRegisterLogging, OnExecuteLogging, false /* async */, cq_,
+      restricted_kv);
 }
 
 void
@@ -1732,13 +1796,16 @@ CommonHandler::RegisterSystemSharedMemoryStatus()
         TRITONSERVER_ErrorDelete(err);
       };
 
+  const auto it = restricted_keys_.find("shared-memory");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<
           inference::SystemSharedMemoryStatusResponse>,
       inference::SystemSharedMemoryStatusRequest,
       inference::SystemSharedMemoryStatusResponse>(
       "SystemSharedMemoryStatus", 0, OnRegisterSystemSharedMemoryStatus,
-      OnExecuteSystemSharedMemoryStatus, false /* async */, cq_);
+      OnExecuteSystemSharedMemoryStatus, false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -1768,13 +1835,17 @@ CommonHandler::RegisterSystemSharedMemoryRegister()
         TRITONSERVER_ErrorDelete(err);
       };
 
+  const auto it = restricted_keys_.find("shared-memory");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<
           inference::SystemSharedMemoryRegisterResponse>,
       inference::SystemSharedMemoryRegisterRequest,
       inference::SystemSharedMemoryRegisterResponse>(
       "SystemSharedMemoryRegister", 0, OnRegisterSystemSharedMemoryRegister,
-      OnExecuteSystemSharedMemoryRegister, false /* async */, cq_);
+      OnExecuteSystemSharedMemoryRegister, false /* async */, cq_,
+      restricted_kv);
 }
 
 void
@@ -1808,13 +1879,17 @@ CommonHandler::RegisterSystemSharedMemoryUnregister()
         TRITONSERVER_ErrorDelete(err);
       };
 
+  const auto it = restricted_keys_.find("shared-memory");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<
           inference::SystemSharedMemoryUnregisterResponse>,
       inference::SystemSharedMemoryUnregisterRequest,
       inference::SystemSharedMemoryUnregisterResponse>(
       "SystemSharedMemoryUnregister", 0, OnRegisterSystemSharedMemoryUnregister,
-      OnExecuteSystemSharedMemoryUnregister, false /* async */, cq_);
+      OnExecuteSystemSharedMemoryUnregister, false /* async */, cq_,
+      restricted_kv);
 }
 
 void
@@ -1871,13 +1946,17 @@ CommonHandler::RegisterCudaSharedMemoryStatus()
         GrpcStatusUtil::Create(status, err);
         TRITONSERVER_ErrorDelete(err);
       };
+
+  const auto it = restricted_keys_.find("shared-memory");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<
           inference::CudaSharedMemoryStatusResponse>,
       inference::CudaSharedMemoryStatusRequest,
       inference::CudaSharedMemoryStatusResponse>(
       "CudaSharedMemoryStatus", 0, OnRegisterCudaSharedMemoryStatus,
-      OnExecuteCudaSharedMemoryStatus, false /* async */, cq_);
+      OnExecuteCudaSharedMemoryStatus, false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -1919,13 +1998,16 @@ CommonHandler::RegisterCudaSharedMemoryRegister()
         TRITONSERVER_ErrorDelete(err);
       };
 
+  const auto it = restricted_keys_.find("shared-memory");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<
           inference::CudaSharedMemoryRegisterResponse>,
       inference::CudaSharedMemoryRegisterRequest,
       inference::CudaSharedMemoryRegisterResponse>(
       "CudaSharedMemoryRegister", 0, OnRegisterCudaSharedMemoryRegister,
-      OnExecuteCudaSharedMemoryRegister, false /* async */, cq_);
+      OnExecuteCudaSharedMemoryRegister, false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -1959,13 +2041,17 @@ CommonHandler::RegisterCudaSharedMemoryUnregister()
         TRITONSERVER_ErrorDelete(err);
       };
 
+  const auto it = restricted_keys_.find("shared-memory");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<
           inference::CudaSharedMemoryUnregisterResponse>,
       inference::CudaSharedMemoryUnregisterRequest,
       inference::CudaSharedMemoryUnregisterResponse>(
       "CudaSharedMemoryUnregister", 0, OnRegisterCudaSharedMemoryUnregister,
-      OnExecuteCudaSharedMemoryUnregister, false /* async */, cq_);
+      OnExecuteCudaSharedMemoryUnregister, false /* async */, cq_,
+      restricted_kv);
 }
 
 void
@@ -2061,11 +2147,14 @@ CommonHandler::RegisterRepositoryIndex()
         TRITONSERVER_ErrorDelete(err);
       };
 
+  const auto it = restricted_keys_.find("model-repository");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<inference::RepositoryIndexResponse>,
       inference::RepositoryIndexRequest, inference::RepositoryIndexResponse>(
       "RepositoryIndex", 0, OnRegisterRepositoryIndex, OnExecuteRepositoryIndex,
-      false /* async */, cq_);
+      false /* async */, cq_, restricted_kv);
 }
 
 void
@@ -2170,12 +2259,15 @@ CommonHandler::RegisterRepositoryModelLoad()
         TRITONSERVER_ErrorDelete(err);
       };
 
+  const auto it = restricted_keys_.find("model-repository");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<inference::RepositoryModelLoadResponse>,
       inference::RepositoryModelLoadRequest,
       inference::RepositoryModelLoadResponse>(
       "RepositoryModelLoad", 0, OnRegisterRepositoryModelLoad,
-      OnExecuteRepositoryModelLoad, true /* async */, cq_);
+      OnExecuteRepositoryModelLoad, true /* async */, cq_, restricted_kv);
 }
 
 void
@@ -2236,13 +2328,16 @@ CommonHandler::RegisterRepositoryModelUnload()
         TRITONSERVER_ErrorDelete(err);
       };
 
+  const auto it = restricted_keys_.find("model-repository");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys_.end()) ? empty_restricted_key_ : it->second;
   new CommonCallData<
       ::grpc::ServerAsyncResponseWriter<
           inference::RepositoryModelUnloadResponse>,
       inference::RepositoryModelUnloadRequest,
       inference::RepositoryModelUnloadResponse>(
       "RepositoryModelUnload", 0, OnRegisterRepositoryModelUnload,
-      OnExecuteRepositoryModelUnload, true /* async */, cq_);
+      OnExecuteRepositoryModelUnload, true /* async */, cq_, restricted_kv);
 }
 
 //=========================================================================
@@ -2710,7 +2805,8 @@ class InferHandler : public Server::HandlerBase {
       const std::string& name,
       const std::shared_ptr<TRITONSERVER_Server>& tritonserver,
       ServiceType* service, ::grpc::ServerCompletionQueue* cq,
-      size_t max_state_bucket_count);
+      size_t max_state_bucket_count,
+      std::pair<std::string, std::string> restricted_kv);
   virtual ~InferHandler();
 
   // Descriptive name of of the handler.
@@ -2768,6 +2864,7 @@ class InferHandler : public Server::HandlerBase {
 
   virtual void StartNewRequest() = 0;
   virtual bool Process(State* state, bool rpc_ok) = 0;
+  bool ExecutePrecondition(InferHandler::State* state);
 
   const std::string name_;
   std::shared_ptr<TRITONSERVER_Server> tritonserver_;
@@ -2783,6 +2880,8 @@ class InferHandler : public Server::HandlerBase {
   // of creating a state for every new request.
   const size_t max_state_bucket_count_;
   std::vector<State*> state_bucket_;
+
+  std::pair<std::string, std::string> restricted_kv_;
 };
 
 template <
@@ -2793,9 +2892,11 @@ InferHandler<ServiceType, ServerResponderType, RequestType, ResponseType>::
         const std::string& name,
         const std::shared_ptr<TRITONSERVER_Server>& tritonserver,
         ServiceType* service, ::grpc::ServerCompletionQueue* cq,
-        size_t max_state_bucket_count)
+        size_t max_state_bucket_count,
+        std::pair<std::string, std::string> restricted_kv)
     : name_(name), tritonserver_(tritonserver), service_(service), cq_(cq),
-      max_state_bucket_count_(max_state_bucket_count)
+      max_state_bucket_count_(max_state_bucket_count),
+      restricted_kv_(restricted_kv)
 {
 }
 
@@ -2858,6 +2959,21 @@ InferHandler<
   LOG_VERBOSE(1) << "Thread exited for " << Name();
 }
 
+template <
+    typename ServiceType, typename ServerResponderType, typename RequestType,
+    typename ResponseType>
+bool
+InferHandler<ServiceType, ServerResponderType, RequestType, ResponseType>::
+    ExecutePrecondition(InferHandler::State* state)
+{
+  if (!restricted_kv_.first.empty()) {
+    const auto& metadata = state->context_->ctx_->client_metadata();
+    const auto it = metadata.find(restricted_kv_.first);
+    return (it != metadata.end()) && (it->second == restricted_kv_.second);
+  }
+  return true;
+}
+
 //
 // Infer utilities
 //
@@ -3886,8 +4002,11 @@ class ModelInferHandler
       const std::shared_ptr<SharedMemoryManager>& shm_manager,
       inference::GRPCInferenceService::AsyncService* service,
       ::grpc::ServerCompletionQueue* cq, size_t max_state_bucket_count,
-      grpc_compression_level compression_level)
-      : InferHandler(name, tritonserver, service, cq, max_state_bucket_count),
+      grpc_compression_level compression_level,
+      std::pair<std::string, std::string> restricted_kv)
+      : InferHandler(
+            name, tritonserver, service, cq, max_state_bucket_count,
+            restricted_kv),
         trace_manager_(trace_manager), shm_manager_(shm_manager),
         compression_level_(compression_level)
   {
@@ -3920,6 +4039,7 @@ class ModelInferHandler
   bool Process(State* state, bool rpc_ok) override;
 
  private:
+  void Execute(State* state);
   static void InferResponseComplete(
       TRITONSERVER_InferenceResponse* response, const uint32_t flags,
       void* userp);
@@ -3974,11 +4094,7 @@ ModelInferHandler::Process(InferHandler::State* state, bool rpc_ok)
     finished = true;
   }
 
-  const inference::ModelInferRequest& request = state->request_;
-  auto response_queue = state->response_queue_;
-
   if (state->step_ == Steps::START) {
-    TRITONSERVER_Error* err = nullptr;
 #ifdef TRITON_ENABLE_TRACING
     // Can't create trace as we don't know the model to be requested,
     // track timestamps in 'state'
@@ -3991,123 +4107,150 @@ ModelInferHandler::Process(InferHandler::State* state, bool rpc_ok)
       StartNewRequest();
     }
 
-    int64_t requested_model_version;
-    if (err == nullptr) {
-      err = GetModelVersionFromString(
-          request.model_version(), &requested_model_version);
-    }
-
-    if (err == nullptr) {
-      uint32_t txn_flags;
-      err = TRITONSERVER_ServerModelTransactionProperties(
-          tritonserver_.get(), request.model_name().c_str(),
-          requested_model_version, &txn_flags, nullptr /* voidp */);
-      if ((err == nullptr) && (txn_flags & TRITONSERVER_TXN_DECOUPLED) != 0) {
-        err = TRITONSERVER_ErrorNew(
-            TRITONSERVER_ERROR_UNSUPPORTED,
-            "ModelInfer RPC doesn't support models with decoupled "
-            "transaction policy");
-      }
-    }
-
-    // Create the inference request which contains all the
-    // input information needed for an inference.
-    TRITONSERVER_InferenceRequest* irequest = nullptr;
-    if (err == nullptr) {
-      err = TRITONSERVER_InferenceRequestNew(
-          &irequest, tritonserver_.get(), request.model_name().c_str(),
-          requested_model_version);
-    }
+    if (ExecutePrecondition(state)) {
+      Execute(state);
+    } else {
+      ::grpc::Status status = ::grpc::Status(
+          ::grpc::StatusCode::UNAVAILABLE,
+          std::string("This protocol is restricted, expecting header '") +
+              restricted_kv_.first + "'");
 
-    if (err == nullptr) {
-      err = SetInferenceRequestMetadata(irequest, request);
-    }
 
-    // Will be used to hold the serialized data in case explicit string
-    // tensors are present in the request.
-    std::list<std::string> serialized_data;
+#ifdef TRITON_ENABLE_TRACING
+      state->trace_timestamps_.emplace_back(
+          std::make_pair("GRPC_SEND_START", TraceManager::CaptureTimestamp()));
+#endif  // TRITON_ENABLE_TRACING
 
-    if (err == nullptr) {
-      err = InferGRPCToInput(
-          tritonserver_, shm_manager_, request, &serialized_data, irequest);
-    }
-    if (err == nullptr) {
-      err = InferAllocatorPayload<inference::ModelInferResponse>(
-          tritonserver_, shm_manager_, request, std::move(serialized_data),
-          response_queue, &state->alloc_payload_);
-    }
-    if (err == nullptr) {
-      err = TRITONSERVER_InferenceRequestSetReleaseCallback(
-          irequest, InferRequestComplete, nullptr /* request_release_userp */);
-    }
-    if (err == nullptr) {
-      err = TRITONSERVER_InferenceRequestSetResponseCallback(
-          irequest, allocator_,
-          &state->alloc_payload_ /* response_allocator_userp */,
-          InferResponseComplete, reinterpret_cast<void*>(state));
-    }
-    // Get request ID for logging in case of error.
-    const char* request_id = "";
-    if (irequest != nullptr) {
-      LOG_TRITONSERVER_ERROR(
-          TRITONSERVER_InferenceRequestId(irequest, &request_id),
-          "unable to retrieve request ID string");
+      state->step_ = COMPLETE;
+      state->context_->responder_->Finish(
+          inference::ModelInferResponse(), status, state);
     }
 
-    if (!strncmp(request_id, "", 1)) {
-      request_id = "<id_unknown>";
-    }
-    if (err == nullptr) {
-      TRITONSERVER_InferenceTrace* triton_trace = nullptr;
+  } else if (state->step_ == Steps::COMPLETE) {
 #ifdef TRITON_ENABLE_TRACING
-      state->trace_ =
-          std::move(trace_manager_->SampleTrace(request.model_name()));
-      if (state->trace_ != nullptr) {
-        triton_trace = state->trace_->trace_;
-      }
+    state->trace_timestamps_.emplace_back(
+        std::make_pair("GRPC_SEND_END", TraceManager::CaptureTimestamp()));
 #endif  // TRITON_ENABLE_TRACING
 
-      state->step_ = ISSUED;
-      err = TRITONSERVER_ServerInferAsync(
-          tritonserver_.get(), irequest, triton_trace);
+    state->step_ = Steps::FINISH;
+    finished = true;
+  }
+
+  return !finished;
+}
+
+void
+ModelInferHandler::Execute(InferHandler::State* state)
+{
+  TRITONSERVER_Error* err = nullptr;
+  const inference::ModelInferRequest& request = state->request_;
+  auto response_queue = state->response_queue_;
+  int64_t requested_model_version;
+  if (err == nullptr) {
+    err = GetModelVersionFromString(
+        request.model_version(), &requested_model_version);
+  }
+
+  if (err == nullptr) {
+    uint32_t txn_flags;
+    err = TRITONSERVER_ServerModelTransactionProperties(
+        tritonserver_.get(), request.model_name().c_str(),
+        requested_model_version, &txn_flags, nullptr /* voidp */);
+    if ((err == nullptr) && (txn_flags & TRITONSERVER_TXN_DECOUPLED) != 0) {
+      err = TRITONSERVER_ErrorNew(
+          TRITONSERVER_ERROR_UNSUPPORTED,
+          "ModelInfer RPC doesn't support models with decoupled "
+          "transaction policy");
     }
+  }
 
-    // If not error then state->step_ == ISSUED and inference request
-    // has initiated... completion callback will transition to
-    // COMPLETE. If error go immediately to COMPLETE.
-    if (err != nullptr) {
-      LOG_VERBOSE(1) << "[request id: " << request_id << "] "
-                     << "Infer failed: " << TRITONSERVER_ErrorMessage(err);
+  // Create the inference request which contains all the
+  // input information needed for an inference.
+  TRITONSERVER_InferenceRequest* irequest = nullptr;
+  if (err == nullptr) {
+    err = TRITONSERVER_InferenceRequestNew(
+        &irequest, tritonserver_.get(), request.model_name().c_str(),
+        requested_model_version);
+  }
 
-      LOG_TRITONSERVER_ERROR(
-          TRITONSERVER_InferenceRequestDelete(irequest),
-          "deleting GRPC inference request");
+  if (err == nullptr) {
+    err = SetInferenceRequestMetadata(irequest, request);
+  }
 
-      ::grpc::Status status;
-      GrpcStatusUtil::Create(&status, err);
-      TRITONSERVER_ErrorDelete(err);
+  // Will be used to hold the serialized data in case explicit string
+  // tensors are present in the request.
+  std::list<std::string> serialized_data;
 
-      inference::ModelInferResponse error_response;
+  if (err == nullptr) {
+    err = InferGRPCToInput(
+        tritonserver_, shm_manager_, request, &serialized_data, irequest);
+  }
+  if (err == nullptr) {
+    err = InferAllocatorPayload<inference::ModelInferResponse>(
+        tritonserver_, shm_manager_, request, std::move(serialized_data),
+        response_queue, &state->alloc_payload_);
+  }
+  if (err == nullptr) {
+    err = TRITONSERVER_InferenceRequestSetReleaseCallback(
+        irequest, InferRequestComplete, nullptr /* request_release_userp */);
+  }
+  if (err == nullptr) {
+    err = TRITONSERVER_InferenceRequestSetResponseCallback(
+        irequest, allocator_,
+        &state->alloc_payload_ /* response_allocator_userp */,
+        InferResponseComplete, reinterpret_cast<void*>(state));
+  }
+  // Get request ID for logging in case of error.
+  const char* request_id = "";
+  if (irequest != nullptr) {
+    LOG_TRITONSERVER_ERROR(
+        TRITONSERVER_InferenceRequestId(irequest, &request_id),
+        "unable to retrieve request ID string");
+  }
 
+  if (!strncmp(request_id, "", 1)) {
+    request_id = "<id_unknown>";
+  }
+  if (err == nullptr) {
+    TRITONSERVER_InferenceTrace* triton_trace = nullptr;
 #ifdef TRITON_ENABLE_TRACING
-      state->trace_timestamps_.emplace_back(
-          std::make_pair("GRPC_SEND_START", TraceManager::CaptureTimestamp()));
+    state->trace_ =
+        std::move(trace_manager_->SampleTrace(request.model_name()));
+    if (state->trace_ != nullptr) {
+      triton_trace = state->trace_->trace_;
+    }
 #endif  // TRITON_ENABLE_TRACING
 
-      state->step_ = COMPLETE;
-      state->context_->responder_->Finish(error_response, status, state);
-    }
-  } else if (state->step_ == Steps::COMPLETE) {
+    state->step_ = ISSUED;
+    err = TRITONSERVER_ServerInferAsync(
+        tritonserver_.get(), irequest, triton_trace);
+  }
+
+  // If not error then state->step_ == ISSUED and inference request
+  // has initiated... completion callback will transition to
+  // COMPLETE. If error go immediately to COMPLETE.
+  if (err != nullptr) {
+    LOG_VERBOSE(1) << "[request id: " << request_id << "] "
+                   << "Infer failed: " << TRITONSERVER_ErrorMessage(err);
+
+    LOG_TRITONSERVER_ERROR(
+        TRITONSERVER_InferenceRequestDelete(irequest),
+        "deleting GRPC inference request");
+
+    ::grpc::Status status;
+    GrpcStatusUtil::Create(&status, err);
+    TRITONSERVER_ErrorDelete(err);
+
+    inference::ModelInferResponse error_response;
+
 #ifdef TRITON_ENABLE_TRACING
     state->trace_timestamps_.emplace_back(
-        std::make_pair("GRPC_SEND_END", TraceManager::CaptureTimestamp()));
+        std::make_pair("GRPC_SEND_START", TraceManager::CaptureTimestamp()));
 #endif  // TRITON_ENABLE_TRACING
 
-    state->step_ = Steps::FINISH;
-    finished = true;
+    state->step_ = COMPLETE;
+    state->context_->responder_->Finish(error_response, status, state);
   }
-
-  return !finished;
 }
 
 void
@@ -4281,8 +4424,11 @@ class ModelStreamInferHandler
       const std::shared_ptr<SharedMemoryManager>& shm_manager,
       inference::GRPCInferenceService::AsyncService* service,
       ::grpc::ServerCompletionQueue* cq, size_t max_state_bucket_count,
-      grpc_compression_level compression_level)
-      : InferHandler(name, tritonserver, service, cq, max_state_bucket_count),
+      grpc_compression_level compression_level,
+      std::pair<std::string, std::string> restricted_kv)
+      : InferHandler(
+            name, tritonserver, service, cq, max_state_bucket_count,
+            restricted_kv),
         trace_manager_(trace_manager), shm_manager_(shm_manager),
         compression_level_(compression_level)
   {
@@ -4372,11 +4518,23 @@ ModelStreamInferHandler::Process(InferHandler::State* state, bool rpc_ok)
     // Start a new request to replace this one...
     StartNewRequest();
 
-    // Since this is the start of a connection, 'state' hasn't been
-    // used yet so use it to read a request off the connection.
-    state->context_->step_ = Steps::READ;
-    state->step_ = Steps::READ;
-    state->context_->responder_->Read(&state->request_, state);
+    if (ExecutePrecondition(state)) {
+      // Since this is the start of a connection, 'state' hasn't been
+      // used yet so use it to read a request off the connection.
+      state->context_->step_ = Steps::READ;
+      state->step_ = Steps::READ;
+      state->context_->responder_->Read(&state->request_, state);
+    } else {
+      // Precondition is not satisfied, cancel the stream
+      state->context_->step_ = Steps::COMPLETE;
+      state->step_ = Steps::COMPLETE;
+      ::grpc::Status status = ::grpc::Status(
+          ::grpc::StatusCode::UNAVAILABLE,
+          std::string("This protocol is restricted, expecting header '") +
+              restricted_kv_.first + "'");
+      state->context_->responder_->Finish(status, state);
+      return !finished;
+    }
 
   } else if (state->step_ == Steps::READ) {
     TRITONSERVER_Error* err = nullptr;
@@ -4906,18 +5064,41 @@ Server::Server(
   model_infer_cq_ = builder_.AddCompletionQueue();
   model_stream_infer_cq_ = builder_.AddCompletionQueue();
 
+  // Read and set restriction for each protocol specified
+  // map from protocol name to a pair of header to look for and the key
+  std::map<std::string, std::pair<std::string, std::string>> restricted_keys;
+  for (const auto& pg : options.protocol_groups_) {
+    for (const auto& p : pg.protocols_) {
+      if (restricted_keys.find(p) != restricted_keys.end()) {
+        throw std::invalid_argument(
+            std::string("protocol '") + p +
+            "' can not be "
+            "specified in multiple config group");
+      }
+      const auto header = std::string(kRestrictedProtocolHeaderTemplate) +
+                          pg.restricted_key_.first;
+      restricted_keys[p] = std::make_pair(header, pg.restricted_key_.second);
+    }
+  }
+
   // A common Handler for other non-inference requests
   common_handler_.reset(new CommonHandler(
       "CommonHandler", tritonserver_, shm_manager_, trace_manager_, &service_,
-      &health_service_, common_cq_.get()));
+      &health_service_, common_cq_.get(), restricted_keys));
 
+  // [FIXME] "register" logic is different for infer
   // Handler for model inference requests.
+  const auto it = restricted_keys.find("inference");
+  std::pair<std::string, std::string> restricted_kv =
+      (it == restricted_keys.end())
+          ? std::pair<std::string, std::string>{"", ""}
+          : it->second;
   for (int i = 0; i < REGISTER_GRPC_INFER_THREAD_COUNT; ++i) {
     model_infer_handlers_.emplace_back(new ModelInferHandler(
         "ModelInferHandler", tritonserver_, trace_manager_, shm_manager_,
         &service_, model_infer_cq_.get(),
         options.infer_allocation_pool_size_ /* max_state_bucket_count */,
-        options.infer_compression_level_));
+        options.infer_compression_level_, restricted_kv));
   }
 
   // Handler for streaming inference requests. Keeps one handler for streaming
@@ -4926,7 +5107,7 @@ Server::Server(
       "ModelStreamInferHandler", tritonserver_, trace_manager_, shm_manager_,
       &service_, model_stream_infer_cq_.get(),
       options.infer_allocation_pool_size_ /* max_state_bucket_count */,
-      options.infer_compression_level_));
+      options.infer_compression_level_, restricted_kv));
 }
 
 Server::~Server()
@@ -4943,8 +5124,14 @@ Server::Create(
 {
   const std::string addr = server_options.socket_.address_ + ":" +
                            std::to_string(server_options.socket_.port_);
-  server->reset(
-      new Server(tritonserver, trace_manager, shm_manager, server_options));
+  try {
+    server->reset(
+        new Server(tritonserver, trace_manager, shm_manager, server_options));
+  }
+  catch (const std::invalid_argument& pe) {
+    return TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INVALID_ARG, pe.what());
+    ;
+  }
 
   return nullptr;  // success
 }
diff --git a/src/grpc_server.h b/src/grpc_server.h
index a27d3d9206..f4d60cbcff 100644
--- a/src/grpc_server.h
+++ b/src/grpc_server.h
@@ -26,6 +26,7 @@
 #pragma once
 
 #include <grpc++/grpc++.h>
+#include <vector>
 #include "grpc_service.grpc.pb.h"
 #include "health.grpc.pb.h"
 #include "shared_memory_manager.h"
@@ -34,6 +35,10 @@
 
 namespace triton { namespace server { namespace grpc {
 
+// GRPC uses HTTP2 which requires header to be in lowercase, so the Triton
+// specific header that may be set for GRPC is defined to be all lowercases
+constexpr char kRestrictedProtocolHeaderTemplate[] = "triton-grpc-protocol-";
+
 struct SocketOptions {
   std::string address_{"0.0.0.0"};
   int32_t port_{8001};
@@ -63,6 +68,12 @@ struct KeepAliveOptions {
   int http2_max_ping_strikes_{2};
 };
 
+struct ProtocolGroup {
+  std::string name_{""};
+  std::set<std::string> protocols_{};
+  std::pair<std::string, std::string> restricted_key_{"", ""};
+};
+
 struct Options {
   SocketOptions socket_;
   SslOptions ssl_;
@@ -73,6 +84,7 @@ struct Options {
   // requests doesn't exceed this value there will be no
   // allocation/deallocation of request/response objects.
   int infer_allocation_pool_size_{8};
+  std::vector<ProtocolGroup> protocol_groups_{};
 };
 
 class Server {