Fix query attribute lifecycle and prepared statement encoding

Author: IuraHD

client/conn.go

@@ -715,7 +715,7 @@ func (c *Conn) exec(query string) (*mysql.Result, error) {
 // https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_com_query.html
 func (c *Conn) execSend(query string) error {
 	var buf bytes.Buffer
-	defer clear(c.queryAttributes)
+	defer func() { c.queryAttributes = nil }()
 
 	if c.capability&mysql.CLIENT_QUERY_ATTRIBUTES > 0 {
 		if c.includeLine >= 0 {

client/query_attribute_lifecycle_test.go

@@ -0,0 +1,107 @@
+package client
+
+import (
+	"bytes"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/go-mysql-org/go-mysql/mysql"
+	"github.com/go-mysql-org/go-mysql/packet"
+)
+
+func captureAttributeRequest(t *testing.T, c *Conn, peer *packet.Conn, send func() error) []byte {
+	t.Helper()
+	done := make(chan error, 1)
+	go func() { done <- send() }()
+	peer.ResetSequence()
+	payload, err := peer.ReadPacket()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := peer.WritePacket([]byte{0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0}); err != nil {
+		t.Fatal(err)
+	}
+	if err := <-done; err != nil {
+		t.Fatal(err)
+	}
+	return payload
+}
+
+func attributeConnection(t *testing.T) (*Conn, *packet.Conn) {
+	a, b := net.Pipe()
+	_ = a.SetDeadline(time.Now().Add(5 * time.Second))
+	_ = b.SetDeadline(time.Now().Add(5 * time.Second))
+	t.Cleanup(func() { a.Close(); b.Close() })
+	return &Conn{Conn: packet.NewConn(a), capability: mysql.CLIENT_QUERY_ATTRIBUTES, includeLine: -1}, packet.NewConn(b)
+}
+
+func TestQueryAttributeAttributesConsumed(t *testing.T) {
+	c, p := attributeConnection(t)
+	c.SetQueryAttributes(mysql.QueryAttribute{Name: "trace", Value: "first"})
+	first := captureAttributeRequest(t, c, p, func() error { _, err := c.Execute("SELECT 1"); return err })
+	second := captureAttributeRequest(t, c, p, func() error { _, err := c.Execute("SELECT 2"); return err })
+	t.Logf("first=%x second=%x", first, second)
+	if second[1] != 0 {
+		t.Fatalf("next query still announces %d attributes; want zero", second[1])
+	}
+}
+
+func TestQueryAttributeZeroParameterAttributes(t *testing.T) {
+	c, p := attributeConnection(t)
+	c.SetQueryAttributes(mysql.QueryAttribute{Name: "trace", Value: "marker"})
+	s := &Stmt{conn: c}
+	data := captureAttributeRequest(t, c, p, func() error { _, err := s.Execute(); return err })
+	t.Logf("execute=%x", data)
+	if !bytes.Contains(data, []byte("marker")) {
+		t.Fatal("zero-parameter execution lost its query attribute")
+	}
+}
+
+func TestQueryAttributeNullParameterAttributes(t *testing.T) {
+	c, p := attributeConnection(t)
+	c.SetQueryAttributes(mysql.QueryAttribute{Name: "trace", Value: "marker"})
+	s := &Stmt{conn: c}
+	s.Params = 1
+	data := captureAttributeRequest(t, c, p, func() error { _, err := s.Execute(nil); return err })
+	t.Logf("execute=%x", data)
+	if !bytes.Contains(data, []byte("marker")) {
+		t.Fatal("NULL-parameter execution lost its query attribute")
+	}
+}
+
+func TestQueryAttributeCallerAttributesReusable(t *testing.T) {
+	c, p := attributeConnection(t)
+	attrs := []mysql.QueryAttribute{{Name: "trace", Value: "marker"}}
+	c.SetQueryAttributes(attrs...)
+	captureAttributeRequest(t, c, p, func() error { _, err := c.Execute("SELECT 1"); return err })
+	if attrs[0].Name != "trace" || attrs[0].Value != "marker" {
+		t.Fatal("sending query mutated caller attribute slice")
+	}
+}
+
+func TestQueryAttributeLineAttributeDoesNotAccumulate(t *testing.T) {
+	c, p := attributeConnection(t)
+	c.IncludeLine(0)
+	for i := 0; i < 3; i++ {
+		data := captureAttributeRequest(t, c, p, func() error { _, err := c.Execute("SELECT 1"); return err })
+		if data[1] != 1 {
+			t.Fatalf("query %d has %d attributes; want one line attribute", i, data[1])
+		}
+	}
+}
+
+func TestQueryAttributeUnsupportedCapability(t *testing.T) {
+	c, p := attributeConnection(t)
+	c.capability = 0
+	c.SetQueryAttributes(mysql.QueryAttribute{Name: "trace", Value: "marker"})
+	s := &Stmt{conn: c}
+	s.Params = 1
+	data := captureAttributeRequest(t, c, p, func() error { _, err := s.Execute(7); return err })
+	if data[5] != 0 {
+		t.Fatalf("legacy execute has unsupported flags %x", data[5])
+	}
+	if bytes.Contains(data, []byte("marker")) {
+		t.Fatal("attributes encoded without negotiated support")
+	}
+}

client/stmt.go

@@ -102,7 +102,7 @@ func (s *Stmt) Close() error {
 
 // https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_com_stmt_execute.html
 func (s *Stmt) write(args ...any) error {
-	defer clear(s.conn.queryAttributes)
+	defer func() { s.conn.queryAttributes = nil }()
 	paramsNum := s.Params
 
 	if len(args) != paramsNum {
@@ -121,6 +121,9 @@ func (s *Stmt) write(args ...any) error {
 	}
 
 	qaLen := len(s.conn.queryAttributes)
+	if s.conn.capability&mysql.CLIENT_QUERY_ATTRIBUTES == 0 {
+		qaLen = 0
+	}
 	paramTypes := make([][]byte, paramsNum+qaLen)
 	paramFlags := make([][]byte, paramsNum+qaLen)
 	paramValues := make([][]byte, paramsNum+qaLen)
@@ -215,12 +218,13 @@ func (s *Stmt) write(args ...any) error {
 
 		length += len(paramValues[i])
 	}
-	for i, qa := range s.conn.queryAttributes {
+	for i, qa := range s.conn.queryAttributes[:qaLen] {
 		tf := qa.TypeAndFlag()
 		paramTypes[(i + paramsNum)] = []byte{tf[0]}
 		paramFlags[i+paramsNum] = []byte{tf[1]}
 		paramValues[i+paramsNum] = qa.ValueBytes()
 		paramNames[i+paramsNum] = mysql.PutLengthEncodedString([]byte(qa.Name))
+		newParamBoundFlag = 1
 	}
 
 	data := utils.BytesBufferGet()
@@ -236,7 +240,7 @@ func (s *Stmt) write(args ...any) error {
 	data.Write([]byte{byte(s.ID), byte(s.ID >> 8), byte(s.ID >> 16), byte(s.ID >> 24)})
 
 	flags := mysql.CURSOR_TYPE_NO_CURSOR
-	if paramsNum > 0 {
+	if s.conn.capability&mysql.CLIENT_QUERY_ATTRIBUTES > 0 && paramsNum+qaLen > 0 {
 		flags |= mysql.PARAMETER_COUNT_AVAILABLE
 	}
 	data.WriteByte(flags)
@@ -246,7 +250,7 @@ func (s *Stmt) write(args ...any) error {
 
 	if paramsNum > 0 || (s.conn.capability&mysql.CLIENT_QUERY_ATTRIBUTES > 0 && (flags&mysql.PARAMETER_COUNT_AVAILABLE > 0)) {
 		if s.conn.capability&mysql.CLIENT_QUERY_ATTRIBUTES > 0 {
-			paramsNum += len(s.conn.queryAttributes)
+			paramsNum += qaLen
 			data.Write(mysql.PutLengthEncodedInt(uint64(paramsNum)))
 		}
 		if paramsNum > 0 {