remove packetWriter and simplify tests

Author: methane

compress.go

@@ -86,7 +86,7 @@ func zCompress(src []byte, dst io.Writer) error {
 	return nil
 }
 
-type compressor struct {
+type decompressor struct {
 	mc *mysqlConn
 	// read buffer (FIFO).
 	// We can not reuse already-read buffer until dropping Go 1.20 support.
@@ -95,13 +95,13 @@ type compressor struct {
 	bytesBuf []byte
 }
 
-func newCompressor(mc *mysqlConn) *compressor {
-	return &compressor{
+func newDecompressor(mc *mysqlConn) *decompressor {
+	return &decompressor{
 		mc: mc,
 	}
 }
 
-func (c *compressor) readNext(need int) ([]byte, error) {
+func (c *decompressor) readNext(need int) ([]byte, error) {
 	for len(c.bytesBuf) < need {
 		if err := c.uncompressPacket(); err != nil {
 			return nil, err
@@ -113,7 +113,7 @@ func (c *compressor) readNext(need int) ([]byte, error) {
 	return data, nil
 }
 
-func (c *compressor) uncompressPacket() error {
+func (c *decompressor) uncompressPacket() error {
 	header, err := c.mc.buf.readNext(7) // size of compressed header
 	if err != nil {
 		return err
@@ -166,9 +166,11 @@ func (c *compressor) uncompressPacket() error {
 
 const maxPayloadLen = maxPacketSize - 4
 
-func (c *compressor) Write(data []byte) (int, error) {
-	totalBytes := len(data)
-	dataLen := len(data)
+// writeCompressed sends one or some packets with compression.
+// Use this instead of mc.netConn.Write() when mc.compress is true.
+func (mc *mysqlConn) writeCompressed(packets []byte) (int, error) {
+	totalBytes := len(packets)
+	dataLen := len(packets)
 	blankHeader := make([]byte, 7)
 	var buf bytes.Buffer
 
@@ -177,7 +179,7 @@ func (c *compressor) Write(data []byte) (int, error) {
 		if payloadLen > maxPayloadLen {
 			payloadLen = maxPayloadLen
 		}
-		payload := data[:payloadLen]
+		payload := packets[:payloadLen]
 		uncompressedLen := payloadLen
 
 		if _, err := buf.Write(blankHeader); err != nil {
@@ -194,11 +196,11 @@ func (c *compressor) Write(data []byte) (int, error) {
 			zCompress(payload, &buf)
 		}
 
-		if err := c.writeCompressedPacket(buf.Bytes(), uncompressedLen); err != nil {
+		if err := mc.writeCompressedPacket(buf.Bytes(), uncompressedLen); err != nil {
 			return 0, err
 		}
 		dataLen -= payloadLen
-		data = data[payloadLen:]
+		packets = packets[payloadLen:]
 		buf.Reset()
 	}
 
@@ -207,32 +209,32 @@ func (c *compressor) Write(data []byte) (int, error) {
 
 // writeCompressedPacket writes a compressed packet with header.
 // data should start with 7 size space for header followed by payload.
-func (c *compressor) writeCompressedPacket(data []byte, uncompressedLen int) error {
+func (mc *mysqlConn) writeCompressedPacket(data []byte, uncompressedLen int) error {
 	comprLength := len(data) - 7
 	if debugTrace {
-		c.mc.cfg.Logger.Print(
+		mc.cfg.Logger.Print(
 			fmt.Sprintf(
 				"writeCompressedPacket: comprLength=%v, uncompressedLen=%v, seq=%v",
-				comprLength, uncompressedLen, c.mc.compressSequence))
+				comprLength, uncompressedLen, mc.compressSequence))
 	}
 
 	// compression header
 	data[0] = byte(0xff & comprLength)
 	data[1] = byte(0xff & (comprLength >> 8))
 	data[2] = byte(0xff & (comprLength >> 16))
 
-	data[3] = c.mc.compressSequence
+	data[3] = mc.compressSequence
 
 	// this value is never greater than maxPayloadLength
 	data[4] = byte(0xff & uncompressedLen)
 	data[5] = byte(0xff & (uncompressedLen >> 8))
 	data[6] = byte(0xff & (uncompressedLen >> 16))
 
-	if _, err := c.mc.netConn.Write(data); err != nil {
-		c.mc.cfg.Logger.Print(err)
+	if _, err := mc.netConn.Write(data); err != nil {
+		mc.cfg.Logger.Print(err)
 		return err
 	}
 
-	c.mc.compressSequence++
+	mc.compressSequence++
 	return nil
 }

compress_test.go

@@ -13,7 +13,6 @@ import (
 	"crypto/rand"
 	"fmt"
 	"io"
-	"net"
 	"testing"
 )
 
@@ -23,85 +22,37 @@ func makeRandByteSlice(size int) []byte {
 	return randBytes
 }
 
-func newMockConn() *mysqlConn {
-	newConn := &mysqlConn{cfg: NewConfig()}
-	return newConn
-}
-
-func newMockBuf(data []byte) buffer {
-	return buffer{
-		buf:    data,
-		length: len(data),
-	}
-}
-
-type dummyConn struct {
-	buf bytes.Buffer
-	net.Conn
-}
-
-func (c *dummyConn) Write(data []byte) (int, error) {
-	return c.buf.Write(data)
-}
-
 // compressHelper compresses uncompressedPacket and checks state variables
 func compressHelper(t *testing.T, mc *mysqlConn, uncompressedPacket []byte) []byte {
-	// get status variables
-
-	cs := mc.compressSequence
-
-	var b dummyConn
-	mc.netConn = &b
-	cw := newCompressor(mc)
-
-	n, err := cw.Write(uncompressedPacket)
+	conn := new(mockConn)
+	mc.netConn = conn
 
+	n, err := mc.writeCompressed(uncompressedPacket)
 	if err != nil {
-		t.Fatal(err.Error())
+		t.Fatal(err)
 	}
-
 	if n != len(uncompressedPacket) {
 		t.Fatalf("expected to write %d bytes, wrote %d bytes", len(uncompressedPacket), n)
 	}
-
-	if len(uncompressedPacket) > 0 {
-		if mc.compressSequence != (cs + 1) {
-			t.Fatalf("mc.compressionSequence updated incorrectly, expected %d and saw %d", (cs + 1), mc.compressSequence)
-		}
-
-	} else {
-		if mc.compressSequence != cs {
-			t.Fatalf("mc.compressionSequence updated incorrectly for case of empty write, expected %d and saw %d", cs, mc.compressSequence)
-		}
-	}
-
-	return b.buf.Bytes()
+	return conn.written
 }
 
 // uncompressHelper uncompresses compressedPacket and checks state variables
 func uncompressHelper(t *testing.T, mc *mysqlConn, compressedPacket []byte, expSize int) []byte {
-	// get status variables
-	cs := mc.compressSequence
-
 	// mocking out buf variable
-	mc.buf = newMockBuf(compressedPacket)
-	cr := newCompressor(mc)
+	conn := new(mockConn)
+	conn.data = compressedPacket
+	mc.buf.nc = conn
+	cr := newDecompressor(mc)
 
 	uncompressedPacket, err := cr.readNext(expSize)
 	if err != nil {
 		if err != io.EOF {
 			t.Fatalf("non-nil/non-EOF error when reading contents: %s", err.Error())
 		}
 	}
-
-	if expSize > 0 {
-		if mc.compressSequence != (cs + 1) {
-			t.Fatalf("mc.compressionSequence updated incorrectly, expected %d and saw %d", (cs + 1), mc.compressSequence)
-		}
-	} else {
-		if mc.compressSequence != cs {
-			t.Fatalf("mc.compressionSequence updated incorrectly for case of empty read, expected %d and saw %d", cs, mc.compressSequence)
-		}
+	if len(uncompressedPacket) != expSize {
+		t.Errorf("uncompressed size is unexpected. expected %d but got %d", expSize, len(uncompressedPacket))
 	}
 	return uncompressedPacket
 }
@@ -141,20 +92,33 @@ func TestRoundtrip(t *testing.T) {
 		{uncompressed: makeRandByteSlice(32768),
 			desc: "32768 rand bytes",
 		},
-		{uncompressed: makeRandByteSlice(33000),
-			desc: "33000 rand bytes",
+		{uncompressed: bytes.Repeat(makeRandByteSlice(100), 10000),
+			desc: "100 rand * 10000 repeat bytes",
 		},
 	}
 
-	cSend := newMockConn()
-	cReceive := newMockConn()
+	_, cSend := newRWMockConn(0)
+	cSend.compress = true
+	_, cReceive := newRWMockConn(0)
+	cReceive.compress = true
 
 	for _, test := range tests {
 		s := fmt.Sprintf("Test roundtrip with %s", test.desc)
+		cSend.resetSequenceNr()
+		cReceive.resetSequenceNr()
 
 		uncompressed := roundtripHelper(t, cSend, cReceive, test.uncompressed)
 		if !bytes.Equal(uncompressed, test.uncompressed) {
 			t.Fatalf("%s: roundtrip failed", s)
 		}
+
+		if cSend.sequence != cReceive.sequence {
+			t.Errorf("inconsistent sequence number: send=%v recv=%v",
+				cSend.sequence, cReceive.sequence)
+		}
+		if cSend.compressSequence != cReceive.compressSequence {
+			t.Errorf("inconsistent compress sequence number: send=%v recv=%v",
+				cSend.compressSequence, cReceive.compressSequence)
+		}
 	}
 }

connection.go

@@ -26,7 +26,6 @@ type mysqlConn struct {
 	rawConn          net.Conn    // underlying connection when netConn is TLS connection.
 	result           mysqlResult // managed by clearResult() and handleOkPacket().
 	packetReader     packetReader
-	packetWriter     io.Writer
 	cfg              *Config
 	connector        *connector
 	maxAllowedPacket int

connection_test.go

@@ -169,7 +169,6 @@ func TestPingMarkBadConnection(t *testing.T) {
 		netConn:          nc,
 		buf:              buf,
 		packetReader:     &buf,
-		packetWriter:     nc,
 		maxAllowedPacket: defaultMaxAllowedPacket,
 	}
 
@@ -188,7 +187,6 @@ func TestPingErrInvalidConn(t *testing.T) {
 		netConn:          nc,
 		buf:              buf,
 		packetReader:     &buf,
-		packetWriter:     nc,
 		maxAllowedPacket: defaultMaxAllowedPacket,
 		closech:          make(chan struct{}),
 		cfg:              NewConfig(),

connector.go

@@ -125,7 +125,6 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	mc.buf = newBuffer(mc.netConn)
 	// packet reader and writer in handshake are never compressed
 	mc.packetReader = &mc.buf
-	mc.packetWriter = mc.netConn
 
 	// Set I/O timeouts
 	mc.buf.timeout = mc.cfg.ReadTimeout
@@ -168,10 +167,9 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 		return nil, err
 	}
 
-	if mc.compress {
-		cmpr := newCompressor(mc)
-		mc.packetReader = cmpr
-		mc.packetWriter = cmpr
+	if mc.cfg.compress && mc.flags&clientCompress == clientCompress {
+		mc.compress = true
+		mc.packetReader = newDecompressor(mc)
 	}
 	if mc.cfg.MaxAllowedPacket > 0 {
 		mc.maxAllowedPacket = mc.cfg.MaxAllowedPacket

packets.go

@@ -131,7 +131,15 @@ func (mc *mysqlConn) writePacket(data []byte) error {
 			}
 		}
 
-		n, err := mc.packetWriter.Write(data[:4+size])
+		var (
+			n   int
+			err error
+		)
+		if mc.compress {
+			n, err = mc.writeCompressed(data[:4+size])
+		} else {
+			n, err = mc.netConn.Write(data[:4+size])
+		}
 		if err == nil && n == 4+size {
 			mc.sequence++
 			if size != maxPacketSize {
@@ -278,7 +286,6 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
 	}
 	if mc.cfg.compress && mc.flags&clientCompress == clientCompress {
 		clientFlags |= clientCompress
-		mc.compress = true
 	}
 	// To enable TLS / SSL
 	if mc.cfg.TLS != nil {
@@ -368,7 +375,6 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
 		mc.rawConn = mc.netConn
 		mc.netConn = tlsConn
 		mc.buf.nc = tlsConn
-		mc.packetWriter = mc.netConn
 	}
 
 	// User [null terminated string]

packets_test.go

@@ -101,7 +101,6 @@ func newRWMockConn(sequence uint8) (*mockConn, *mysqlConn) {
 	mc := &mysqlConn{
 		buf:              buf,
 		packetReader:     &buf,
-		packetWriter:     conn,
 		cfg:              connector.cfg,
 		connector:        connector,
 		netConn:          conn,