Enable sending connection attributes

Summary: In order to be able track Go client usage we enable sending connection attributes as specified in https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_connection_phase_packets_protocol_handshake_response.html (CLIENT_CONNECT_ATTRS). This commit is based on PR go-sql-driver#1241 which is based on PR go-sql-driver#737 (which are still open as of 11.04.2023)

Test Plan:
manual
integration test will be added to the main repo

Reviewers: amakarovych-ua, okramarenko-ua

Reviewed By: okramarenko-ua

Subscribers: engineering-list

Differential Revision: https://grizzly.internal.memcompute.com/D61981

Author: Pavlo

.arcconfig

@@ -0,0 +1,5 @@
+{
+  "project_id" : "go-singlestore-driver",
+  "conduit_uri" : "https:\/\/grizzly.internal.memcompute.com\/api\/"
+}
+

AUTHORS

@@ -16,6 +16,7 @@ Achille Roussel <achille.roussel at gmail.com>
 Alex Snast <alexsn at fb.com>
 Alexey Palazhchenko <alexey.palazhchenko at gmail.com>
 Andrew Reid <andrew.reid at tixtrack.com>
+Andy Grunwald <andygrunwald at gmail.com>
 Animesh Ray <mail.rayanimesh at gmail.com>
 Arne Hormann <arnehormann at gmail.com>
 Ariel Mashraki <ariel at mashraki.co.il>
@@ -92,6 +93,7 @@ Tan Jinhua <312841925 at qq.com>
 Thomas Wodarek <wodarekwebpage at gmail.com>
 Tim Ruffles <timruffles at gmail.com>
 Tom Jenkinson <tom at tjenkinson.me>
+Vasily Fedoseyev <vasilyfedoseyev at gmail.com>
 Vladimir Kovpak <cn007b at gmail.com>
 Vladyslav Zhelezniak <zhvladi at gmail.com>
 Xiangyu Hu <xiangyu.hu at outlook.com>
@@ -114,5 +116,6 @@ Keybase Inc.
 Multiplay Ltd.
 Percona LLC
 Pivotal Inc.
+SingleStore Inc.
 Stripe Inc.
 Zendesk Inc.

README.md

@@ -242,6 +242,16 @@ SELECT u.id FROM users as u
 
 will return `u.id` instead of just `id` if `columnsWithAlias=true`.
 
+##### `connectAttrs`
+
+```
+Type:           map
+Valid Values:   comma-separated list of attribute:value pairs
+Default:        empty
+```
+
+Allows setting of connection attributes, for example `connectAttrs=program_name:YourProgramName,program_version:1.2.3`. They are sent in the [Protocol::HandshakeResponse41](https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_connection_phase_packets_protocol_handshake_response.html) packet. Attributes and their values must not contain `:` and `,` symbols.
+
 ##### `interpolateParams`
 
 ```

driver_test.go

@@ -77,8 +77,11 @@ func init() {
 	dsn = fmt.Sprintf("%s:%s@%s/%s?timeout=30s", user, pass, netAddr, dbname)
 	c, err := net.Dial(prot, addr)
 	if err == nil {
+		fmt.Println("Running driver tests")
 		available = true
 		c.Close()
+	} else {
+		fmt.Printf("Not running driver tests. %s is not connectable\n", addr)
 	}
 }
 

dsn.go

@@ -41,6 +41,7 @@ type Config struct {
 	DBName           string            // Database name
 	Params           map[string]string // Connection parameters
 	Collation        string            // Connection collation
+	ConnectAttrs     map[string]string // Connection attributes
 	Loc              *time.Location    // Location for time.Time values
 	MaxAllowedPacket int               // Max packet size allowed
 	ServerPubKey     string            // Server public key name
@@ -272,6 +273,30 @@ func (cfg *Config) FormatDSN() string {
 		writeDSNParam(&buf, &hasParam, "maxAllowedPacket", strconv.Itoa(cfg.MaxAllowedPacket))
 	}
 
+	if len(cfg.ConnectAttrs) > 0 {
+		// connectAttrs=program_name:Login Server,other_name:other
+		if hasParam {
+			buf.WriteString("&connectAttrs=")
+		} else {
+			hasParam = true
+			buf.WriteString("?connectAttrs=")
+		}
+
+		var attr_names []string
+		for attr_name := range cfg.ConnectAttrs {
+			attr_names = append(attr_names, attr_name)
+		}
+		sort.Strings(attr_names)
+		for index, attr_name := range attr_names {
+			if index > 0 {
+				buf.WriteByte(',')
+			}
+			buf.WriteString(url.QueryEscape(attr_name))
+			buf.WriteByte(':')
+			buf.WriteString(url.QueryEscape(cfg.ConnectAttrs[attr_name]))
+		}
+	}
+
 	// other params
 	if cfg.Params != nil {
 		var params []string
@@ -537,6 +562,24 @@ func parseDSNParams(cfg *Config, params string) (err error) {
 			if err != nil {
 				return
 			}
+		case "connectAttrs":
+			if cfg.ConnectAttrs == nil {
+				cfg.ConnectAttrs = make(map[string]string)
+			}
+
+			var ConnectAttrs string
+			if ConnectAttrs, err = url.QueryUnescape(value); err != nil {
+				return
+			}
+
+			// program_name:Name,foo:bar
+			for _, attr_str := range strings.Split(ConnectAttrs, ",") {
+				attr := strings.SplitN(attr_str, ":", 2)
+				if len(attr) != 2 {
+					continue
+				}
+				cfg.ConnectAttrs[attr[0]] = attr[1]
+			}
 		default:
 			// lazy init
 			if cfg.Params == nil {

dsn_test.go

@@ -71,6 +71,9 @@ var testDSNs = []struct {
 }, {
 	"tcp(de:ad:be:ef::ca:fe)/dbname",
 	&Config{Net: "tcp", Addr: "[de:ad:be:ef::ca:fe]:3306", DBName: "dbname", Collation: "utf8mb4_general_ci", Loc: time.UTC, MaxAllowedPacket: defaultMaxAllowedPacket, AllowNativePasswords: true, CheckConnLiveness: true},
+}, {
+	"tcp(127.0.0.1)/dbname?connectAttrs=program_name:SomeService",
+	&Config{Net: "tcp", Addr: "127.0.0.1:3306", DBName: "dbname", ConnectAttrs: map[string]string{"program_name": "SomeService"}, Collation: "utf8mb4_general_ci", Loc: time.UTC, MaxAllowedPacket: defaultMaxAllowedPacket, AllowNativePasswords: true, CheckConnLiveness: true},
 },
 }
 

packets.go

@@ -18,6 +18,9 @@ import (
 	"fmt"
 	"io"
 	"math"
+	"os"
+	"runtime"
+	"strconv"
 	"time"
 )
 
@@ -235,10 +238,15 @@ func (mc *mysqlConn) readHandshakePacket() (data []byte, plugin string, err erro
 	if len(data) > pos {
 		// character set [1 byte]
 		// status flags [2 bytes]
+		// ignore character set and status 
+		pos += 1 + 2
 		// capability flags (upper 2 bytes) [2 bytes]
+		mc.flags |= clientFlag(uint32(binary.LittleEndian.Uint16(data[pos:pos+2])) << 16)
+		pos += 2
+
 		// length of auth-plugin-data [1 byte]
 		// reserved (all [00]) [10 bytes]
-		pos += 1 + 2 + 2 + 1 + 10
+		pos += 1 + 10
 
 		// second part of the password cipher [mininum 13 bytes],
 		// where len=MAX(13, length of auth-plugin-data - 8)
@@ -312,6 +320,34 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
 	}
 
 	pktLen := 4 + 4 + 1 + 23 + len(mc.cfg.User) + 1 + len(authRespLEI) + len(authResp) + 21 + 1
+	const MAX_CONN_ATTRS_LEN = 2000
+
+	connectAttrsBuf := make([]byte, 0, MAX_CONN_ATTRS_LEN)
+	if mc.flags&clientConnectAttrs != 0 {
+		clientFlags |= clientConnectAttrs
+
+		// Set default connection attributes
+		// See https://dev.mysql.com/doc/refman/8.0/en/performance-schema-connection-attribute-tables.html#performance-schema-connection-attributes-available
+		connectAttrsBuf = appendLengthEncodedString(connectAttrsBuf, []byte("_client_name"))
+		connectAttrsBuf = appendLengthEncodedString(connectAttrsBuf, []byte("memsql/go-singlestore-driver"))
+
+		connectAttrsBuf = appendLengthEncodedString(connectAttrsBuf, []byte("_os"))
+		connectAttrsBuf = appendLengthEncodedString(connectAttrsBuf, []byte(runtime.GOOS))
+
+		connectAttrsBuf = appendLengthEncodedString(connectAttrsBuf, []byte("_platform"))
+		connectAttrsBuf = appendLengthEncodedString(connectAttrsBuf, []byte(runtime.GOARCH))
+
+		connectAttrsBuf = appendLengthEncodedString(connectAttrsBuf, []byte("_pid"))
+		connectAttrsBuf = appendLengthEncodedString(connectAttrsBuf, []byte(strconv.Itoa(os.Getpid())))
+
+		for k, v := range mc.cfg.ConnectAttrs {
+			connectAttrsBuf = appendLengthEncodedString(connectAttrsBuf, []byte(k))
+			connectAttrsBuf = appendLengthEncodedString(connectAttrsBuf, []byte(v))
+		}
+		// reserve additional 8 bytes to write connectAttrsBuf length
+		connectAttrsBuf = appendLengthEncodedString(make([]byte, 0, 8 + MAX_CONN_ATTRS_LEN), connectAttrsBuf)
+		pktLen += len(connectAttrsBuf)
+	}
 
 	// To specify a db name
 	if n := len(mc.cfg.DBName); n > 0 {
@@ -395,6 +431,10 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
 	data[pos] = 0x00
 	pos++
 
+	if clientFlags&clientConnectAttrs != 0 {
+		pos += copy(data[pos:], connectAttrsBuf)
+	}
+
 	// Send Auth packet
 	return mc.writePacket(data[:pos])
 }

utils.go

@@ -626,6 +626,12 @@ func appendLengthEncodedInteger(b []byte, n uint64) []byte {
 		byte(n>>32), byte(n>>40), byte(n>>48), byte(n>>56))
 }
 
+// encodes a bytes slice with prepended length-encoded size and appends it to the given bytes slice
+func appendLengthEncodedString(b []byte, str []byte) []byte {
+	b = appendLengthEncodedInteger(b, uint64(len(str)))
+	return append(b, str...)
+}
+
 // reserveBuffer checks cap(buf) and expand buffer to len(buf) + appendSize.
 // If cap(buf) is not enough, reallocate new buffer.
 func reserveBuffer(buf []byte, appendSize int) []byte {