net: fix packet sequence when sequence overflows (#378)

Author: djshow832

pkg/proxy/net/compress.go

@@ -25,21 +25,11 @@ const (
 	CompressionZstd
 )
 
-type rwStatus int
-
-const (
-	rwNone rwStatus = iota
-	rwRead
-	rwWrite
-)
-
 const (
 	// maxCompressedSize is the max uncompressed data size for a compressed packet.
 	// Packets bigger than maxCompressedSize will be split into multiple compressed packets.
 	// MySQL has 16K for the first packet. The rest packets and MySQL Connector/J are 16M.
-	// Two restrictions for the length:
-	// - it should be smaller than 16M so that the length can fit in the 3 byte field in the header.
-	// - it should be larger than 4M so that the compressed sequence can fit in the 3 byte field when max_allowed_packet is 1G.
+	// It should be smaller than 16M so that the length can fit in the 3 byte field in the header.
 	maxCompressedSize = 1<<24 - 1
 	// minCompressSize is the min uncompressed data size for compressed data.
 	// Packets smaller than minCompressSize won't be compressed.
@@ -83,26 +73,25 @@ func newCompressedReadWriter(rw packetReadWriter, algorithm CompressAlgorithm, z
 	}
 }
 
-func (crw *compressedReadWriter) SetSequence(sequence uint8) {
-	crw.packetReadWriter.SetSequence(sequence)
+func (crw *compressedReadWriter) ResetSequence() {
+	crw.packetReadWriter.ResetSequence()
 	// Reset the compressed sequence before the next command.
-	if sequence == 0 {
-		crw.sequence = 0
-		crw.rwStatus = rwNone
-	}
+	// Sequence wraps around once it hits 0xFF, so we need ResetSequence() to know that it's reset instead of overflow.
+	crw.sequence = 0
+	crw.rwStatus = rwNone
 }
 
+// BeginRW implements packetReadWriter.BeginRW.
 // Uncompressed sequence of MySQL doesn't follow the spec: it's set to the compressed sequence when
 // the client/server begins reading or writing.
-func (crw *compressedReadWriter) beginRW(status rwStatus) {
+func (crw *compressedReadWriter) BeginRW(status rwStatus) {
 	if crw.rwStatus != status {
 		crw.packetReadWriter.SetSequence(crw.sequence)
 		crw.rwStatus = status
 	}
 }
 
 func (crw *compressedReadWriter) Read(p []byte) (n int, err error) {
-	crw.beginRW(rwRead)
 	// Read from the connection to fill the buffer if the buffer is empty.
 	if crw.readBuffer.Len() == 0 {
 		if err = crw.readFromConn(); err != nil {
@@ -156,7 +145,6 @@ func (crw *compressedReadWriter) readFromConn() error {
 }
 
 func (crw *compressedReadWriter) Write(data []byte) (n int, err error) {
-	crw.beginRW(rwWrite)
 	for {
 		remainingLen := maxCompressedSize - crw.writeBuffer.Len()
 		if len(data) <= remainingLen {
@@ -234,7 +222,6 @@ func (crw *compressedReadWriter) DirectWrite(data []byte) (n int, err error) {
 // Peek won't be used.
 // Notice: the peeked data may be discarded if an error is returned.
 func (crw *compressedReadWriter) Peek(n int) (data []byte, err error) {
-	crw.beginRW(rwRead)
 	for crw.readBuffer.Len() < n {
 		if err = crw.readFromConn(); err != nil {
 			return
@@ -247,7 +234,6 @@ func (crw *compressedReadWriter) Peek(n int) (data []byte, err error) {
 
 // Discard won't be used.
 func (crw *compressedReadWriter) Discard(n int) (d int, err error) {
-	crw.beginRW(rwRead)
 	for crw.readBuffer.Len() < n {
 		if err = crw.readFromConn(); err != nil {
 			return

pkg/proxy/net/compress_test.go

@@ -235,6 +235,7 @@ func TestReadWriteError(t *testing.T) {
 }
 
 func fillAndWrite(t *testing.T, crw *compressedReadWriter, b byte, length int) {
+	crw.BeginRW(rwWrite)
 	data := fillData(b, length)
 	_, err := crw.Write(data)
 	require.NoError(t, err)
@@ -250,6 +251,7 @@ func fillData(b byte, length int) []byte {
 }
 
 func readAndCheck(t *testing.T, crw *compressedReadWriter, b byte, length int) {
+	crw.BeginRW(rwRead)
 	data := make([]byte, length)
 	_, err := io.ReadFull(crw, data)
 	require.NoError(t, err)

pkg/proxy/net/packetio.go

@@ -50,6 +50,14 @@ const (
 	defaultReaderSize = 16 * 1024
 )
 
+type rwStatus int
+
+const (
+	rwNone rwStatus = iota
+	rwRead
+	rwWrite
+)
+
 // packetReadWriter acts like a net.Conn with read and write buffer.
 type packetReadWriter interface {
 	net.Conn
@@ -65,6 +73,10 @@ type packetReadWriter interface {
 	IsPeerActive() bool
 	SetSequence(uint8)
 	Sequence() uint8
+	// ResetSequence is called before executing a command.
+	ResetSequence()
+	// BeginRW is called before reading or writing packets.
+	BeginRW(status rwStatus)
 }
 
 var _ packetReadWriter = (*basicReadWriter)(nil)
@@ -123,6 +135,13 @@ func (brw *basicReadWriter) OutBytes() uint64 {
 	return brw.outBytes
 }
 
+func (brw *basicReadWriter) BeginRW(rwStatus) {
+}
+
+func (brw *basicReadWriter) ResetSequence() {
+	brw.sequence = 0
+}
+
 func (brw *basicReadWriter) TLSConnectionState() tls.ConnectionState {
 	return tls.ConnectionState{}
 }
@@ -190,7 +209,7 @@ func (p *PacketIO) RemoteAddr() net.Addr {
 }
 
 func (p *PacketIO) ResetSequence() {
-	p.readWriter.SetSequence(0)
+	p.readWriter.ResetSequence()
 }
 
 // GetSequence is used in tests to assert that the sequences on the client and server are equal.
@@ -219,6 +238,7 @@ func (p *PacketIO) readOnePacket() ([]byte, bool, error) {
 
 // ReadPacket reads data and removes the header
 func (p *PacketIO) ReadPacket() (data []byte, err error) {
+	p.readWriter.BeginRW(rwRead)
 	for more := true; more; {
 		var buf []byte
 		buf, more, err = p.readOnePacket()
@@ -262,6 +282,7 @@ func (p *PacketIO) writeOnePacket(data []byte) (int, bool, error) {
 
 // WritePacket writes data without a header
 func (p *PacketIO) WritePacket(data []byte, flush bool) (err error) {
+	p.readWriter.BeginRW(rwWrite)
 	for more := true; more; {
 		var n int
 		n, more, err = p.writeOnePacket(data)

pkg/proxy/net/packetio_test.go

@@ -362,3 +362,55 @@ func TestProxyTLSCompress(t *testing.T) {
 		}
 	}
 }
+
+// Test the sequence is correct with the compression protocol.
+func TestPacketSequence(t *testing.T) {
+	write := func(p *PacketIO, flush bool) {
+		require.NoError(t, p.WritePacket([]byte{0}, flush))
+	}
+	read := func(p *PacketIO) {
+		_, err := p.ReadPacket()
+		require.NoError(t, err)
+	}
+	loops := 1024
+	testTCPConn(t,
+		func(t *testing.T, cli *PacketIO) {
+			require.NoError(t, cli.SetCompressionAlgorithm(CompressionZlib, 0))
+			read(cli)
+			// uncompressed sequence = compressed sequence
+			write(cli, false)
+			write(cli, true)
+			// uncompressed sequence wraps around (1000 writes + 1 flush)
+			for i := 0; i < loops; i++ {
+				write(cli, false)
+			}
+			require.NoError(t, cli.Flush())
+			// compressed sequence wraps around (1000 writes + 1000 flushes)
+			for i := 0; i < loops; i++ {
+				write(cli, true)
+			}
+			// reset sequence
+			cli.ResetSequence()
+			write(cli, true)
+		},
+		func(t *testing.T, srv *PacketIO) {
+			require.NoError(t, srv.SetCompressionAlgorithm(CompressionZlib, 0))
+			write(srv, true)
+			// uncompressed sequence = compressed sequence
+			read(srv)
+			read(srv)
+			// uncompressed sequence wraps around
+			for i := 0; i < loops; i++ {
+				read(srv)
+			}
+			// compressed sequence wraps around
+			for i := 0; i < loops; i++ {
+				read(srv)
+			}
+			// reset sequence
+			srv.ResetSequence()
+			read(srv)
+		},
+		1,
+	)
+}