proxy: simplified packetIO and proxyv2 en(de)coding

pkg/proxy/backend/authenticator.go

@@ -54,7 +54,7 @@ func (auth *Authenticator) handshakeFirstTime(clientIO, backendIO *pnet.PacketIO
 	// Read initial handshake packet from the backend.
 	serverPkt, serverCapability, err = auth.readInitialHandshake(backendIO)
 	if serverPkt != nil {
-		writeErr := clientIO.WritePacket(serverPkt)
+		writeErr := clientIO.WritePacket(serverPkt, false)
 		if writeErr != nil {
 			return false, writeErr
 		}
@@ -80,7 +80,7 @@ func (auth *Authenticator) handshakeFirstTime(clientIO, backendIO *pnet.PacketIO
 	sslEnabled := uint32(capability)&mysql.ClientSSL > 0
 	if sslEnabled {
 		// Upgrade TLS with the client if SSL is enabled.
-		if err = clientIO.UpgradeToServerTLS(serverTLSConfig); err != nil {
+		if _, err = clientIO.UpgradeToServerTLS(serverTLSConfig); err != nil {
 			return false, err
 		}
 	} else {
@@ -90,7 +90,7 @@ func (auth *Authenticator) handshakeFirstTime(clientIO, backendIO *pnet.PacketIO
 		copy(pktWithSSL[2:], clientPkt[2:])
 		clientPkt = pktWithSSL
 	}
-	if err = backendIO.WritePacket(clientPkt); err != nil {
+	if err = backendIO.WritePacket(clientPkt, false); err != nil {
 		return false, err
 	}
 	if err = backendIO.Flush(); err != nil {
@@ -108,7 +108,7 @@ func (auth *Authenticator) handshakeFirstTime(clientIO, backendIO *pnet.PacketIO
 		}
 	}
 	// Send the response again.
-	if err = backendIO.WritePacket(clientPkt); err != nil {
+	if err = backendIO.WritePacket(clientPkt, false); err != nil {
 		return false, err
 	}
 	if err = backendIO.Flush(); err != nil {
@@ -142,7 +142,7 @@ func forwardMsg(srcIO, destIO *pnet.PacketIO) (data []byte, err error) {
 	if err != nil {
 		return
 	}
-	err = destIO.WritePacket(data)
+	err = destIO.WritePacket(data, false)
 	if err != nil {
 		return
 	}
@@ -357,7 +357,7 @@ func (auth *Authenticator) writeAuthHandshake(backendIO *pnet.PacketIO, authData
 	}
 
 	// Send TLS / SSL request packet. The server must have supported TLS.
-	err := backendIO.WritePacket(data[:pos])
+	err := backendIO.WritePacket(data[:pos], false)
 	if err != nil {
 		return err
 	}
@@ -397,7 +397,7 @@ func (auth *Authenticator) writeAuthHandshake(backendIO *pnet.PacketIO, authData
 		pos += copy(data[pos:], auth.attrs)
 	}
 
-	if err = backendIO.WritePacket(data); err != nil {
+	if err = backendIO.WritePacket(data, false); err != nil {
 		return err
 	}
 	return backendIO.Flush()

pkg/proxy/backend/backend_conn.go

@@ -61,8 +61,7 @@ func (bc *BackendConnectionImpl) Connect() error {
 		return errors.New("dial backend error")
 	}
 
-	bufReadConn := pnet.NewBufferedReadConn(cn)
-	pkt := pnet.NewPacketIO(bufReadConn)
+	pkt := pnet.NewPacketIO(cn)
 	bc.pkt = pkt
 	return nil
 }

pkg/proxy/backend/cmd_processor.go

@@ -55,7 +55,7 @@ func (cp *CmdProcessor) executeCmd(request []byte, clientIO, backendIO *pnet.Pac
 		if _, response, err = cp.query(backendIO, "COMMIT"); err != nil {
 			// If commit fails, forward the response to the client.
 			if _, ok := err.(*gomysql.MyError); ok {
-				if err = clientIO.WritePacket(response); err == nil {
+				if err = clientIO.WritePacket(response, false); err == nil {
 					err = clientIO.Flush()
 				}
 			}
@@ -67,7 +67,7 @@ func (cp *CmdProcessor) executeCmd(request []byte, clientIO, backendIO *pnet.Pac
 		return
 	}
 
-	if err = backendIO.WritePacket(request); err != nil {
+	if err = backendIO.WritePacket(request, false); err != nil {
 		return
 	}
 	if err = backendIO.Flush(); err != nil {
@@ -119,7 +119,7 @@ func forwardOnePacket(clientIO, backendIO *pnet.PacketIO) (response []byte, err
 	if response, err = backendIO.ReadPacket(); err != nil {
 		return
 	}
-	err = clientIO.WritePacket(response)
+	err = clientIO.WritePacket(response, false)
 	return
 }
 
@@ -358,7 +358,7 @@ func (cp *CmdProcessor) query(packetIO *pnet.PacketIO, sql string) (result *gomy
 	request := make([]byte, 0, 1+len(data))
 	request = append(request, mysql.ComQuery)
 	request = append(request, data...)
-	if err = packetIO.WritePacket(request); err != nil {
+	if err = packetIO.WritePacket(request, false); err != nil {
 		return
 	}
 	if err = packetIO.Flush(); err != nil {

pkg/proxy/client/client_conn.go

@@ -30,22 +30,19 @@ import (
 )
 
 type ClientConnectionImpl struct {
-	serverTLSConfig  *tls.Config            // the TLS config to connect to clients.
-	backendTLSConfig *tls.Config            // the TLS config to connect to TiDB server.
-	pkt              *pnet.PacketIO         // a helper to read and write data in packet format.
-	bufReadConn      *pnet.BufferedReadConn // a buffered-read net.Conn or buffered-read tls.Conn.
+	serverTLSConfig  *tls.Config    // the TLS config to connect to clients.
+	backendTLSConfig *tls.Config    // the TLS config to connect to TiDB server.
+	pkt              *pnet.PacketIO // a helper to read and write data in packet format.
 	queryCtx         driver.QueryCtx
 	connectionID     uint64
 }
 
-func NewClientConnectionImpl(queryCtx driver.QueryCtx, conn net.Conn, connectionID uint64, serverTLSConfig, backendTLSConfig *tls.Config) driver.ClientConnection {
-	bufReadConn := pnet.NewBufferedReadConn(conn)
-	pkt := pnet.NewPacketIO(bufReadConn)
+func NewClientConnectionImpl(queryCtx driver.QueryCtx, conn net.Conn, connectionID uint64, serverTLSConfig *tls.Config, backendTLSConfig *tls.Config) driver.ClientConnection {
+	pkt := pnet.NewPacketIO(conn)
 	return &ClientConnectionImpl{
 		queryCtx:         queryCtx,
 		serverTLSConfig:  serverTLSConfig,
 		backendTLSConfig: backendTLSConfig,
-		bufReadConn:      bufReadConn,
 		pkt:              pkt,
 		connectionID:     connectionID,
 	}
@@ -56,7 +53,7 @@ func (cc *ClientConnectionImpl) ConnectionID() uint64 {
 }
 
 func (cc *ClientConnectionImpl) Addr() string {
-	return cc.bufReadConn.RemoteAddr().String()
+	return cc.pkt.RemoteAddr().String()
 }
 
 func (cc *ClientConnectionImpl) Run(ctx context.Context) {

pkg/proxy/net/error.go

@@ -0,0 +1,26 @@
+// Copyright 2022 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package net
+
+import "github.com/pingcap/TiProxy/pkg/util/errors"
+
+var (
+	ErrExpectSSLRequest = errors.New("expect a SSLRequest packet")
+	ErrReadConn         = errors.New("failed to read the connection")
+	ErrWriteConn        = errors.New("failed to write the connection")
+	ErrFlushConn        = errors.New("failed to flush the connection")
+	ErrCloseConn        = errors.New("failed to close the connection")
+	ErrHandshakeTLS     = errors.New("failed to complete tls handshake")
+)

pkg/proxy/net/mysql.go

@@ -0,0 +1,95 @@
+// Copyright 2022 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package net
+
+import (
+	"encoding/binary"
+
+	"github.com/pingcap/TiProxy/pkg/util/errors"
+	"github.com/pingcap/tidb/parser/mysql"
+)
+
+func (p *PacketIO) WriteInitialHandshake(capability uint32, status uint16, salt []byte) error {
+	data := make([]byte, 0, 128)
+
+	// min version 10
+	data = append(data, 10)
+	// server version[NUL]
+	data = append(data, serverVersion...)
+	data = append(data, 0)
+	// connection id
+	data = append(data, byte(capability), byte(capability>>8), byte(capability>>16), byte(capability>>24))
+	// auth-plugin-data-part-1
+	data = append(data, salt[:]...)
+	// filler [00]
+	data = append(data, 0)
+	// capability flag lower 2 bytes, using default capability here
+	data = append(data, byte(capability), byte(capability>>8))
+	// charset
+	data = append(data, utf8mb4BinID)
+	// status
+	data = append(data, byte(status), byte(status>>8))
+	// capability flag upper 2 bytes, using default capability here
+	data = append(data, byte(capability>>16), byte(capability>>24))
+	// length of auth-plugin-data
+	data = append(data, byte(len(salt)+1))
+	// reserved 10 [00]
+	data = append(data, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
+	// auth-plugin-data-part-2
+	data = append(data, salt[8:]...)
+	data = append(data, 0)
+	// auth-plugin name
+	data = append(data, []byte("mysql_native_password")...)
+	data = append(data, 0)
+
+	if err := p.WritePacket(data, true); err != nil {
+		return err
+	}
+
+	return p.Flush()
+}
+
+func (p *PacketIO) ReadSSLRequest() ([]byte, error) {
+	pkt, err := p.ReadPacket()
+	if err != nil {
+		return nil, err
+	}
+
+	if len(pkt) < 2 {
+		return nil, errors.WithStack(errors.Errorf("%w: but got less than 2 bytes", ErrExpectSSLRequest))
+	}
+
+	capability := uint32(binary.LittleEndian.Uint16(pkt[:2]))
+	if capability&mysql.ClientSSL == 0 {
+		return nil, errors.WithStack(errors.Errorf("%w: but capability flags has no SSL", ErrExpectSSLRequest))
+	}
+
+	return pkt, nil
+}
+
+func (p *PacketIO) WriteErrPacket(code uint16, msg string) error {
+	data := make([]byte, 0, 64)
+	data = append(data, 0xff)
+	data = append(data, byte(code), byte(code>>8))
+	// fill spaces for 421 protocol for compatibility
+	data = append(data, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20)
+	data = append(data, msg...)
+
+	if err := p.WritePacket(data, true); err != nil {
+		return err
+	}
+
+	return p.Flush()
+}