From c48aa9410319900ad5dc1d1c5f93b89fb546c24c Mon Sep 17 00:00:00 2001
From: longwenwen <304330668@qq.com>
Date: Wed, 12 Dec 2018 17:41:16 +0800
Subject: [PATCH] fix compress bug

---
 lib/mysql/connector/network.py | 65 ++++++++++++++++++----------------
 1 file changed, 34 insertions(+), 31 deletions(-)

diff --git a/lib/mysql/connector/network.py b/lib/mysql/connector/network.py
index 078e3f3..e6bd90d 100644
--- a/lib/mysql/connector/network.py
+++ b/lib/mysql/connector/network.py
@@ -87,6 +87,7 @@ def __init__(self):
         self._compressed_packet_number = -1
         self._packet_queue = deque()
         self.recvsize = 8192
+        self._resi_bytes = None
 
     @property
     def next_packet_number(self):
@@ -316,13 +317,22 @@ def recv_py26_plain(self):
 
     def _split_zipped_payload(self, packet_bunch):
         """Split compressed payload"""
+        if self._resi_bytes is not None:
+            packet_bunch = self._resi_bytes + packet_bunch
+            self._resi_bytes = None
         while packet_bunch:
+            if len(packet_bunch) < 3:
+                self._resi_bytes = packet_bunch
+                return
             if PY2:
                 payload_length = struct.unpack_from(
                     "<I",
                     packet_bunch[0:3] + b'\x00')[0]  # pylint: disable=E0602
             else:
                 payload_length = struct.unpack("<I", packet_bunch[0:3] + b'\x00')[0]
+            if payload_length > len(packet_bunch) - 4:
+                self._resi_bytes = packet_bunch
+                return
 
             self._packet_queue.append(packet_bunch[0:payload_length + 4])
             packet_bunch = packet_bunch[payload_length + 4:]
@@ -336,14 +346,15 @@ def recv_compressed(self):
         except IndexError:
             pass
 
-        header = bytearray(b'')
-        packets = []
-        try:
-            abyte = self.sock.recv(1)
-            while abyte and len(header) < 7:
-                header += abyte
+        while True:
+            header = bytearray(b'')
+            packets = []
+            try:
                 abyte = self.sock.recv(1)
-            while header:
+                while abyte and len(header) < 7:
+                    header += abyte
+                    abyte = self.sock.recv(1)
+
                 if len(header) < 7:
                     raise errors.InterfaceError(errno=2013)
 
@@ -372,32 +383,24 @@ def recv_compressed(self):
 
                 packets.append((payload_length, zip_payload))
 
-                if zip_payload_length <= 16384:
-                    # We received the full compressed packet
-                    break
-
-                # Get next compressed packet
-                header = init_bytearray(b'')
-                abyte = self.sock.recv(1)
-                while abyte and len(header) < 7:
-                    header += abyte
-                    abyte = self.sock.recv(1)
+            except IOError as err:
+                raise errors.OperationalError(
+                    errno=2055, values=(self.get_address(), _strioerror(err)))
 
-        except IOError as err:
-            raise errors.OperationalError(
-                errno=2055, values=(self.get_address(), _strioerror(err)))
+            # Compressed packet can contain more than 1 MySQL packets
+            # We decompress and make one so we can split it up
+            tmp = init_bytearray(b'')
+            for payload_length, payload in packets:
+                # payload_length can not be 0; this was previously handled
+                if PY2:
+                    tmp += zlib.decompress(buffer(payload))  # pylint: disable=E0602
+                else:
+                    tmp += zlib.decompress(payload)
+            self._split_zipped_payload(tmp)
+            del tmp
 
-        # Compressed packet can contain more than 1 MySQL packets
-        # We decompress and make one so we can split it up
-        tmp = init_bytearray(b'')
-        for payload_length, payload in packets:
-            # payload_length can not be 0; this was previously handled
-            if PY2:
-                tmp += zlib.decompress(buffer(payload))  # pylint: disable=E0602
-            else:
-                tmp += zlib.decompress(payload)
-        self._split_zipped_payload(tmp)
-        del tmp
+            if len(self._packet_queue) != 0:
+                break
 
         try:
             pkt = self._packet_queue.popleft()