From 13bf2ee27d0cb0fcfb603b7a0c9adb57e89161cb Mon Sep 17 00:00:00 2001
From: GuanLuo <41310872+GuanLuo@users.noreply.github.com>
Date: Thu, 24 Mar 2022 11:13:39 -0700
Subject: [PATCH] Fix inference header handling for Vertex and SageMaker
 (#4099)

* Fix inference header handling for Vertex and SageMaker

* Fix copyright
---
 src/sagemaker_server.cc | 7 ++++---
 src/vertex_ai_server.cc | 7 ++++---
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/sagemaker_server.cc b/src/sagemaker_server.cc
index e69216c118..63a24eb0e8 100644
--- a/src/sagemaker_server.cc
+++ b/src/sagemaker_server.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2021-2022, 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
@@ -34,8 +34,9 @@ TRITONSERVER_Error*
 SagemakerAPIServer::GetInferenceHeaderLength(
     evhtp_request_t* req, int32_t content_length, size_t* header_length)
 {
-  // Check mime type and set inference header length. If missing set to 0
-  *header_length = 0;
+  // Check mime type and set inference header length.
+  // Set to content length in case that it is not specified
+  *header_length = content_length;
   const char* content_type_c_str =
       evhtp_kv_find(req->headers_in, kContentTypeHeader);
   if (content_type_c_str != NULL) {
diff --git a/src/vertex_ai_server.cc b/src/vertex_ai_server.cc
index a7f7ee9407..cbc2439411 100644
--- a/src/vertex_ai_server.cc
+++ b/src/vertex_ai_server.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2021-2022, 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
@@ -52,8 +52,9 @@ TRITONSERVER_Error*
 VertexAiAPIServer::GetInferenceHeaderLength(
     evhtp_request_t* req, int32_t content_length, size_t* header_length)
 {
-  // Check mime type and set inference header length. If missing set to 0
-  *header_length = 0;
+  // Check mime type and set inference header length.
+  // Set to content length in case that it is not specified
+  *header_length = content_length;
   const char* content_type_c_str =
       evhtp_kv_find(req->headers_in, kContentTypeHeader);
   if (content_type_c_str != NULL) {