diff --git a/glance/store/vmware_datastore.py b/glance/store/vmware_datastore.py index 62672ecadf..b5d45783ef 100644 --- a/glance/store/vmware_datastore.py +++ b/glance/store/vmware_datastore.py @@ -108,16 +108,40 @@ def http_response_iterator(conn, response, size): class _Reader(object): - def __init__(self, data, checksum): + def __init__(self, data, checksum, blocksize=8192): self.data = data self.checksum = checksum self._size = 0 + self.blocksize = blocksize + self.current_chunk = "" + self.closed = False - def read(self, length): - result = self.data.read(length) - self._size += len(result) - self.checksum.update(result) - return result + def read(self, size=None): + ret = "" + while size is None or size >= len(self.current_chunk): + ret += self.current_chunk + if size is not None: + size -= len(self.current_chunk) + if self.closed: + self.current_chunk = "" + break + self._get_chunk() + else: + ret += self.current_chunk[:size] + self.current_chunk = self.current_chunk[size:] + return ret + + def _get_chunk(self): + if not self.closed: + chunk = self.data.read(self.blocksize) + chunk_len = len(chunk) + self._size += chunk_len + self.checksum.update(chunk) + if chunk: + self.current_chunk = '%x\r\n%s\r\n' % (chunk_len, chunk) + else: + self.current_chunk = '0\r\n\r\n' + self.closed = True @property def size(self): @@ -269,7 +293,9 @@ class Store(glance.store.base.Store): 'image_id': image_id}) cookie = self._build_vim_cookie_header( self._session.vim.client.options.transport.cookiejar) - headers = {'Cookie': cookie, 'Content-Length': image_size} + headers = {'Connection': 'Keep-Alive', + 'Cookie': cookie, + 'Transfer-Encoding': 'chunked'} try: conn = self._get_http_conn('PUT', loc, headers, content=image_file) diff --git a/glance/tests/unit/test_vmware_store.py b/glance/tests/unit/test_vmware_store.py index fb9dd96d7f..cb8dd95944 100644 --- a/glance/tests/unit/test_vmware_store.py +++ b/glance/tests/unit/test_vmware_store.py @@ -250,3 +250,61 @@ class TestStore(base.StoreClearingUnitTest): with mock.patch('httplib.HTTPConnection') as HttpConn: HttpConn.return_value = FakeHTTPConnection(status=404) self.assertRaises(exception.NotFound, self.store.get_size, loc) + + def test_reader_image_fits_in_blocksize(self): + """ + Test that the image file reader returns the expected chunk of data + when the block size is larger than the image. + """ + content = 'XXX' + image = six.StringIO(content) + expected_checksum = hashlib.md5(content).hexdigest() + checksum = hashlib.md5() + reader = vm_store._Reader(image, checksum) + ret = reader.read() + expected_chunk = '%x\r\n%s\r\n' % (len(content), content) + last_chunk = '0\r\n\r\n' + self.assertEqual('%s%s' % (expected_chunk, last_chunk), ret) + self.assertEqual(image.len, reader.size) + self.assertEqual(expected_checksum, reader.checksum.hexdigest()) + self.assertTrue(reader.closed) + ret = reader.read() + self.assertEqual(image.len, reader.size) + self.assertEqual(expected_checksum, reader.checksum.hexdigest()) + self.assertTrue(reader.closed) + self.assertEqual('', ret) + + def test_reader_image_larger_blocksize(self): + """ + Test that the image file reader returns the expected chunks when + the block size specified is smaller than the image. + """ + content = 'XXX' + image = six.StringIO(content) + expected_checksum = hashlib.md5(content).hexdigest() + checksum = hashlib.md5() + last_chunk = '0\r\n\r\n' + reader = vm_store._Reader(image, checksum, blocksize=1) + ret = reader.read() + expected_chunk = '1\r\nX\r\n' + self.assertEqual('%s%s%s%s' % (expected_chunk, expected_chunk, + expected_chunk, last_chunk), ret) + self.assertEqual(expected_checksum, reader.checksum.hexdigest()) + self.assertEqual(image.len, reader.size) + self.assertTrue(reader.closed) + + def test_reader_size(self): + """Test that the image reader takes into account the specified size.""" + content = 'XXX' + image = six.StringIO(content) + expected_checksum = hashlib.md5(content).hexdigest() + checksum = hashlib.md5() + reader = vm_store._Reader(image, checksum, blocksize=1) + ret = reader.read(size=3) + self.assertEqual('1\r\n', ret) + ret = reader.read(size=1) + self.assertEqual('X', ret) + ret = reader.read() + self.assertEqual(expected_checksum, reader.checksum.hexdigest()) + self.assertEqual(image.len, reader.size) + self.assertTrue(reader.closed)