fixing bad connection error when reading large compressed packets

.golangci.yml

@@ -2,13 +2,10 @@ linters:
   disable-all: true
   enable:
     # All code is ready for:
-    - deadcode
     - errcheck
     - staticcheck
-    - structcheck
     - typecheck
     - unused
-    - varcheck
     - misspell
     - nolintlint
     - goimports

packet/conn.go

@@ -9,6 +9,7 @@ import (
 	"crypto/sha1"
 	"crypto/x509"
 	"encoding/pem"
+	goErrors "errors"
 	"io"
 	"net"
 	"sync"
@@ -108,41 +109,17 @@ func (c *Conn) ReadPacketReuseMem(dst []byte) ([]byte, error) {
 
 	if c.Compression != MYSQL_COMPRESS_NONE {
 		if !c.compressedReaderActive {
-			if _, err := io.ReadFull(c.reader, c.compressedHeader[:7]); err != nil {
-				return nil, errors.Wrapf(ErrBadConn, "io.ReadFull(compressedHeader) failed. err %v", err)
-			}
-
-			compressedSequence := c.compressedHeader[3]
-			uncompressedLength := int(uint32(c.compressedHeader[4]) | uint32(c.compressedHeader[5])<<8 | uint32(c.compressedHeader[6])<<16)
-			if compressedSequence != c.CompressedSequence {
-				return nil, errors.Errorf("invalid compressed sequence %d != %d",
-					compressedSequence, c.CompressedSequence)
-			}
-
-			if uncompressedLength > 0 {
-				var err error
-				switch c.Compression {
-				case MYSQL_COMPRESS_ZLIB:
-					c.compressedReader, err = zlib.NewReader(c.reader)
-				case MYSQL_COMPRESS_ZSTD:
-					c.compressedReader, err = zstd.NewReader(c.reader)
-				}
-				if err != nil {
-					return nil, err
-				}
+			var err error
+			c.compressedReader, err = c.newCompressedPacketReader()
+			if err != nil {
+				return nil, err
 			}
 			c.compressedReaderActive = true
 		}
 	}
 
-	if c.compressedReader != nil {
-		if err := c.ReadPacketTo(buf, c.compressedReader); err != nil {
-			return nil, errors.Trace(err)
-		}
-	} else {
-		if err := c.ReadPacketTo(buf, c.reader); err != nil {
-			return nil, errors.Trace(err)
-		}
+	if err := c.ReadPacketTo(buf); err != nil {
+		return nil, errors.Trace(err)
 	}
 
 	readBytes := buf.Bytes()
@@ -167,6 +144,41 @@ func (c *Conn) ReadPacketReuseMem(dst []byte) ([]byte, error) {
 	return result, nil
 }
 
+// newCompressedPacketReader creates a new compressed packet reader.
+func (c *Conn) newCompressedPacketReader() (io.Reader, error) {
+	if _, err := io.ReadFull(c.reader, c.compressedHeader[:7]); err != nil {
+		return nil, errors.Wrapf(ErrBadConn, "io.ReadFull(compressedHeader) failed. err %v", err)
+	}
+
+	compressedSequence := c.compressedHeader[3]
+	if compressedSequence != c.CompressedSequence {
+		return nil, errors.Errorf("invalid compressed sequence %d != %d",
+			compressedSequence, c.CompressedSequence)
+	}
+
+	compressedLength := int(uint32(c.compressedHeader[0]) | uint32(c.compressedHeader[1])<<8 | uint32(c.compressedHeader[2])<<16)
+	uncompressedLength := int(uint32(c.compressedHeader[4]) | uint32(c.compressedHeader[5])<<8 | uint32(c.compressedHeader[6])<<16)
+	if uncompressedLength > 0 {
+		limitedReader := io.LimitReader(c.reader, int64(compressedLength))
+		switch c.Compression {
+		case MYSQL_COMPRESS_ZLIB:
+			return zlib.NewReader(limitedReader)
+		case MYSQL_COMPRESS_ZSTD:
+			return zstd.NewReader(limitedReader)
+		}
+	}
+
+	return nil, nil
+}
+
+func (c *Conn) currentPacketReader() io.Reader {
+	if c.Compression == MYSQL_COMPRESS_NONE || c.compressedReader == nil {
+		return c.reader
+	} else {
+		return c.compressedReader
+	}
+}
+
 func (c *Conn) copyN(dst io.Writer, src io.Reader, n int64) (written int64, err error) {
 	for n > 0 {
 		bcap := cap(c.copyNBuf)
@@ -175,9 +187,29 @@ func (c *Conn) copyN(dst io.Writer, src io.Reader, n int64) (written int64, err
 		}
 		buf := c.copyNBuf[:bcap]
 
-		rd, err := io.ReadAtLeast(src, buf, bcap)
+		var rd int
+		rd, err = io.ReadAtLeast(src, buf, bcap)
+
 		n -= int64(rd)
 
+		// if we've read to EOF, and we have compression then advance the sequence number
+		// and reset the compressed reader to continue reading the remaining bytes
+		// in the next compressed packet.
+		if c.Compression != MYSQL_COMPRESS_NONE && rd < bcap &&
+			(goErrors.Is(err, io.ErrUnexpectedEOF) || goErrors.Is(err, io.EOF)) {
+			// we have read to EOF and read an incomplete uncompressed packet
+			// so advance the compressed sequence number and reset the compressed reader
+			// to get the remaining unread uncompressed bytes from the next compressed packet.
+			c.CompressedSequence++
+			if c.compressedReader, err = c.newCompressedPacketReader(); err != nil {
+				return written, errors.Trace(err)
+			}
+
+			// now read the remaining bytes into the buffer containing the first read bytes
+			rd, err = io.ReadAtLeast(c.currentPacketReader(), buf[rd:], bcap-rd)
+			n -= int64(rd)
+		}
+
 		if err != nil {
 			return written, errors.Trace(err)
 		}
@@ -192,9 +224,21 @@ func (c *Conn) copyN(dst io.Writer, src io.Reader, n int64) (written int64, err
 	return written, nil
 }
 
-func (c *Conn) ReadPacketTo(w io.Writer, r io.Reader) error {
-	if _, err := io.ReadFull(r, c.header[:4]); err != nil {
+func (c *Conn) ReadPacketTo(w io.Writer) error {
+	b := utils.BytesBufferGet()
+	defer func() {
+		utils.BytesBufferPut(b)
+	}()
+
+	// packets that come in a compressed packet may be partial
+	// so use the copyN function to read the packet header into a
+	// buffer, since copyN is capable of getting the next compressed
+	// packet and updating the Conn state with a new compressedReader.
+	if _, err := c.copyN(b, c.currentPacketReader(), 4); err != nil {
 		return errors.Wrapf(ErrBadConn, "io.ReadFull(header) failed. err %v", err)
+	} else {
+		// copy was successful so copy the 4 bytes from the buffer to the header
+		copy(c.header[:4], b.Bytes()[:4])
 	}
 
 	length := int(uint32(c.header[0]) | uint32(c.header[1])<<8 | uint32(c.header[2])<<16)
@@ -211,7 +255,7 @@ func (c *Conn) ReadPacketTo(w io.Writer, r io.Reader) error {
 		buf.Grow(length)
 	}
 
-	if n, err := c.copyN(w, r, int64(length)); err != nil {
+	if n, err := c.copyN(w, c.currentPacketReader(), int64(length)); err != nil {
 		return errors.Wrapf(ErrBadConn, "io.CopyN failed. err %v, copied %v, expected %v", err, n, length)
 	} else if n != int64(length) {
 		return errors.Wrapf(ErrBadConn, "io.CopyN failed(n != int64(length)). %v bytes copied, while %v expected", n, length)
@@ -220,7 +264,7 @@ func (c *Conn) ReadPacketTo(w io.Writer, r io.Reader) error {
 			return nil
 		}
 
-		if err = c.ReadPacketTo(w, r); err != nil {
+		if err = c.ReadPacketTo(w); err != nil {
 			return errors.Wrap(err, "ReadPacketTo failed")
 		}
 	}