From 5f31fe11e73075e60e32236663ccce19e60cf3b6 Mon Sep 17 00:00:00 2001
From: zhaomoran <imrzhao@qq.com>
Date: Sat, 31 Jul 2021 01:53:20 +0800
Subject: [PATCH 1/3] Support SaslHandshakeRequest v0

---
 broker.go | 66 ++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 65 insertions(+), 1 deletion(-)

diff --git a/broker.go b/broker.go
index dd01e4ef1..32c61d67a 100644
--- a/broker.go
+++ b/broker.go
@@ -941,7 +941,7 @@ func (b *Broker) authenticateViaSASL() error {
 	case SASLTypeOAuth:
 		return b.sendAndReceiveSASLOAuth(b.conf.Net.SASL.TokenProvider)
 	case SASLTypeSCRAMSHA256, SASLTypeSCRAMSHA512:
-		return b.sendAndReceiveSASLSCRAMv1()
+		return b.sendAndReceiveSASLSCRAM()
 	case SASLTypeGSSAPI:
 		return b.sendAndReceiveKerberos()
 	default:
@@ -1180,6 +1180,70 @@ func (b *Broker) sendClientMessage(message []byte) (bool, error) {
 	return isChallenge, err
 }
 
+func (b *Broker) sendAndReceiveSASLSCRAM() error {
+	if b.conf.Net.SASL.Version == SASLHandshakeV0 {
+		return b.sendAndReceiveSASLSCRAMv0()
+	}
+	return b.sendAndReceiveSASLSCRAMv1()
+}
+
+func (b *Broker) sendAndReceiveSASLSCRAMv0() error {
+    if err := b.sendAndReceiveSASLHandshake(b.conf.Net.SASL.Mechanism, SASLHandshakeV0); err != nil {
+    	return err
+    }
+
+	scramClient := b.conf.Net.SASL.SCRAMClientGeneratorFunc()
+	if err := scramClient.Begin(b.conf.Net.SASL.User, b.conf.Net.SASL.Password, b.conf.Net.SASL.SCRAMAuthzID); err != nil {
+		return fmt.Errorf("failed to start SCRAM exchange with the server: %s", err.Error())
+	}
+
+	msg, err := scramClient.Step("")
+	if err != nil {
+		return fmt.Errorf("failed to advance the SCRAM exchange: %s", err.Error())
+	}
+
+	for !scramClient.Done() {
+		requestTime := time.Now()
+		// Will be decremented in updateIncomingCommunicationMetrics (except error)
+		b.addRequestInFlightMetrics(1)
+		length := len(msg)
+		authBytes := make([]byte, length+4) //4 byte length header + auth data
+		binary.BigEndian.PutUint32(authBytes, uint32(length))
+		copy(authBytes[4:], []byte(msg))
+		_, err := b.write(authBytes)
+		b.updateOutgoingCommunicationMetrics(length + 4)
+		if err != nil {
+			b.addRequestInFlightMetrics(-1)
+			Logger.Printf("Failed to write SASL auth header to broker %s: %s\n", b.addr, err.Error())
+			return err
+		}
+		b.correlationID++
+		header := make([]byte, 4)
+		_, err = b.readFull(header)
+		if err != nil {
+			b.addRequestInFlightMetrics(-1)
+			Logger.Printf("Failed to read response header while authenticating with SASL to broker %s: %s\n", b.addr, err.Error())
+			return err
+		}
+		payload := make([]byte, int32(binary.BigEndian.Uint32(header)))
+		n, err := b.readFull(payload)
+		if err != nil {
+			b.addRequestInFlightMetrics(-1)
+			Logger.Printf("Failed to read response payload while authenticating with SASL to broker %s: %s\n", b.addr, err.Error())
+			return err
+		}
+		b.updateIncomingCommunicationMetrics(n+4, time.Since(requestTime))
+		msg, err = scramClient.Step(string(payload))
+		if err != nil {
+			Logger.Println("SASL authentication failed", err)
+			return err
+		}
+	}
+
+	Logger.Println("SASL authentication succeeded")
+	return nil
+}
+
 func (b *Broker) sendAndReceiveSASLSCRAMv1() error {
 	if err := b.sendAndReceiveSASLHandshake(b.conf.Net.SASL.Mechanism, SASLHandshakeV1); err != nil {
 		return err

From d5de7bebc5b658cf4c2505332f04d103de928bfa Mon Sep 17 00:00:00 2001
From: mrzhao <imrzhao@qq.com>
Date: Tue, 10 Aug 2021 20:26:59 +0800
Subject: [PATCH 2/3] fix the TestSASLSCRAMSHAXXX for SASLHandshakeV1

---
 broker_test.go | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/broker_test.go b/broker_test.go
index 2fa40ceb4..a91e35254 100644
--- a/broker_test.go
+++ b/broker_test.go
@@ -359,9 +359,11 @@ func TestSASLSCRAMSHAXXX(t *testing.T) {
 
 			conf := NewTestConfig()
 			conf.Net.SASL.Mechanism = SASLTypeSCRAMSHA512
+			conf.Net.SASL.Version = SASLHandshakeV1
 			conf.Net.SASL.SCRAMClientGeneratorFunc = func() SCRAMClient { return test.scramClient }
 
 			broker.conf = conf
+			broker.conf.Version = V1_0_0_0
 			dialer := net.Dialer{
 				Timeout:   conf.Net.DialTimeout,
 				KeepAlive: conf.Net.KeepAlive,

From e7071f3672bf950c6df9019fecbdaf9e08b76db6 Mon Sep 17 00:00:00 2001
From: Vlad Gorodetsky <v@gor.io>
Date: Wed, 8 Sep 2021 12:53:05 +0300
Subject: [PATCH 3/3] Fix go fmt

---
 broker.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/broker.go b/broker.go
index 32c61d67a..d30840b1e 100644
--- a/broker.go
+++ b/broker.go
@@ -1188,9 +1188,9 @@ func (b *Broker) sendAndReceiveSASLSCRAM() error {
 }
 
 func (b *Broker) sendAndReceiveSASLSCRAMv0() error {
-    if err := b.sendAndReceiveSASLHandshake(b.conf.Net.SASL.Mechanism, SASLHandshakeV0); err != nil {
-    	return err
-    }
+	if err := b.sendAndReceiveSASLHandshake(b.conf.Net.SASL.Mechanism, SASLHandshakeV0); err != nil {
+		return err
+	}
 
 	scramClient := b.conf.Net.SASL.SCRAMClientGeneratorFunc()
 	if err := scramClient.Begin(b.conf.Net.SASL.User, b.conf.Net.SASL.Password, b.conf.Net.SASL.SCRAMAuthzID); err != nil {