fix: remove prohibited characters from error response (#62)

Author: ajaykarpur

src/sagemaker_inference/transformer.py

@@ -85,10 +85,12 @@ def handle_error(context, inference_exception, trace):
         Returns:
             str: The error message and stacktrace from the exception.
         """
-        context.set_response_status(
-            code=inference_exception.status_code, phrase=inference_exception.phrase
-        )
-        return ["{}\n{}".format(inference_exception.message, trace)]
+        phrase = utils.remove_crlf(inference_exception.phrase)
+        message = utils.remove_crlf(inference_exception.message)
+        stack_trace = utils.remove_crlf(trace)
+
+        context.set_response_status(code=inference_exception.status_code, phrase=phrase)
+        return ["{} {}".format(message, stack_trace)]
 
     def transform(self, data, context):
         """Take a request with input data, deserialize it, make a prediction, and return a

src/sagemaker_inference/utils.py

@@ -79,3 +79,24 @@ def parse_accept(accept):
             understand.
     """
     return accept.replace(" ", "").split(",")
+
+
+def remove_crlf(illegal_string):
+    """Removes characters prohibited by the MMS dependency Netty.
+
+    https://github.com/netty/netty/issues/8312
+
+    Args:
+        illegal_string: The string containing prohibited characters.
+
+    Returns:
+        str: The input string with the prohibited characters removed.
+    """
+    prohibited = ("\r", "\n")
+
+    sanitized_string = illegal_string
+
+    for character in prohibited:
+        sanitized_string = sanitized_string.replace(character, " ")
+
+    return sanitized_string

test/unit/test_utils.py

@@ -16,6 +16,7 @@
 from sagemaker_inference.utils import (
     parse_accept,
     read_file,
+    remove_crlf,
     retrieve_content_type_header,
     write_file,
 )
@@ -92,3 +93,10 @@ def test_content_type_header(content_type_key):
 def test_parse_accept(input, expected):
     actual = parse_accept(input)
     assert actual == expected
+
+
+def test_remove_crlf():
+    illegal_string = "test:\r\nstring"
+    sanitized_string = "test:  string"
+
+    assert sanitized_string == remove_crlf(illegal_string)