Fix for multibyte inputs to custom search commands

Author: amysutedja

splunklib/searchcommands/search_command.py

@@ -852,15 +852,23 @@ def _execute(self, ifile, process):
     @staticmethod
     def _read_chunk(ifile):
         # noinspection PyBroadException
+        if six.PY2:
+            istream = ifile
+        else:
+            try:
+                istream = ifile.buffer
+            except AttributeError as error:
+                raise RuntimeError('Failed to get underlying buffer: {}'.format(error))
+
         try:
-            header = ifile.readline()
+            header = istream.readline()
         except Exception as error:
             raise RuntimeError('Failed to read transport header: {}'.format(error))
 
         if not header:
             return None
 
-        match = SearchCommand._header.match(header)
+        match = SearchCommand._header.match(six.ensure_str(header))
 
         if match is None:
             raise RuntimeError('Failed to parse transport header: {}'.format(header))
@@ -870,14 +878,14 @@ def _read_chunk(ifile):
         body_length = int(body_length)
 
         try:
-            metadata = ifile.read(metadata_length)
+            metadata = istream.read(metadata_length)
         except Exception as error:
             raise RuntimeError('Failed to read metadata of length {}: {}'.format(metadata_length, error))
 
         decoder = MetadataDecoder()
 
         try:
-            metadata = decoder.decode(metadata)
+            metadata = decoder.decode(six.ensure_str(metadata))
         except Exception as error:
             raise RuntimeError('Failed to parse metadata of length {}: {}'.format(metadata_length, error))
 
@@ -887,11 +895,11 @@ def _read_chunk(ifile):
         body = ""
         try:
             if body_length > 0:
-                body = ifile.read(body_length)
+                body = istream.read(body_length)
         except Exception as error:
             raise RuntimeError('Failed to read body of length {}: {}'.format(body_length, error))
 
-        return metadata, body
+        return metadata, six.ensure_str(body)
 
     _header = re.compile(r'chunked\s+1.0\s*,\s*(\d+)\s*,\s*(\d+)\s*\n')
 

tests/searchcommands/test_search_command.py

@@ -33,8 +33,21 @@
 import os
 import re
 
+from io import TextIOWrapper
+
 import pytest
 
+def build_command_input(getinfo_metadata, execute_metadata, execute_body):
+    header = ('chunked 1.0,{},0\n{}'.format(len(six.ensure_binary(getinfo_metadata)), getinfo_metadata) +
+              'chunked 1.0,{},{}\n{}{}'.format(len(six.ensure_binary(execute_metadata)), len(six.ensure_binary(execute_body)), execute_metadata, execute_body))
+
+    ifile = BytesIO(six.ensure_binary(header))
+
+    if not six.PY2:
+        ifile = TextIOWrapper(ifile)
+
+    return ifile
+
 @Configuration()
 class TestCommand(SearchCommand):
 
@@ -428,11 +441,9 @@ def test_process_scpv2(self):
             show_configuration=('true' if show_configuration is True else 'false'))
 
         execute_metadata = '{"action":"execute","finished":true}'
-        execute_body = 'test\r\ndata\r\n'
+        execute_body = 'test\r\ndata\r\n测试\r\n'
 
-        ifile = StringIO(
-            'chunked 1.0,{},0\n{}'.format(len(getinfo_metadata), getinfo_metadata) +
-            'chunked 1.0,{},{}\n{}{}'.format(len(execute_metadata), len(execute_body), execute_metadata, execute_body))
+        ifile = build_command_input(getinfo_metadata, execute_metadata, execute_body)
 
         command = TestCommand()
         result = BytesIO()
@@ -455,12 +466,17 @@ def test_process_scpv2(self):
         self.assertEqual(command.required_option_1, 'value_1')
         self.assertEqual(command.required_option_2, 'value_2')
 
-        self.assertEqual(
+        expected = (
             'chunked 1.0,68,0\n'
             '{"inspector":{"messages":[["INFO","test command configuration: "]]}}\n'
-            'chunked 1.0,17,23\n'
+            'chunked 1.0,17,32\n'
             '{"finished":true}test,__mv_test\r\n'
-            'data,\r\n',
+            'data,\r\n'
+            '测试,\r\n'
+        )
+
+        self.assertEqual(
+            expected,
             result.getvalue().decode('utf-8'))
 
         self.assertEqual(command.protocol_version, 2)
@@ -620,11 +636,9 @@ def test_process_scpv2(self):
             show_configuration=show_configuration)
 
         execute_metadata = '{"action":"execute","finished":true}'
-        execute_body = 'test\r\ndata\r\n'
+        execute_body = 'test\r\ndata\r\n测试\r\n'
 
-        ifile = StringIO(
-            'chunked 1.0,{},0\n{}'.format(len(getinfo_metadata), getinfo_metadata) +
-            'chunked 1.0,{},{}\n{}{}'.format(len(execute_metadata), len(execute_body), execute_metadata, execute_body))
+        ifile = build_command_input(getinfo_metadata, execute_metadata, execute_body)
 
         command = TestCommand()
         result = BytesIO()
@@ -666,11 +680,9 @@ def test_process_scpv2(self):
             show_configuration=('true' if show_configuration is True else 'false'))
 
         execute_metadata = '{"action":"execute","finished":true}'
-        execute_body = 'action\r\nraise_exception\r\n'
+        execute_body = 'action\r\nraise_exception\r\n测试\r\n'
 
-        ifile = StringIO(
-            'chunked 1.0,{},0\n{}'.format(len(getinfo_metadata), getinfo_metadata) +
-            'chunked 1.0,{},{}\n{}{}'.format(len(execute_metadata), len(execute_body), execute_metadata, execute_body))
+        ifile = build_command_input(getinfo_metadata, execute_metadata, execute_body)
 
         command = TestCommand()
         result = BytesIO()