Skip to content

Commit

Permalink
Fix BufferReader.read_strings mishandling feeding from buffer (#72)
Browse files Browse the repository at this point in the history
* Fix BufferReader.read_strings mishandling feeding from buffer

This solves an issue where the code would incorrectly read past the
number of read characters that came from the network stream.
  • Loading branch information
mitsuhiko authored and xzkostyan committed Feb 14, 2019
1 parent 674996d commit 479dae1
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 0 deletions.
1 change: 1 addition & 0 deletions clickhouse_driver/bufferedreader.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,7 @@ def read_strings(self, n_items, decode=None):
buffer = self.buffer
buffer_view = self.buffer_view
current_buffer_size = self.current_buffer_size
position = min(position, current_buffer_size)
rv += buffer_view[0:position].tobytes()

else:
Expand Down
89 changes: 89 additions & 0 deletions tests/test_connect.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,16 @@
# coding: utf-8
import socket

from mock import patch
from io import BytesIO

from clickhouse_driver import errors
from clickhouse_driver.client import Client
from clickhouse_driver.protocol import ClientPacketTypes, ServerPacketTypes
from clickhouse_driver.bufferedreader import BufferedReader
from clickhouse_driver.writer import write_binary_str
from tests.testcase import BaseTestCase
from unittest import TestCase


class PacketsTestCase(BaseTestCase):
Expand Down Expand Up @@ -120,3 +125,87 @@ def side_effect(*args, **kwargs):

rv = self.client.execute('SELECT 1')
self.assertEqual(rv, [(1, )])


class FakeBufferedReader(BufferedReader):
def __init__(self, inputs, bufsize=128):
super(FakeBufferedReader, self).__init__(bufsize)
self._inputs = inputs
self._counter = 0

def read_into_buffer(self):
try:
value = self._inputs[self._counter]
except IndexError:
value = b''
else:
self._counter += 1

self.current_buffer_size = len(value)
self.buffer[:len(value)] = value

if self.current_buffer_size == 0:
raise EOFError('Unexpected EOF while reading bytes')


class TestBufferedReader(TestCase):

def test_corner_case_read(self):
rdr = FakeBufferedReader([
b'\x00' * 10,
b'\xff' * 10,
])

self.assertEqual(rdr.read(5), b'\x00' * 5)
self.assertEqual(rdr.read(10), b'\x00' * 5 + b'\xff' * 5)
self.assertEqual(rdr.read(5), b'\xff' * 5)

self.assertRaises(EOFError, rdr.read, 10)

def test_cornder_case_read_to_end_of_buffer(self):
rdr = FakeBufferedReader([
b'\x00' * 10,
b'\xff' * 10,
])

self.assertEqual(rdr.read(5), b'\x00' * 5)
self.assertEqual(rdr.read(5), b'\x00' * 5)
self.assertEqual(rdr.read(5), b'\xff' * 5)
self.assertEqual(rdr.read(5), b'\xff' * 5)

self.assertRaises(EOFError, rdr.read, 10)

def test_corner_case_exact_buffer(self):
rdr = FakeBufferedReader([
b'\x00' * 10,
b'\xff' * 10,
], bufsize=10)

self.assertEqual(rdr.read(5), b'\x00' * 5)
self.assertEqual(rdr.read(10), b'\x00' * 5 + b'\xff' * 5)
self.assertEqual(rdr.read(5), b'\xff' * 5)

def test_read_strings(self):
strings = [
u'Yoyodat' * 10,
u'Peter Maffay' * 10,
]

buf = BytesIO()
for name in strings:
write_binary_str(name, buf)
buf = buf.getvalue()

ref_values = [x.encode('utf-8') for x in strings]

for split in range(1, len(buf) - 1):
for split_2 in range(split + 1, len(buf) - 2):
assert buf[:split] + buf[split:split_2] + buf[split_2:] == buf
bufs = [
buf[:split],
buf[split:split_2],
buf[split_2:],
]
rdr = FakeBufferedReader(bufs, bufsize=4096)
read_values = rdr.read_strings(2)
self.assertEqual(repr(ref_values), repr(read_values))

0 comments on commit 479dae1

Please sign in to comment.