diff --git a/pylsp_jsonrpc/streams.py b/pylsp_jsonrpc/streams.py index 40048a9..7e7427a 100644 --- a/pylsp_jsonrpc/streams.py +++ b/pylsp_jsonrpc/streams.py @@ -53,10 +53,11 @@ def _read_message(self): if not line: return None - content_length = self._content_length(line) + content_length = None - # Blindly consume all header lines while line and line.strip(): + if content_length is None: + content_length = self._content_length(line) line = self._rfile.readline() if not line: @@ -68,8 +69,8 @@ def _read_message(self): @staticmethod def _content_length(line): """Extract the content length from an input line.""" - if line.startswith(b'Content-Length: '): - _, value = line.split(b'Content-Length: ') + header, separator, value = line.partition(b': ') + if separator and header.lower() == b'content-length': value = value.strip() try: return int(value) diff --git a/test/test_streams.py b/test/test_streams.py index 14fe1bb..5e8df64 100644 --- a/test/test_streams.py +++ b/test/test_streams.py @@ -51,6 +51,42 @@ def test_reader(rfile, reader): }) +def test_reader_accepts_content_length_in_any_header_position_and_case(rfile, reader): + first_body = b'{"jsonrpc":"2.0","method":"first"}' + second_body = b'{"jsonrpc":"2.0","method":"second"}' + third_body = b'{"jsonrpc":"2.0","method":"third"}' + first_frame = b''.join(( + b'Content-Type: application/vscode-jsonrpc; charset=utf-8\r\n', + b'Content-Length: ', + str(len(first_body)).encode('ascii'), + b'\r\n\r\n', + first_body, + )) + second_frame = b''.join(( + b'cOnTeNt-LeNgTh: ', + str(len(second_body)).encode('ascii'), + b'\r\n\r\n', + second_body, + )) + third_frame = b''.join(( + b'Content-Length: ', + str(len(third_body)).encode('ascii'), + b'\r\n\r\n', + third_body, + )) + rfile.write(first_frame + second_frame + third_frame) + rfile.seek(0) + + consumer = mock.Mock() + reader.listen(consumer) + + assert consumer.call_args_list == [ + mock.call({'jsonrpc': '2.0', 'method': 'first'}), + mock.call({'jsonrpc': '2.0', 'method': 'second'}), + mock.call({'jsonrpc': '2.0', 'method': 'third'}), + ] + + def test_reader_bad_message(rfile, reader): rfile.write(b'Hello world') rfile.seek(0)