sha256_password auth support (#808)

* auth: add sha256_password implementation

* auth: fix sha256_password implementation

* auth: add sha256_password tests

* auth: remove test pub key by ascii encoded version

* packets: allow auth responses longer than 255 bytes

* utils: correct naming of registries

* auth: allow using a local server pub key

README.md

@@ -301,6 +301,19 @@ other cases. You should ensure your application will never cause an ERROR 1290
 except for `read-only` mode when enabling this option.
 
 
+##### `serverPubKey`
+
+```
+Type:           string
+Valid Values:   <name>
+Default:        none
+```
+
+Server public keys can be registered with [`mysql.RegisterServerPubKey`](https://godoc.org/github.com/go-sql-driver/mysql#RegisterServerPubKey), which can then be used by the assigned name in the DSN.
+Public keys are used to transmit encrypted data, e.g. for authentication.
+If the server's public key is known, it should be set manually to avoid expensive and potentially insecure transmissions of the public key from the server to the client each time it is required.
+
+
 ##### `timeout`
 
 ```

auth.go

@@ -15,8 +15,72 @@ import (
 	"crypto/sha256"
 	"crypto/x509"
 	"encoding/pem"
+	"sync"
 )
 
+// server pub keys registry
+var (
+	serverPubKeyLock     sync.RWMutex
+	serverPubKeyRegistry map[string]*rsa.PublicKey
+)
+
+// RegisterServerPubKey registers a server RSA public key which can be used to
+// send data in a secure manner to the server without receiving the public key
+// in a potentially insecure way from the server first.
+// Registered keys can afterwards be used adding serverPubKey=<name> to the DSN.
+//
+// Note: The provided rsa.PublicKey instance is exclusively owned by the driver
+// after registering it and may not be modified.
+//
+//  data, err := ioutil.ReadFile("mykey.pem")
+//  if err != nil {
+//  	log.Fatal(err)
+//  }
+//
+//  block, _ := pem.Decode(data)
+//  if block == nil || block.Type != "PUBLIC KEY" {
+//  	log.Fatal("failed to decode PEM block containing public key")
+//  }
+//
+//  pub, err := x509.ParsePKIXPublicKey(block.Bytes)
+//  if err != nil {
+//  	log.Fatal(err)
+//  }
+//
+//  if rsaPubKey, ok := pub.(*rsa.PublicKey); ok {
+//  	mysql.RegisterServerPubKey("mykey", rsaPubKey)
+//  } else {
+//  	log.Fatal("not a RSA public key")
+//  }
+//
+func RegisterServerPubKey(name string, pubKey *rsa.PublicKey) {
+	serverPubKeyLock.Lock()
+	if serverPubKeyRegistry == nil {
+		serverPubKeyRegistry = make(map[string]*rsa.PublicKey)
+	}
+
+	serverPubKeyRegistry[name] = pubKey
+	serverPubKeyLock.Unlock()
+}
+
+// DeregisterServerPubKey removes the public key registered with the given name.
+func DeregisterServerPubKey(name string) {
+	serverPubKeyLock.Lock()
+	if serverPubKeyRegistry != nil {
+		delete(serverPubKeyRegistry, name)
+	}
+	serverPubKeyLock.Unlock()
+}
+
+func getServerPubKey(name string) (pubKey *rsa.PublicKey) {
+	serverPubKeyLock.RLock()
+	if v, ok := serverPubKeyRegistry[name]; ok {
+		pubKey = v
+	}
+	serverPubKeyLock.RUnlock()
+	return
+}
+
 // Hash password using pre 4.1 (old password) method
 // https://github.com/atcurtis/mariadb/blob/master/mysys/my_rnd.c
 type myRnd struct {
@@ -154,6 +218,25 @@ func scrambleSHA256Password(scramble []byte, password string) []byte {
 	return message1
 }
 
+func encryptPassword(password string, seed []byte, pub *rsa.PublicKey) ([]byte, error) {
+	plain := make([]byte, len(password)+1)
+	copy(plain, password)
+	for i := range plain {
+		j := i % len(seed)
+		plain[i] ^= seed[j]
+	}
+	sha1 := sha1.New()
+	return rsa.EncryptOAEP(sha1, rand.Reader, pub, plain, nil)
+}
+
+func (mc *mysqlConn) sendEncryptedPassword(seed []byte, pub *rsa.PublicKey) error {
+	enc, err := encryptPassword(mc.cfg.Passwd, seed, pub)
+	if err != nil {
+		return err
+	}
+	return mc.writeAuthSwitchPacket(enc, false)
+}
+
 func (mc *mysqlConn) auth(authData []byte, plugin string) ([]byte, bool, error) {
 	switch plugin {
 	case "caching_sha2_password":
@@ -187,6 +270,25 @@ func (mc *mysqlConn) auth(authData []byte, plugin string) ([]byte, bool, error)
 		authResp := scramblePassword(authData[:20], mc.cfg.Passwd)
 		return authResp, false, nil
 
+	case "sha256_password":
+		if len(mc.cfg.Passwd) == 0 {
+			return nil, true, nil
+		}
+		if mc.cfg.tls != nil || mc.cfg.Net == "unix" {
+			// write cleartext auth packet
+			return []byte(mc.cfg.Passwd), true, nil
+		}
+
+		pubKey := mc.cfg.pubKey
+		if pubKey == nil {
+			// request public key from server
+			return []byte{1}, false, nil
+		}
+
+		// encrypted password
+		enc, err := encryptPassword(mc.cfg.Passwd, authData, pubKey)
+		return enc, false, err
+
 	default:
 		errLog.Print("unknown auth plugin:", plugin)
 		return nil, false, ErrUnknownPlugin
@@ -206,6 +308,9 @@ func (mc *mysqlConn) handleAuthResult(oldAuthData []byte, plugin string) error {
 		// sent and we have to keep using the cipher sent in the init packet.
 		if authData == nil {
 			authData = oldAuthData
+		} else {
+			// copy data from read buffer to owned slice
+			copy(oldAuthData, authData)
 		}
 
 		plugin = newPlugin
@@ -223,6 +328,7 @@ func (mc *mysqlConn) handleAuthResult(oldAuthData []byte, plugin string) error {
 		if err != nil {
 			return err
 		}
+
 		// Do not allow to change the auth plugin more than once
 		if newPlugin != "" {
 			return ErrMalformPkt
@@ -251,48 +357,34 @@ func (mc *mysqlConn) handleAuthResult(oldAuthData []byte, plugin string) error {
 						return err
 					}
 				} else {
-					seed := oldAuthData
-
-					// TODO: allow to specify a local file with the pub key via
-					// the DSN
-
-					// request public key
-					data := mc.buf.takeSmallBuffer(4 + 1)
-					data[4] = cachingSha2PasswordRequestPublicKey
-					mc.writePacket(data)
-
-					// parse public key
-					data, err := mc.readPacket()
-					if err != nil {
-						return err
-					}
-
-					block, _ := pem.Decode(data[1:])
-					pub, err := x509.ParsePKIXPublicKey(block.Bytes)
-					if err != nil {
-						return err
+					pubKey := mc.cfg.pubKey
+					if pubKey == nil {
+						// request public key from server
+						data := mc.buf.takeSmallBuffer(4 + 1)
+						data[4] = cachingSha2PasswordRequestPublicKey
+						mc.writePacket(data)
+
+						// parse public key
+						data, err := mc.readPacket()
+						if err != nil {
+							return err
+						}
+
+						block, _ := pem.Decode(data[1:])
+						pkix, err := x509.ParsePKIXPublicKey(block.Bytes)
+						if err != nil {
+							return err
+						}
+						pubKey = pkix.(*rsa.PublicKey)
 					}
 
 					// send encrypted password
-					plain := make([]byte, len(mc.cfg.Passwd)+1)
-					copy(plain, mc.cfg.Passwd)
-					for i := range plain {
-						j := i % len(seed)
-						plain[i] ^= seed[j]
-					}
-					sha1 := sha1.New()
-					enc, err := rsa.EncryptOAEP(sha1, rand.Reader, pub.(*rsa.PublicKey), plain, nil)
+					err = mc.sendEncryptedPassword(oldAuthData, pubKey)
 					if err != nil {
 						return err
 					}
-
-					if err = mc.writeAuthSwitchPacket(enc, false); err != nil {
-						return err
-					}
-				}
-				if err = mc.readResultOK(); err == nil {
-					return nil // auth successful
 				}
+				return mc.readResultOK()
 
 			default:
 				return ErrMalformPkt
@@ -301,6 +393,25 @@ func (mc *mysqlConn) handleAuthResult(oldAuthData []byte, plugin string) error {
 			return ErrMalformPkt
 		}
 
+	case "sha256_password":
+		switch len(authData) {
+		case 0:
+			return nil // auth successful
+		default:
+			block, _ := pem.Decode(authData)
+			pub, err := x509.ParsePKIXPublicKey(block.Bytes)
+			if err != nil {
+				return err
+			}
+
+			// send encrypted password
+			err = mc.sendEncryptedPassword(oldAuthData, pub.(*rsa.PublicKey))
+			if err != nil {
+				return err
+			}
+			return mc.readResultOK()
+		}
+
 	default:
 		return nil // auth successful
 	}

auth_test.go

@@ -10,38 +10,34 @@ package mysql
 
 import (
 	"bytes"
+	"crypto/rsa"
 	"crypto/tls"
+	"crypto/x509"
+	"encoding/pem"
 	"fmt"
 	"testing"
 )
 
-var serverPubKey = []byte{1, 45, 45, 45, 45, 45, 66, 69, 71, 73, 78, 32, 80, 85,
-	66, 76, 73, 67, 32, 75, 69, 89, 45, 45, 45, 45, 45, 10, 77, 73, 73, 66, 73,
-	106, 65, 78, 66, 103, 107, 113, 104, 107, 105, 71, 57, 119, 48, 66, 65, 81,
-	69, 70, 65, 65, 79, 67, 65, 81, 56, 65, 77, 73, 73, 66, 67, 103, 75, 67, 65,
-	81, 69, 65, 51, 72, 115, 120, 83, 53, 80, 47, 72, 97, 88, 80, 118, 109, 51,
-	109, 50, 65, 68, 110, 10, 98, 117, 54, 71, 81, 102, 112, 83, 71, 111, 55,
-	104, 50, 103, 104, 56, 49, 112, 109, 97, 120, 107, 67, 110, 68, 67, 119,
-	102, 54, 109, 109, 101, 72, 55, 76, 75, 104, 115, 110, 89, 110, 78, 52, 81,
-	48, 99, 122, 49, 81, 69, 47, 98, 104, 100, 80, 117, 54, 106, 115, 43, 86,
-	97, 89, 52, 10, 67, 99, 77, 117, 98, 80, 78, 49, 103, 79, 75, 97, 89, 118,
-	78, 99, 103, 69, 87, 112, 116, 73, 67, 105, 50, 88, 84, 116, 116, 66, 55,
-	117, 104, 43, 118, 67, 77, 106, 76, 118, 106, 65, 77, 100, 54, 47, 68, 109,
-	120, 100, 98, 85, 66, 48, 122, 80, 71, 113, 68, 79, 103, 105, 76, 68, 10,
-	75, 82, 79, 79, 53, 113, 100, 55, 115, 104, 98, 55, 49, 82, 47, 88, 74, 69,
-	70, 118, 76, 120, 71, 88, 69, 70, 48, 90, 116, 104, 72, 101, 78, 111, 57,
-	102, 69, 118, 120, 70, 81, 111, 109, 98, 49, 107, 90, 57, 74, 56, 110, 66,
-	119, 116, 101, 53, 83, 70, 53, 89, 108, 113, 86, 50, 10, 66, 66, 53, 113,
-	108, 97, 122, 43, 51, 81, 83, 78, 118, 109, 67, 49, 105, 87, 102, 108, 106,
-	88, 98, 89, 53, 107, 51, 47, 97, 54, 109, 107, 77, 47, 76, 97, 87, 104, 97,
-	117, 78, 53, 80, 82, 51, 115, 67, 120, 53, 85, 117, 49, 77, 102, 100, 115,
-	86, 105, 107, 53, 102, 88, 77, 77, 10, 100, 120, 107, 102, 70, 43, 88, 51,
-	99, 104, 107, 65, 110, 119, 73, 51, 70, 117, 119, 119, 50, 87, 71, 109, 87,
-	79, 71, 98, 75, 116, 109, 73, 101, 85, 109, 51, 98, 73, 82, 109, 100, 70,
-	85, 113, 97, 108, 81, 105, 70, 104, 113, 101, 90, 50, 105, 107, 106, 104,
-	103, 86, 73, 57, 112, 76, 10, 119, 81, 73, 68, 65, 81, 65, 66, 10, 45, 45,
-	45, 45, 45, 69, 78, 68, 32, 80, 85, 66, 76, 73, 67, 32, 75, 69, 89, 45, 45,
-	45, 45, 45, 10}
+var testPubKey = []byte("-----BEGIN PUBLIC KEY-----\n" +
+	"MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAol0Z8G8U+25Btxk/g/fm\n" +
+	"UAW/wEKjQCTjkibDE4B+qkuWeiumg6miIRhtilU6m9BFmLQSy1ltYQuu4k17A4tQ\n" +
+	"rIPpOQYZges/qsDFkZh3wyK5jL5WEFVdOasf6wsfszExnPmcZS4axxoYJfiuilrN\n" +
+	"hnwinBAqfi3S0sw5MpSI4Zl1AbOrHG4zDI62Gti2PKiMGyYDZTS9xPrBLbN95Kby\n" +
+	"FFclQLEzA9RJcS1nHFsWtRgHjGPhhjCQxEm9NQ1nePFhCfBfApyfH1VM2VCOQum6\n" +
+	"Ci9bMuHWjTjckC84mzF99kOxOWVU7mwS6gnJqBzpuz8t3zq8/iQ2y7QrmZV+jTJP\n" +
+	"WQIDAQAB\n" +
+	"-----END PUBLIC KEY-----\n")
+
+var testPubKeyRSA *rsa.PublicKey
+
+func init() {
+	block, _ := pem.Decode(testPubKey)
+	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
+	if err != nil {
+		panic(err)
+	}
+	testPubKeyRSA = pub.(*rsa.PublicKey)
+}
 
 func TestScrambleOldPass(t *testing.T) {
 	scramble := []byte{9, 8, 7, 6, 5, 4, 3, 2}
@@ -204,7 +200,7 @@ func TestAuthFastCachingSHA256PasswordFullRSA(t *testing.T) {
 	}
 	conn.queuedReplies = [][]byte{
 		// pub key response
-		append([]byte{byte(len(serverPubKey)), 1, 0, 4}, serverPubKey...),
+		append([]byte{byte(1 + len(testPubKey)), 1, 0, 4, 1}, testPubKey...),
 
 		// OK
 		{7, 0, 0, 6, 0, 0, 0, 2, 0, 0, 0},
@@ -221,6 +217,59 @@ func TestAuthFastCachingSHA256PasswordFullRSA(t *testing.T) {
 	}
 }
 
+func TestAuthFastCachingSHA256PasswordFullRSAWithKey(t *testing.T) {
+	conn, mc := newRWMockConn(1)
+	mc.cfg.User = "root"
+	mc.cfg.Passwd = "secret"
+	mc.cfg.pubKey = testPubKeyRSA
+
+	authData := []byte{6, 81, 96, 114, 14, 42, 50, 30, 76, 47, 1, 95, 126, 81,
+		62, 94, 83, 80, 52, 85}
+	plugin := "caching_sha2_password"
+
+	// Send Client Authentication Packet
+	authResp, addNUL, err := mc.auth(authData, plugin)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = mc.writeHandshakeResponsePacket(authResp, addNUL, plugin)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// check written auth response
+	authRespStart := 4 + 4 + 4 + 1 + 23 + len(mc.cfg.User) + 1
+	authRespEnd := authRespStart + 1 + len(authResp)
+	writtenAuthRespLen := conn.written[authRespStart]
+	writtenAuthResp := conn.written[authRespStart+1 : authRespEnd]
+	expectedAuthResp := []byte{171, 201, 138, 146, 89, 159, 11, 170, 0, 67, 165,
+		49, 175, 94, 218, 68, 177, 109, 110, 86, 34, 33, 44, 190, 67, 240, 70,
+		110, 40, 139, 124, 41}
+	if writtenAuthRespLen != 32 || !bytes.Equal(writtenAuthResp, expectedAuthResp) {
+		t.Fatalf("unexpected written auth response (%d bytes): %v", writtenAuthRespLen, writtenAuthResp)
+	}
+	conn.written = nil
+
+	// auth response
+	conn.data = []byte{
+		2, 0, 0, 2, 1, 4, // Perform Full Authentication
+	}
+	conn.queuedReplies = [][]byte{
+		// OK
+		{7, 0, 0, 4, 0, 0, 0, 2, 0, 0, 0},
+	}
+	conn.maxReads = 2
+
+	// Handle response to auth packet
+	if err := mc.handleAuthResult(authData, plugin); err != nil {
+		t.Errorf("got error: %v", err)
+	}
+
+	if !bytes.HasPrefix(conn.written, []byte{0, 1, 0, 3}) {
+		t.Errorf("unexpected written data: %v", conn.written)
+	}
+}
+
 func TestAuthFastCachingSHA256PasswordFullSecure(t *testing.T) {
 	conn, mc := newRWMockConn(1)
 	mc.cfg.User = "root"
@@ -481,6 +530,183 @@ func TestAuthFastNativePasswordEmpty(t *testing.T) {
 	}
 }
 
+func TestAuthFastSHA256PasswordEmpty(t *testing.T) {
+	conn, mc := newRWMockConn(1)
+	mc.cfg.User = "root"
+	mc.cfg.Passwd = ""
+
+	authData := []byte{6, 81, 96, 114, 14, 42, 50, 30, 76, 47, 1, 95, 126, 81,
+		62, 94, 83, 80, 52, 85}
+	plugin := "sha256_password"
+
+	// Send Client Authentication Packet
+	authResp, addNUL, err := mc.auth(authData, plugin)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = mc.writeHandshakeResponsePacket(authResp, addNUL, plugin)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// check written auth response
+	authRespStart := 4 + 4 + 4 + 1 + 23 + len(mc.cfg.User) + 1
+	authRespEnd := authRespStart + 1 + len(authResp)
+	writtenAuthRespLen := conn.written[authRespStart]
+	writtenAuthResp := conn.written[authRespStart+1 : authRespEnd]
+	if writtenAuthRespLen != 0 {
+		t.Fatalf("unexpected written auth response (%d bytes): %v", writtenAuthRespLen, writtenAuthResp)
+	}
+	conn.written = nil
+
+	// auth response (pub key response)
+	conn.data = append([]byte{byte(1 + len(testPubKey)), 1, 0, 2, 1}, testPubKey...)
+	conn.queuedReplies = [][]byte{
+		// OK
+		{7, 0, 0, 4, 0, 0, 0, 2, 0, 0, 0},
+	}
+	conn.maxReads = 2
+
+	// Handle response to auth packet
+	if err := mc.handleAuthResult(authData, plugin); err != nil {
+		t.Errorf("got error: %v", err)
+	}
+
+	if !bytes.HasPrefix(conn.written, []byte{0, 1, 0, 3}) {
+		t.Errorf("unexpected written data: %v", conn.written)
+	}
+}
+
+func TestAuthFastSHA256PasswordRSA(t *testing.T) {
+	conn, mc := newRWMockConn(1)
+	mc.cfg.User = "root"
+	mc.cfg.Passwd = "secret"
+
+	authData := []byte{6, 81, 96, 114, 14, 42, 50, 30, 76, 47, 1, 95, 126, 81,
+		62, 94, 83, 80, 52, 85}
+	plugin := "sha256_password"
+
+	// Send Client Authentication Packet
+	authResp, addNUL, err := mc.auth(authData, plugin)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = mc.writeHandshakeResponsePacket(authResp, addNUL, plugin)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// check written auth response
+	authRespStart := 4 + 4 + 4 + 1 + 23 + len(mc.cfg.User) + 1
+	authRespEnd := authRespStart + 1 + len(authResp)
+	writtenAuthRespLen := conn.written[authRespStart]
+	writtenAuthResp := conn.written[authRespStart+1 : authRespEnd]
+	expectedAuthResp := []byte{1}
+	if writtenAuthRespLen != 1 || !bytes.Equal(writtenAuthResp, expectedAuthResp) {
+		t.Fatalf("unexpected written auth response (%d bytes): %v", writtenAuthRespLen, writtenAuthResp)
+	}
+	conn.written = nil
+
+	// auth response (pub key response)
+	conn.data = append([]byte{byte(1 + len(testPubKey)), 1, 0, 2, 1}, testPubKey...)
+	conn.queuedReplies = [][]byte{
+		// OK
+		{7, 0, 0, 4, 0, 0, 0, 2, 0, 0, 0},
+	}
+	conn.maxReads = 2
+
+	// Handle response to auth packet
+	if err := mc.handleAuthResult(authData, plugin); err != nil {
+		t.Errorf("got error: %v", err)
+	}
+
+	if !bytes.HasPrefix(conn.written, []byte{0, 1, 0, 3}) {
+		t.Errorf("unexpected written data: %v", conn.written)
+	}
+}
+
+func TestAuthFastSHA256PasswordRSAWithKey(t *testing.T) {
+	conn, mc := newRWMockConn(1)
+	mc.cfg.User = "root"
+	mc.cfg.Passwd = "secret"
+	mc.cfg.pubKey = testPubKeyRSA
+
+	authData := []byte{6, 81, 96, 114, 14, 42, 50, 30, 76, 47, 1, 95, 126, 81,
+		62, 94, 83, 80, 52, 85}
+	plugin := "sha256_password"
+
+	// Send Client Authentication Packet
+	authResp, addNUL, err := mc.auth(authData, plugin)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = mc.writeHandshakeResponsePacket(authResp, addNUL, plugin)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// auth response (OK)
+	conn.data = []byte{7, 0, 0, 2, 0, 0, 0, 2, 0, 0, 0}
+	conn.maxReads = 1
+
+	// Handle response to auth packet
+	if err := mc.handleAuthResult(authData, plugin); err != nil {
+		t.Errorf("got error: %v", err)
+	}
+}
+
+func TestAuthFastSHA256PasswordSecure(t *testing.T) {
+	conn, mc := newRWMockConn(1)
+	mc.cfg.User = "root"
+	mc.cfg.Passwd = "secret"
+
+	// hack to make the caching_sha2_password plugin believe that the connection
+	// is secure
+	mc.cfg.tls = &tls.Config{InsecureSkipVerify: true}
+
+	authData := []byte{6, 81, 96, 114, 14, 42, 50, 30, 76, 47, 1, 95, 126, 81,
+		62, 94, 83, 80, 52, 85}
+	plugin := "sha256_password"
+
+	// send Client Authentication Packet
+	authResp, addNUL, err := mc.auth(authData, plugin)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// unset TLS config to prevent the actual establishment of a TLS wrapper
+	mc.cfg.tls = nil
+
+	err = mc.writeHandshakeResponsePacket(authResp, addNUL, plugin)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// check written auth response
+	authRespStart := 4 + 4 + 4 + 1 + 23 + len(mc.cfg.User) + 1
+	authRespEnd := authRespStart + 1 + len(authResp) + 1
+	writtenAuthRespLen := conn.written[authRespStart]
+	writtenAuthResp := conn.written[authRespStart+1 : authRespEnd]
+	expectedAuthResp := []byte{115, 101, 99, 114, 101, 116, 0}
+	if writtenAuthRespLen != 6 || !bytes.Equal(writtenAuthResp, expectedAuthResp) {
+		t.Fatalf("unexpected written auth response (%d bytes): %v", writtenAuthRespLen, writtenAuthResp)
+	}
+	conn.written = nil
+
+	// auth response (OK)
+	conn.data = []byte{7, 0, 0, 2, 0, 0, 0, 2, 0, 0, 0}
+	conn.maxReads = 1
+
+	// Handle response to auth packet
+	if err := mc.handleAuthResult(authData, plugin); err != nil {
+		t.Errorf("got error: %v", err)
+	}
+
+	if !bytes.Equal(conn.written, []byte{}) {
+		t.Errorf("unexpected written data: %v", conn.written)
+	}
+}
+
 func TestAuthSwitchCachingSHA256PasswordCached(t *testing.T) {
 	conn, mc := newRWMockConn(2)
 	mc.cfg.Passwd = "secret"
@@ -559,7 +785,7 @@ func TestAuthSwitchCachingSHA256PasswordFullRSA(t *testing.T) {
 		{2, 0, 0, 4, 1, 4},
 
 		// Pub Key Response
-		append([]byte{byte(len(serverPubKey)), 1, 0, 6}, serverPubKey...),
+		append([]byte{byte(1 + len(testPubKey)), 1, 0, 6, 1}, testPubKey...),
 
 		// OK
 		{7, 0, 0, 8, 0, 0, 0, 2, 0, 0, 0},
@@ -591,6 +817,48 @@ func TestAuthSwitchCachingSHA256PasswordFullRSA(t *testing.T) {
 	}
 }
 
+func TestAuthSwitchCachingSHA256PasswordFullRSAWithKey(t *testing.T) {
+	conn, mc := newRWMockConn(2)
+	mc.cfg.Passwd = "secret"
+	mc.cfg.pubKey = testPubKeyRSA
+
+	// auth switch request
+	conn.data = []byte{44, 0, 0, 2, 254, 99, 97, 99, 104, 105, 110, 103, 95,
+		115, 104, 97, 50, 95, 112, 97, 115, 115, 119, 111, 114, 100, 0, 101,
+		11, 26, 18, 94, 97, 22, 72, 2, 46, 70, 106, 29, 55, 45, 94, 76, 90, 84,
+		50, 0}
+
+	conn.queuedReplies = [][]byte{
+		// Perform Full Authentication
+		{2, 0, 0, 4, 1, 4},
+
+		// OK
+		{7, 0, 0, 6, 0, 0, 0, 2, 0, 0, 0},
+	}
+	conn.maxReads = 3
+
+	authData := []byte{123, 87, 15, 84, 20, 58, 37, 121, 91, 117, 51, 24, 19,
+		47, 43, 9, 41, 112, 67, 110}
+	plugin := "mysql_native_password"
+
+	if err := mc.handleAuthResult(authData, plugin); err != nil {
+		t.Errorf("got error: %v", err)
+	}
+
+	expectedReplyPrefix := []byte{
+		// 1. Packet: Hash
+		32, 0, 0, 3, 129, 93, 132, 95, 114, 48, 79, 215, 128, 62, 193, 118, 128,
+		54, 75, 208, 159, 252, 227, 215, 129, 15, 242, 97, 19, 159, 31, 20, 58,
+		153, 9, 130,
+
+		// 2. Packet: Encrypted Password
+		0, 1, 0, 5, // [changing bytes]
+	}
+	if !bytes.HasPrefix(conn.written, expectedReplyPrefix) {
+		t.Errorf("got unexpected data: %v", conn.written)
+	}
+}
+
 func TestAuthSwitchCachingSHA256PasswordFullSecure(t *testing.T) {
 	conn, mc := newRWMockConn(2)
 	mc.cfg.Passwd = "secret"
@@ -851,3 +1119,142 @@ func TestAuthSwitchOldPasswordEmpty(t *testing.T) {
 		t.Errorf("got unexpected data: %v", conn.written)
 	}
 }
+
+func TestAuthSwitchSHA256PasswordEmpty(t *testing.T) {
+	conn, mc := newRWMockConn(2)
+	mc.cfg.Passwd = ""
+
+	// auth switch request
+	conn.data = []byte{38, 0, 0, 2, 254, 115, 104, 97, 50, 53, 54, 95, 112, 97,
+		115, 115, 119, 111, 114, 100, 0, 78, 82, 62, 40, 100, 1, 59, 31, 44, 69,
+		33, 112, 8, 81, 51, 96, 65, 82, 16, 114, 0}
+
+	conn.queuedReplies = [][]byte{
+		// OK
+		{7, 0, 0, 4, 0, 0, 0, 2, 0, 0, 0},
+	}
+	conn.maxReads = 3
+
+	authData := []byte{123, 87, 15, 84, 20, 58, 37, 121, 91, 117, 51, 24, 19,
+		47, 43, 9, 41, 112, 67, 110}
+	plugin := "mysql_native_password"
+
+	if err := mc.handleAuthResult(authData, plugin); err != nil {
+		t.Errorf("got error: %v", err)
+	}
+
+	expectedReplyPrefix := []byte{
+		// 1. Packet: Empty Password
+		1, 0, 0, 3, 0,
+	}
+	if !bytes.HasPrefix(conn.written, expectedReplyPrefix) {
+		t.Errorf("got unexpected data: %v", conn.written)
+	}
+}
+
+func TestAuthSwitchSHA256PasswordRSA(t *testing.T) {
+	conn, mc := newRWMockConn(2)
+	mc.cfg.Passwd = "secret"
+
+	// auth switch request
+	conn.data = []byte{38, 0, 0, 2, 254, 115, 104, 97, 50, 53, 54, 95, 112, 97,
+		115, 115, 119, 111, 114, 100, 0, 78, 82, 62, 40, 100, 1, 59, 31, 44, 69,
+		33, 112, 8, 81, 51, 96, 65, 82, 16, 114, 0}
+
+	conn.queuedReplies = [][]byte{
+		// Pub Key Response
+		append([]byte{byte(1 + len(testPubKey)), 1, 0, 4, 1}, testPubKey...),
+
+		// OK
+		{7, 0, 0, 6, 0, 0, 0, 2, 0, 0, 0},
+	}
+	conn.maxReads = 3
+
+	authData := []byte{123, 87, 15, 84, 20, 58, 37, 121, 91, 117, 51, 24, 19,
+		47, 43, 9, 41, 112, 67, 110}
+	plugin := "mysql_native_password"
+
+	if err := mc.handleAuthResult(authData, plugin); err != nil {
+		t.Errorf("got error: %v", err)
+	}
+
+	expectedReplyPrefix := []byte{
+		// 1. Packet: Pub Key Request
+		1, 0, 0, 3, 1,
+
+		// 2. Packet: Encrypted Password
+		0, 1, 0, 5, // [changing bytes]
+	}
+	if !bytes.HasPrefix(conn.written, expectedReplyPrefix) {
+		t.Errorf("got unexpected data: %v", conn.written)
+	}
+}
+
+func TestAuthSwitchSHA256PasswordRSAWithKey(t *testing.T) {
+	conn, mc := newRWMockConn(2)
+	mc.cfg.Passwd = "secret"
+	mc.cfg.pubKey = testPubKeyRSA
+
+	// auth switch request
+	conn.data = []byte{38, 0, 0, 2, 254, 115, 104, 97, 50, 53, 54, 95, 112, 97,
+		115, 115, 119, 111, 114, 100, 0, 78, 82, 62, 40, 100, 1, 59, 31, 44, 69,
+		33, 112, 8, 81, 51, 96, 65, 82, 16, 114, 0}
+
+	conn.queuedReplies = [][]byte{
+		// OK
+		{7, 0, 0, 4, 0, 0, 0, 2, 0, 0, 0},
+	}
+	conn.maxReads = 2
+
+	authData := []byte{123, 87, 15, 84, 20, 58, 37, 121, 91, 117, 51, 24, 19,
+		47, 43, 9, 41, 112, 67, 110}
+	plugin := "mysql_native_password"
+
+	if err := mc.handleAuthResult(authData, plugin); err != nil {
+		t.Errorf("got error: %v", err)
+	}
+
+	expectedReplyPrefix := []byte{
+		// 1. Packet: Encrypted Password
+		0, 1, 0, 3, // [changing bytes]
+	}
+	if !bytes.HasPrefix(conn.written, expectedReplyPrefix) {
+		t.Errorf("got unexpected data: %v", conn.written)
+	}
+}
+
+func TestAuthSwitchSHA256PasswordSecure(t *testing.T) {
+	conn, mc := newRWMockConn(2)
+	mc.cfg.Passwd = "secret"
+
+	// Hack to make the caching_sha2_password plugin believe that the connection
+	// is secure
+	mc.cfg.tls = &tls.Config{InsecureSkipVerify: true}
+
+	// auth switch request
+	conn.data = []byte{38, 0, 0, 2, 254, 115, 104, 97, 50, 53, 54, 95, 112, 97,
+		115, 115, 119, 111, 114, 100, 0, 78, 82, 62, 40, 100, 1, 59, 31, 44, 69,
+		33, 112, 8, 81, 51, 96, 65, 82, 16, 114, 0}
+
+	conn.queuedReplies = [][]byte{
+		// OK
+		{7, 0, 0, 4, 0, 0, 0, 2, 0, 0, 0},
+	}
+	conn.maxReads = 2
+
+	authData := []byte{123, 87, 15, 84, 20, 58, 37, 121, 91, 117, 51, 24, 19,
+		47, 43, 9, 41, 112, 67, 110}
+	plugin := "mysql_native_password"
+
+	if err := mc.handleAuthResult(authData, plugin); err != nil {
+		t.Errorf("got error: %v", err)
+	}
+
+	expectedReplyPrefix := []byte{
+		// 1. Packet: Cleartext Password
+		7, 0, 0, 3, 115, 101, 99, 114, 101, 116, 0,
+	}
+	if !bytes.Equal(conn.written, expectedReplyPrefix) {
+		t.Errorf("got unexpected data: %v", conn.written)
+	}
+}

dsn.go

@@ -10,6 +10,7 @@ package mysql
 
 import (
 	"bytes"
+	"crypto/rsa"
 	"crypto/tls"
 	"errors"
 	"fmt"
@@ -41,6 +42,8 @@ type Config struct {
 	Collation        string            // Connection collation
 	Loc              *time.Location    // Location for time.Time values
 	MaxAllowedPacket int               // Max packet size allowed
+	ServerPubKey     string            // Server public key name
+	pubKey           *rsa.PublicKey    // Server public key
 	TLSConfig        string            // TLS configuration name
 	tls              *tls.Config       // TLS configuration
 	Timeout          time.Duration     // Dial timeout
@@ -254,6 +257,16 @@ func (cfg *Config) FormatDSN() string {
 		}
 	}
 
+	if len(cfg.ServerPubKey) > 0 {
+		if hasParam {
+			buf.WriteString("&serverPubKey=")
+		} else {
+			hasParam = true
+			buf.WriteString("?serverPubKey=")
+		}
+		buf.WriteString(url.QueryEscape(cfg.ServerPubKey))
+	}
+
 	if cfg.Timeout > 0 {
 		if hasParam {
 			buf.WriteString("&timeout=")
@@ -512,6 +525,20 @@ func parseDSNParams(cfg *Config, params string) (err error) {
 				return errors.New("invalid bool value: " + value)
 			}
 
+		// Server public key
+		case "serverPubKey":
+			name, err := url.QueryUnescape(value)
+			if err != nil {
+				return fmt.Errorf("invalid value for server pub key name: %v", err)
+			}
+
+			if pubKey := getServerPubKey(name); pubKey != nil {
+				cfg.ServerPubKey = name
+				cfg.pubKey = pubKey
+			} else {
+				return errors.New("invalid value / unknown server pub key name: " + name)
+			}
+
 		// Strict mode
 		case "strict":
 			panic("strict mode has been removed. See https://github.com/go-sql-driver/mysql/wiki/strict-mode")

dsn_test.go

@@ -135,11 +135,56 @@ func TestDSNReformat(t *testing.T) {
 	}
 }
 
+func TestDSNServerPubKey(t *testing.T) {
+	baseDSN := "User:password@tcp(localhost:5555)/dbname?serverPubKey="
+
+	RegisterServerPubKey("testKey", testPubKeyRSA)
+	defer DeregisterServerPubKey("testKey")
+
+	tst := baseDSN + "testKey"
+	cfg, err := ParseDSN(tst)
+	if err != nil {
+		t.Error(err.Error())
+	}
+
+	if cfg.ServerPubKey != "testKey" {
+		t.Errorf("unexpected cfg.ServerPubKey value: %v", cfg.ServerPubKey)
+	}
+	if cfg.pubKey != testPubKeyRSA {
+		t.Error("pub key pointer doesn't match")
+	}
+
+	// Key is missing
+	tst = baseDSN + "invalid_name"
+	cfg, err = ParseDSN(tst)
+	if err == nil {
+		t.Errorf("invalid name in DSN (%s) but did not error. Got config: %#v", tst, cfg)
+	}
+}
+
+func TestDSNServerPubKeyQueryEscape(t *testing.T) {
+	const name = "&%!:"
+	dsn := "User:password@tcp(localhost:5555)/dbname?serverPubKey=" + url.QueryEscape(name)
+
+	RegisterServerPubKey(name, testPubKeyRSA)
+	defer DeregisterServerPubKey(name)
+
+	cfg, err := ParseDSN(dsn)
+	if err != nil {
+		t.Error(err.Error())
+	}
+
+	if cfg.pubKey != testPubKeyRSA {
+		t.Error("pub key pointer doesn't match")
+	}
+}
+
 func TestDSNWithCustomTLS(t *testing.T) {
 	baseDSN := "User:password@tcp(localhost:5555)/dbname?tls="
 	tlsCfg := tls.Config{}
 
 	RegisterTLSConfig("utils_test", &tlsCfg)
+	defer DeregisterTLSConfig("utils_test")
 
 	// Custom TLS is missing
 	tst := baseDSN + "invalid_tls"
@@ -173,8 +218,6 @@ func TestDSNWithCustomTLS(t *testing.T) {
 	} else if tlsCfg.ServerName != "" {
 		t.Errorf("tlsCfg was mutated ServerName (%s) should be empty parsing DSN (%s).", name, tst)
 	}
-
-	DeregisterTLSConfig("utils_test")
 }
 
 func TestDSNTLSConfig(t *testing.T) {
@@ -212,6 +255,7 @@ func TestDSNWithCustomTLSQueryEscape(t *testing.T) {
 	tlsCfg := tls.Config{ServerName: name}
 
 	RegisterTLSConfig(configKey, &tlsCfg)
+	defer DeregisterTLSConfig(configKey)
 
 	cfg, err := ParseDSN(dsn)
 

packets.go

@@ -270,7 +270,16 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, addNUL bool,
 		clientFlags |= clientMultiStatements
 	}
 
-	pktLen := 4 + 4 + 1 + 23 + len(mc.cfg.User) + 1 + 1 + len(authResp) + 21 + 1
+	// encode length of the auth plugin data
+	var authRespLEIBuf [9]byte
+	authRespLEI := appendLengthEncodedInteger(authRespLEIBuf[:0], uint64(len(authResp)))
+	if len(authRespLEI) > 1 {
+		// if the length can not be written in 1 byte, it must be written as a
+		// length encoded integer
+		clientFlags |= clientPluginAuthLenEncClientData
+	}
+
+	pktLen := 4 + 4 + 1 + 23 + len(mc.cfg.User) + 1 + len(authRespLEI) + len(authResp) + 21 + 1
 	if addNUL {
 		pktLen++
 	}
@@ -342,8 +351,8 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, addNUL bool,
 	pos++
 
 	// Auth Data [length encoded integer]
-	data[pos] = byte(len(authResp))
-	pos += 1 + copy(data[pos+1:], authResp)
+	pos += copy(data[pos:], authRespLEI)
+	pos += copy(data[pos:], authResp)
 	if addNUL {
 		data[pos] = 0x00
 		pos++

utils.go

@@ -20,15 +20,17 @@ import (
 	"time"
 )
 
+// Registry for custom tls.Configs
 var (
 	tlsConfigLock     sync.RWMutex
-	tlsConfigRegister map[string]*tls.Config // Register for custom tls.Configs
+	tlsConfigRegistry map[string]*tls.Config
 )
 
 // RegisterTLSConfig registers a custom tls.Config to be used with sql.Open.
 // Use the key as a value in the DSN where tls=value.
 //
-// Note: The tls.Config provided to needs to be exclusively owned by the driver after registering.
+// Note: The provided tls.Config is exclusively owned by the driver after
+// registering it.
 //
 //  rootCertPool := x509.NewCertPool()
 //  pem, err := ioutil.ReadFile("/path/ca-cert.pem")
@@ -56,11 +58,11 @@ func RegisterTLSConfig(key string, config *tls.Config) error {
 	}
 
 	tlsConfigLock.Lock()
-	if tlsConfigRegister == nil {
-		tlsConfigRegister = make(map[string]*tls.Config)
+	if tlsConfigRegistry == nil {
+		tlsConfigRegistry = make(map[string]*tls.Config)
 	}
 
-	tlsConfigRegister[key] = config
+	tlsConfigRegistry[key] = config
 	tlsConfigLock.Unlock()
 	return nil
 }
@@ -68,15 +70,15 @@ func RegisterTLSConfig(key string, config *tls.Config) error {
 // DeregisterTLSConfig removes the tls.Config associated with key.
 func DeregisterTLSConfig(key string) {
 	tlsConfigLock.Lock()
-	if tlsConfigRegister != nil {
-		delete(tlsConfigRegister, key)
+	if tlsConfigRegistry != nil {
+		delete(tlsConfigRegistry, key)
 	}
 	tlsConfigLock.Unlock()
 }
 
 func getTLSConfigClone(key string) (config *tls.Config) {
 	tlsConfigLock.RLock()
-	if v, ok := tlsConfigRegister[key]; ok {
+	if v, ok := tlsConfigRegistry[key]; ok {
 		config = cloneTLSConfig(v)
 	}
 	tlsConfigLock.RUnlock()