support mysql protocol connection attributes (vitessio#18548)

demmer · demmer · commit fac2b0407bfc · 2025-10-02T12:56:03.000-07:00
Signed-off-by: Michael Demmer &lt;mdemmer@slack-corp.com&gt;
diff --git a/go/mysql/client.go b/go/mysql/client.go
@@ -48,6 +48,15 @@ type connectResult struct {
 // FIXME(alainjobart) once we have more of a server side, add test cases
 // to cover all failure scenarios.
 func Connect(ctx context.Context, params *ConnParams) (*Conn, error) {
+	return ConnectWithAttributes(ctx, params, ConnectionAttributes{})
+}
+
+// ConnectWithAttributes creates a connection to a server with connection attributes.
+// It then handles the initial handshake.
+//
+// If context is canceled before the end of the process, this function
+// will return nil, ctx.Err().
+func ConnectWithAttributes(ctx context.Context, params *ConnParams, attributes ConnectionAttributes) (*Conn, error) {
 	if params.ConnectTimeoutMs != 0 {
 		var cancel context.CancelFunc
 		ctx, cancel = context.WithTimeout(ctx, time.Duration(params.ConnectTimeoutMs)*time.Millisecond)
@@ -116,7 +125,7 @@ func Connect(ctx context.Context, params *ConnParams) (*Conn, error) {
 		// make any read or write just return with an error
 		// right away.
 		status <- connectResult{
-			err: c.clientHandshake(params),
+			err: c.clientHandshake(params, attributes),
 		}
 	}()
 
@@ -198,7 +207,7 @@ func (c *Conn) Ping() error {
 // clientHandshake handles the client side of the handshake.
 // Note the connection can be closed while this is running.
 // Returns a SQLError.
-func (c *Conn) clientHandshake(params *ConnParams) error {
+func (c *Conn) clientHandshake(params *ConnParams, attributes ConnectionAttributes) error {
 	// if EnableQueryInfo is set, make sure that all queries starting with the handshake
 	// will actually process the INFO fields in QUERY_OK packets
 	if params.EnableQueryInfo {
@@ -295,9 +304,14 @@ func (c *Conn) clientHandshake(params *ConnParams) error {
 		return sqlerror.NewSQLError(sqlerror.CRSSLConnectionError, sqlerror.SSUnknownSQLState, "server doesn't support ClientSessionTrack but client asked for it")
 	}
 
+	// Connection attributes.
+	if capabilities&CapabilityClientConnAttr != 0 && len(attributes) > 0 {
+		c.Capabilities |= CapabilityClientConnAttr
+	}
+
 	// Build and send our handshake response 41.
 	// Note this one will never have SSL flag on.
-	if err := c.writeHandshakeResponse41(capabilities, scrambledPassword, uint8(params.Charset), params); err != nil {
+	if err := c.writeHandshakeResponse41(capabilities, scrambledPassword, uint8(params.Charset), params, attributes); err != nil {
 		return err
 	}
 
@@ -527,7 +541,7 @@ const CapabilityFlagsSsl = CapabilityFlags |
 
 // writeHandshakeResponse41 writes the handshake response.
 // Returns a SQLError.
-func (c *Conn) writeHandshakeResponse41(capabilities uint32, scrambledPassword []byte, characterSet uint8, params *ConnParams) error {
+func (c *Conn) writeHandshakeResponse41(capabilities uint32, scrambledPassword []byte, characterSet uint8, params *ConnParams, attributes ConnectionAttributes) error {
 	// Build our flags.
 	capabilityFlags := CapabilityFlags |
 		// If the server supported
@@ -564,6 +578,17 @@ func (c *Conn) writeHandshakeResponse41(capabilities uint32, scrambledPassword [
 		length++
 	}
 
+	// If the server supports CapabilityClientConnAttr and there are attributes to be
+	// sent, then calculate the length of the attributes and include it in the overall length.
+	var attrLength int
+	if capabilities&CapabilityClientConnAttr != 0 && len(attributes) > 0 {
+		capabilityFlags |= CapabilityClientConnAttr
+		for key, value := range attributes {
+			attrLength += lenEncStringSize(key) + lenEncStringSize(value)
+		}
+		length += lenEncIntSize(uint64(attrLength)) + attrLength
+	}
+
 	data, pos := c.startEphemeralPacketWithHeader(length)
 
 	// Client capability flags.
@@ -600,6 +625,16 @@ func (c *Conn) writeHandshakeResponse41(capabilities uint32, scrambledPassword [
 	// Assume native client during response
 	pos = writeNullString(data, pos, string(c.authPluginName))
 
+	// Client conn attributes
+	if attrLength > 0 {
+		pos = writeLenEncInt(data, pos, uint64(attrLength))
+
+		for key, value := range attributes {
+			pos = writeLenEncString(data, pos, key)
+			pos = writeLenEncString(data, pos, value)
+		}
+	}
+
 	// Sanity-check the length.
 	if pos != len(data) {
 		return sqlerror.NewSQLError(sqlerror.CRMalformedPacket, sqlerror.SSUnknownSQLState, "writeHandshakeResponse41: only packed %v bytes, out of %v allocated", pos, len(data))
diff --git a/go/mysql/conn.go b/go/mysql/conn.go
@@ -128,6 +128,10 @@ type Conn struct {
 	// It is set during the initial handshake.
 	UserData Getter
 
+	// ConnectionAttributes stores arbitrary client-supplied attributes sent in the
+	// connection handshake.
+	Attributes ConnectionAttributes
+
 	bufferedReader *bufio.Reader
 	flushTimer     *time.Timer
 	flushDelay     time.Duration
diff --git a/go/mysql/constants.go b/go/mysql/constants.go
@@ -37,6 +37,10 @@ const (
 // implemented authentication methods.
 type AuthMethodDescription string
 
+// ConnectionAttributes is a map of key/value pairs sent by the client during
+// the connection phase.
+type ConnectionAttributes map[string]string
+
 // Supported auth forms.
 const (
 	// MysqlNativePassword uses a salt and transmits a hash on the wire.
diff --git a/go/mysql/server.go b/go/mysql/server.go
@@ -688,7 +688,7 @@ func (c *Conn) writeHandshakeV10(serverVersion string, authServer AuthServer, ch
 }
 
 // parseClientHandshakePacket parses the handshake sent by the client.
-// Returns the username, auth method, auth data, error.
+// Returns the username, auth method, auth data, connection attributes, error.
 // The original data is not pointed at, and can be freed.
 func (l *Listener) parseClientHandshakePacket(c *Conn, firstTime bool, data []byte) (string, AuthMethodDescription, []byte, error) {
 	pos := 0
@@ -806,58 +806,43 @@ func (l *Listener) parseClientHandshakePacket(c *Conn, firstTime bool, data []by
 
 	// Decode connection attributes send by the client
 	if clientFlags&CapabilityClientConnAttr != 0 {
-		if _, _, err := parseConnAttrs(data, pos); err != nil {
+		clientAttributes, _, err := parseConnAttrs(data, pos)
+		if err != nil {
 			log.Warningf("Decode connection attributes send by the client: %v", err)
 		}
+
+		c.Attributes = clientAttributes
 	}
 
 	return username, AuthMethodDescription(authMethod), authResponse, nil
 }
 
-func parseConnAttrs(data []byte, pos int) (map[string]string, int, error) {
-	var attrLen uint64
+func parseConnAttrs(data []byte, pos int) (ConnectionAttributes, int, error) {
+	attrs := make(map[string]string)
 
 	attrLen, pos, ok := readLenEncInt(data, pos)
 	if !ok {
 		return nil, 0, vterrors.Errorf(vtrpc.Code_INTERNAL, "parseClientHandshakePacket: can't read connection attributes variable length")
 	}
 
-	var attrLenRead uint64
-
-	attrs := make(map[string]string)
-
-	for attrLenRead < attrLen {
-		var keyLen byte
-		keyLen, pos, ok = readByte(data, pos)
-		if !ok {
-			return nil, 0, vterrors.Errorf(vtrpc.Code_INTERNAL, "parseClientHandshakePacket: can't read connection attribute key length")
-		}
-		attrLenRead += uint64(keyLen) + 1
+	addrEndPos := pos + int(attrLen)
 
-		var connAttrKey []byte
-		connAttrKey, pos, ok = readBytes(data, pos, int(keyLen))
+	var key, value string
+	for pos < addrEndPos {
+		key, pos, ok = readLenEncString(data, pos)
 		if !ok {
 			return nil, 0, vterrors.Errorf(vtrpc.Code_INTERNAL, "parseClientHandshakePacket: can't read connection attribute key")
 		}
 
-		var valLen byte
-		valLen, pos, ok = readByte(data, pos)
-		if !ok {
-			return nil, 0, vterrors.Errorf(vtrpc.Code_INTERNAL, "parseClientHandshakePacket: can't read connection attribute value length")
-		}
-		attrLenRead += uint64(valLen) + 1
-
-		var connAttrVal []byte
-		connAttrVal, pos, ok = readBytes(data, pos, int(valLen))
+		value, pos, ok = readLenEncString(data, pos)
 		if !ok {
 			return nil, 0, vterrors.Errorf(vtrpc.Code_INTERNAL, "parseClientHandshakePacket: can't read connection attribute value")
 		}
 
-		attrs[string(connAttrKey[:])] = string(connAttrVal[:])
+		attrs[key] = value
 	}
 
 	return attrs, pos, nil
-
 }
 
 // writeAuthSwitchRequest writes an auth switch request packet.
diff --git a/go/mysql/server_test.go b/go/mysql/server_test.go
@@ -471,6 +471,91 @@ func TestClientFoundRows(t *testing.T) {
 	c.Close()
 }
 
+func TestConnAttrs(t *testing.T) {
+	ctx := utils.LeakCheckContext(t)
+	th := &testHandler{}
+
+	authServer := NewAuthServerStatic("", "", 0)
+	authServer.entries["user1"] = []*AuthServerStaticEntry{{
+		Password: "password1",
+		UserData: "userData1",
+	}}
+	defer authServer.close()
+
+	l, err := NewListener("tcp", "127.0.0.1:", authServer, th, 0, 0, false, false, 0, 0, false)
+	require.NoError(t, err, "NewListener failed")
+	host, port := getHostPort(t, l.Addr())
+
+	// Test with attrs.
+	params := &ConnParams{
+		Host:  host,
+		Port:  port,
+		Uname: "user1",
+		Pass:  "password1",
+	}
+
+	attributes := ConnectionAttributes{
+		"key1": "value1",
+		"k2":   "v2",
+	}
+
+	go l.Accept()
+	defer cleanupListener(ctx, l, params)
+
+	clientConn, err := ConnectWithAttributes(ctx, params, attributes)
+	require.NoError(t, err, "Connect failed")
+
+	serverConn := th.LastConn()
+	assert.Equal(t, uint32(CapabilityClientConnAttr), clientConn.Capabilities&CapabilityClientConnAttr, "ConnAttr flag: %x, bit must be set", th.LastConn().Capabilities)
+	assert.Equal(t, serverConn.Attributes, attributes, "attributes should be sent and parsed")
+
+	clientConn.Close()
+	assert.True(t, clientConn.IsClosed(), "IsClosed should be true on Close-d connection.")
+
+	// Empty attrs do not even set the capability flag
+	params = &ConnParams{
+		Host:  host,
+		Port:  port,
+		Uname: "user1",
+		Pass:  "password1",
+	}
+
+	clientConn, err = Connect(ctx, params)
+	require.NoError(t, err, "Connect failed")
+
+	serverConn = th.LastConn()
+	assert.Equal(t, uint32(0), clientConn.Capabilities&CapabilityClientConnAttr, "ConnAttr flag: %x, bit must not be set", th.LastConn().Capabilities)
+	assert.Equal(t, 0, len(serverConn.Attributes), "attributes should be empty")
+
+	clientConn.Close()
+	assert.True(t, clientConn.IsClosed(), "IsClosed should be true on Close-d connection.")
+
+	// Test long attributes more than 255 bytes
+	params = &ConnParams{
+		Host:  host,
+		Port:  port,
+		Uname: "user1",
+		Pass:  "password1",
+	}
+
+	longAttributes := ConnectionAttributes{
+		"short":  strings.Repeat("a", 10),
+		"long":   strings.Repeat("b", 256),
+		"longer": strings.Repeat("c", 1024*1024),
+	}
+
+	clientConn, err = ConnectWithAttributes(ctx, params, longAttributes)
+	require.NoError(t, err, "Connect failed")
+
+	serverConn = th.LastConn()
+	assert.Equal(t, uint32(CapabilityClientConnAttr), clientConn.Capabilities&CapabilityClientConnAttr, "ConnAttr flag: %x, bit must be set", th.LastConn().Capabilities)
+	assert.Equal(t, serverConn.Attributes, longAttributes, "attributes should be sent and parsed")
+
+	clientConn.Close()
+	assert.True(t, clientConn.IsClosed(), "IsClosed should be true on Close-d connection.")
+
+}
+
 func TestConnCounts(t *testing.T) {
 	th := &testHandler{}
 
