From 04eb3515165f2e5d46a63d2ab40692deaf03136c Mon Sep 17 00:00:00 2001
From: kthui <18255193+kthui@users.noreply.github.com>
Date: Fri, 23 Feb 2024 14:48:04 -0800
Subject: [PATCH] Add test for enhance bound check for shm offset

---
 qa/L0_shared_memory/shared_memory_test.py | 29 ++++++++++++++++++++---
 qa/L0_shared_memory/test.sh               |  5 ++--
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py
index 6350dc2abe1..5a1384bc354 100755
--- a/qa/L0_shared_memory/shared_memory_test.py
+++ b/qa/L0_shared_memory/shared_memory_test.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2019-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2019-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
@@ -162,6 +162,7 @@ def _basic_inference(
         error_msg,
         big_shm_name="",
         big_shm_size=64,
+        shm_output_offset=0,
     ):
         input0_data = np.arange(start=0, stop=16, dtype=np.int32)
         input1_data = np.ones(shape=16, dtype=np.int32)
@@ -191,8 +192,8 @@ def _basic_inference(
         else:
             inputs[1].set_shared_memory("input1_data", 64)
 
-        outputs[0].set_shared_memory("output0_data", 64)
-        outputs[1].set_shared_memory("output1_data", 64)
+        outputs[0].set_shared_memory("output0_data", 64, offset=shm_output_offset)
+        outputs[1].set_shared_memory("output1_data", 64, offset=shm_output_offset)
 
         try:
             results = triton_client.infer(
@@ -323,6 +324,28 @@ def test_unregisterall(self):
             self.assertTrue(len(status_after.regions) == 0)
         self._cleanup_server(shm_handles)
 
+    def test_infer_offset_out_of_bound(self):
+        # Shared memory offset outside output region - Throws error
+        error_msg = []
+        shm_handles = self._configure_sever()
+        if _protocol == "http":
+            triton_client = httpclient.InferenceServerClient(_url, verbose=True)
+            offset = 18446744073709551584  # -0x20 in int64_t
+        else:
+            triton_client = grpcclient.InferenceServerClient(_url, verbose=True)
+            offset = 9223372036854775807  # grpc will throw if > int64
+        self._basic_inference(
+            shm_handles[0],
+            shm_handles[1],
+            shm_handles[2],
+            shm_handles[3],
+            error_msg,
+            shm_output_offset=offset,
+        )
+        self.assertEqual(len(error_msg), 1)
+        self.assertIn("Invalid offset for shared memory region", error_msg[0])
+        self._cleanup_server(shm_handles)
+
 
 if __name__ == "__main__":
     _protocol = os.environ.get("CLIENT_TYPE", "http")
diff --git a/qa/L0_shared_memory/test.sh b/qa/L0_shared_memory/test.sh
index e30a7dffa7b..30abfca545a 100755
--- a/qa/L0_shared_memory/test.sh
+++ b/qa/L0_shared_memory/test.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright 2019-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2019-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
@@ -49,7 +49,8 @@ for i in \
         test_register_after_inference \
         test_too_big_shm \
         test_mixed_raw_shm \
-        test_unregisterall; do
+        test_unregisterall \
+        test_infer_offset_out_of_bound; do
     for client_type in http grpc; do
         SERVER_ARGS="--model-repository=`pwd`/models --log-verbose=1 ${SERVER_ARGS_EXTRA}"
         SERVER_LOG="./$i.$client_type.server.log"