Add support for length encoded connection attributes

Author: oblitorum

connector.go

@@ -11,7 +11,6 @@ package mysql
 import (
 	"context"
 	"database/sql/driver"
-	"fmt"
 	"net"
 	"os"
 	"strconv"
@@ -24,7 +23,7 @@ type connector struct {
 }
 
 func encodeConnectionAttributes(textAttributes string) string {
-	connAttrsBuf := make([]byte, 0, 251)
+	connAttrsBuf := make([]byte, 0)
 
 	// default connection attributes
 	connAttrsBuf = appendLengthEncodedString(connAttrsBuf, connAttrClientName)
@@ -49,15 +48,12 @@ func encodeConnectionAttributes(textAttributes string) string {
 	return string(connAttrsBuf)
 }
 
-func newConnector(cfg *Config) (*connector, error) {
+func newConnector(cfg *Config) *connector {
 	encodedAttributes := encodeConnectionAttributes(cfg.ConnectionAttributes)
-	if len(encodedAttributes) > 250 {
-		return nil, fmt.Errorf("connection attributes are longer than 250 bytes: %dbytes (%q)", len(encodedAttributes), cfg.ConnectionAttributes)
-	}
 	return &connector{
 		cfg:               cfg,
 		encodedAttributes: encodedAttributes,
-	}, nil
+	}
 }
 
 // Connect implements driver.Connector interface.

connector_test.go

@@ -8,16 +8,13 @@ import (
 )
 
 func TestConnectorReturnsTimeout(t *testing.T) {
-	connector, err := newConnector(&Config{
+	connector := newConnector(&Config{
 		Net:     "tcp",
 		Addr:    "1.1.1.1:1234",
 		Timeout: 10 * time.Millisecond,
 	})
-	if err != nil {
-		t.Fatal(err)
-	}
 
-	_, err = connector.Connect(context.Background())
+	_, err := connector.Connect(context.Background())
 	if err == nil {
 		t.Fatal("error expected")
 	}

driver.go

@@ -83,10 +83,7 @@ func (d MySQLDriver) Open(dsn string) (driver.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	c, err := newConnector(cfg)
-	if err != nil {
-		return nil, err
-	}
+	c := newConnector(cfg)
 	return c.Connect(context.Background())
 }
 
@@ -108,7 +105,7 @@ func NewConnector(cfg *Config) (driver.Connector, error) {
 	if err := cfg.normalize(); err != nil {
 		return nil, err
 	}
-	return newConnector(cfg)
+	return newConnector(cfg), nil
 }
 
 // OpenConnector implements driver.DriverContext.
@@ -117,5 +114,5 @@ func (d MySQLDriver) OpenConnector(dsn string) (driver.Connector, error) {
 	if err != nil {
 		return nil, err
 	}
-	return newConnector(cfg)
+	return newConnector(cfg), nil
 }

packets.go

@@ -292,15 +292,14 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
 		pktLen += n + 1
 	}
 
-	// 1 byte to store length of all key-values
-	// NOTE: Actually, this is length encoded integer.
-	// But we support only len(connAttrBuf) < 251 for now because takeSmallBuffer
-	// doesn't support buffer size more than 4096 bytes.
-	// TODO(methane): Rewrite buffer management.
-	pktLen += 1 + len(mc.connector.encodedAttributes)
+	// encode length of the connection attributes
+	var connAttrsLEIBuf [9]byte
+	connAttrsLen := len(mc.connector.encodedAttributes)
+	connAttrsLEI := appendLengthEncodedInteger(connAttrsLEIBuf[:0], uint64(connAttrsLen))
+	pktLen += len(connAttrsLEI) + len(mc.connector.encodedAttributes)
 
 	// Calculate packet length and get buffer with that size
-	data, err := mc.buf.takeSmallBuffer(pktLen + 4)
+	data, err := mc.buf.takeBuffer(pktLen + 4)
 	if err != nil {
 		// cannot take the buffer. Something must be wrong with the connection
 		mc.cfg.Logger.Print(err)
@@ -380,8 +379,7 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
 	pos++
 
 	// Connection Attributes
-	data[pos] = byte(len(mc.connector.encodedAttributes))
-	pos++
+	pos += copy(data[pos:], connAttrsLEI)
 	pos += copy(data[pos:], []byte(mc.connector.encodedAttributes))
 
 	// Send Auth packet

packets_test.go

@@ -96,10 +96,7 @@ var _ net.Conn = new(mockConn)
 
 func newRWMockConn(sequence uint8) (*mockConn, *mysqlConn) {
 	conn := new(mockConn)
-	connector, err := newConnector(NewConfig())
-	if err != nil {
-		panic(err)
-	}
+	connector := newConnector(NewConfig())
 	mc := &mysqlConn{
 		buf:              newBuffer(conn),
 		cfg:              connector.cfg,