re-implement TestReadPacketFail

Author: shogo82148

packets.go

@@ -34,9 +34,12 @@ func (mc *mysqlConn) readPacket(ctx context.Context) ([]byte, error) {
 		// read packet header
 		err := mc.readFull(ctx, mc.data[:4])
 		if err != nil {
+			if err == context.Canceled || err == context.DeadlineExceeded {
+				return nil, err
+			}
 			mc.cfg.Logger.Print(err)
 			mc.closeContext(ctx)
-			return nil, err
+			return nil, ErrInvalidConn
 		}
 
 		// packet length [24 bit]
@@ -69,6 +72,9 @@ func (mc *mysqlConn) readPacket(ctx context.Context) ([]byte, error) {
 		data := make([]byte, pktLen)
 		err = mc.readFull(ctx, data)
 		if err != nil {
+			if err == context.Canceled || err == context.DeadlineExceeded {
+				return nil, err
+			}
 			mc.cfg.Logger.Print(err)
 			mc.closeContext(ctx)
 			return nil, ErrInvalidConn

packets_test.go

@@ -226,47 +226,50 @@ func TestReadPacketSplit(t *testing.T) {
 	})
 }
 
-// func TestReadPacketFail(t *testing.T) {
-// 	conn := new(mockConn)
-// 	mc := &mysqlConn{
-// 		buf:     newBuffer(conn),
-// 		closech: make(chan struct{}),
-// 		cfg:     NewConfig(),
-// 	}
+func TestReadPacketFail(t *testing.T) {
+	t.Run("illegal empty (stand-alone) packet", func(t *testing.T) {
+		conn, mc := newRWMockConn(t, 0)
+		go func() {
+			conn.Write([]byte{0x00, 0x00, 0x00, 0x00})
+		}()
+		go func() {
+			io.Copy(io.Discard, conn)
+		}()
 
-// 	// illegal empty (stand-alone) packet
-// 	conn.data = []byte{0x00, 0x00, 0x00, 0x00}
-// 	conn.maxReads = 1
-// 	_, err := mc.readPacket()
-// 	if err != ErrInvalidConn {
-// 		t.Errorf("expected ErrInvalidConn, got %v", err)
-// 	}
+		_, err := mc.readPacket(context.Background())
+		if err != ErrInvalidConn {
+			t.Errorf("expected ErrInvalidConn, got %v", err)
+		}
+	})
 
-// 	// reset
-// 	conn.reads = 0
-// 	mc.sequence = 0
-// 	mc.buf = newBuffer(conn)
+	t.Run("fail to read header", func(t *testing.T) {
+		conn, mc := newRWMockConn(t, 0)
+		go func() {
+			conn.Close()
+		}()
 
-// 	// fail to read header
-// 	conn.closed = true
-// 	_, err = mc.readPacket()
-// 	if err != ErrInvalidConn {
-// 		t.Errorf("expected ErrInvalidConn, got %v", err)
-// 	}
+		_, err := mc.readPacket(context.Background())
+		if err != ErrInvalidConn {
+			t.Errorf("expected ErrInvalidConn, got %v", err)
+		}
+	})
 
-// 	// reset
-// 	conn.closed = false
-// 	conn.reads = 0
-// 	mc.sequence = 0
-// 	mc.buf = newBuffer(conn)
+	t.Run("fail to read body", func(t *testing.T) {
+		conn, mc := newRWMockConn(t, 0)
+		go func() {
+			conn.Write([]byte{0x01, 0x00, 0x00, 0x00})
+			conn.Close()
+		}()
+		go func() {
+			io.Copy(io.Discard, conn)
+		}()
 
-// 	// fail to read body
-// 	conn.maxReads = 1
-// 	_, err = mc.readPacket()
-// 	if err != ErrInvalidConn {
-// 		t.Errorf("expected ErrInvalidConn, got %v", err)
-// 	}
-// }
+		_, err := mc.readPacket(context.Background())
+		if err != ErrInvalidConn {
+			t.Errorf("expected ErrInvalidConn, got %v", err)
+		}
+	})
+}
 
 // // https://github.com/go-sql-driver/mysql/pull/801
 // // not-NUL terminated plugin_name in init packet