v3.0.3-alpha.2

Author: YilunZhang39

api/pool.go

@@ -58,6 +58,9 @@ type PoolOption struct {
 
 	// try reconnect times
 	TryReconnectNums *int
+
+	// if enable SCRAM login verify
+	EnableScram bool
 }
 
 // NewDBConnectionPool inits a DBConnectionPool object and configures it with opt, finally returns it.
@@ -107,6 +110,7 @@ func newConn(addr string, opt *PoolOption) (dialer.Conn, error) {
 		HighAvailabilitySites:  opt.HighAvailabilitySites,
 		Reconnect:              opt.Reconnect,
 		TryReconnectNums:       opt.TryReconnectNums,
+		EnableScram:            opt.EnableScram,
 	}
 	conn, err := dialer.NewConn(context.TODO(), addr, bOpt)
 	if err != nil {

dialer/behavior.go

@@ -38,6 +38,9 @@ type BehaviorOptions struct {
 
 	// UsePython specifies whether the session uses a Python parser
 	UsePython bool
+
+	// if enable SCRAM login verify
+	EnableScram bool
 }
 
 // SetPriority sets the priority of the task.

dialer/dialer.go

@@ -69,7 +69,9 @@ type Conn interface {
 	// GetTCPConn returns the TCPConn
 	GetTCPConn() *net.TCPConn
 
+	// for inner use
 	GetReader() protocol.Reader
+	enableScram() bool
 }
 
 type conn struct {
@@ -289,32 +291,9 @@ func (c *conn) connect(addr string) error {
 	c.isConnected = true
 	c.isClosed = false
 	c.refreshHeaderForResponse(h)
-	if c.userID != "" {
-		args := make([]model.DataForm, 2)
-		user, err := model.NewDataType(model.DtString, c.userID)
-		if err != nil {
-			return err
-		}
-		pwd, err := model.NewDataType(model.DtString, c.password)
-		if err != nil {
-			return err
-		}
-		args[0] = model.NewScalar(user)
-		args[1] = model.NewScalar(pwd)
-
-		bo := defaultByteOrder
 
-		_, _, err = c.run(&requestParams{
-			commandType: functionCmd,
-			Command:     generateFunctionCommand("login", bo, args),
-			SessionID:   []byte(c.GetSession()),
-			Args:        args,
-			ByteOrder:   bo,
-		})
-
-		if err != nil {
-			return err
-		}
+	if c.userID != "" {
+		Login(c, c.userID, c.password)
 	}
 
 	return nil
@@ -425,6 +404,7 @@ func (c *conn) run(params *requestParams) (*responseHeader, model.DataForm, erro
 						continue
 					}
 				}
+				time.Sleep(300 * time.Millisecond)
 				err := c.switchDatanode(nil)
 				if err != nil {
 					return nil, nil, err
@@ -481,3 +461,7 @@ func (c *conn) runInternal(params *requestParams) (*responseHeader, model.DataFo
 func (c *conn) refreshHeaderForResponse(h *responseHeader) {
 	c.sessionID = h.sessionID
 }
+
+func (c *conn) enableScram() bool {
+	return c.behaviorOpt.EnableScram
+}

dialer/util.go

@@ -2,6 +2,11 @@ package dialer
 
 import (
 	"bytes"
+	"crypto/hmac"
+	"crypto/rand"
+	"crypto/sha256"
+	"encoding/base64"
+	"fmt"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -10,6 +15,7 @@ import (
 
 	"github.com/dolphindb/api-go/v3/dialer/protocol"
 	"github.com/dolphindb/api-go/v3/model"
+	"golang.org/x/crypto/pbkdf2"
 )
 
 const (
@@ -113,6 +119,14 @@ func parseAddr(raw string) string {
 }
 
 func Login(conn Conn, userID, password string) error {
+	if conn.enableScram() {
+		return scramLogin(conn, userID, password)
+	} else {
+		err := scramLogin(conn, userID, password)
+		if err == nil {
+			return nil
+		}
+	}
 	args := make([]model.DataForm, 2)
 	user, err := model.NewDataType(model.DtString, userID)
 	if err != nil {
@@ -131,3 +145,130 @@ func Login(conn Conn, userID, password string) error {
 	}
 	return nil
 }
+
+func generateNonce(length int) (string, error) {
+	buffer := make([]byte, length)
+	_, err := rand.Read(buffer)
+	if err != nil {
+		return "", err
+	}
+	return base64.StdEncoding.EncodeToString(buffer), nil
+}
+
+func xorBytes(a, b []byte) []byte {
+	result := make([]byte, len(a))
+	for i := range a {
+		result[i] = a[i] ^ b[i]
+	}
+	return result
+}
+
+func scramLogin(conn Conn, userID, password string) error {
+	args := make([]model.DataForm, 2)
+	user, err := model.NewDataType(model.DtString, userID)
+	if err != nil {
+		return fmt.Errorf("SCRAM login failed, %w", err)
+	}
+	clientNonce, err := generateNonce(16)
+	if err != nil {
+		return fmt.Errorf("SCRAM login failed, %w", err)
+	}
+	nonce, err := model.NewDataType(model.DtString, clientNonce)
+	if err != nil {
+		return fmt.Errorf("SCRAM login failed, %w", err)
+	}
+	args[0] = model.NewScalar(user)
+	args[1] = model.NewScalar(nonce)
+
+	result, err := conn.RunFunc("scramClientFirst", args)
+	if err != nil {
+		if strings.Contains(err.Error(), "Can't recognize function name scramClientFirst") {
+			return fmt.Errorf("SCRAM login is unavailable on current server")
+		}
+		if strings.Contains(err.Error(), "sha256 authMode doesn't support scram authMode") {
+			return fmt.Errorf("user '%s' doesn't support scram authMode", userID)
+		}
+		return fmt.Errorf("scramClientFirst failed: %w", err)
+	}
+
+	retVec := result.(*model.Vector)
+
+	if retVec.Rows() != 3 {
+		return fmt.Errorf("SCRAM login failed, server error: get server nonce failed")
+	}
+	saltStr := retVec.Get(0).Value().(*model.Scalar).Value().(string)
+	iterCount := int(retVec.Get(1).Value().(*model.Scalar).Value().(int32))
+	combinedNonce := retVec.Get(2).Value().(*model.Scalar).Value().(string)
+
+	salt, err := base64.StdEncoding.DecodeString(saltStr)
+	if err != nil {
+		return fmt.Errorf("SCRAM login failed, base64 decode failed: %w", err)
+	}
+
+	saltedPassword := pbkdf2.Key([]byte(password), salt, iterCount, 32, sha256.New)
+
+	mac := hmac.New(sha256.New, saltedPassword)
+	_, err = mac.Write([]byte("Client Key"))
+	if err != nil {
+		return fmt.Errorf("SCRAM login failed, HMAC calculation failed: %w", err)
+	}
+	clientKey := mac.Sum(nil)
+
+	storedKey := sha256.Sum256(clientKey)
+
+	authMessage := fmt.Sprintf(`n=%s,r=%s,r=%s,s=%s,i=%d,c=biws,r=%s`,
+		userID, clientNonce, combinedNonce, saltStr, iterCount, combinedNonce)
+
+	mac = hmac.New(sha256.New, storedKey[:])
+	_, err = mac.Write([]byte(authMessage))
+	if err != nil {
+		return fmt.Errorf("SCRAM login failed, HMAC calculation failed: %w", err)
+	}
+	clientSig := mac.Sum(nil)
+
+	proof := xorBytes(clientKey, clientSig)
+
+	finalArgs := make([]model.DataForm, 3)
+	combinedNonceScalar, err := model.NewDataType(model.DtString, combinedNonce)
+	if err != nil {
+		return fmt.Errorf("SCRAM login failed, %w", err)
+	}
+	proofScalar, err := model.NewDataType(model.DtString, base64.StdEncoding.EncodeToString(proof))
+	if err != nil {
+		return fmt.Errorf("SCRAM login failed, %w", err)
+	}
+
+	finalArgs[0] = model.NewScalar(user)
+	finalArgs[1] = model.NewScalar(combinedNonceScalar)
+	finalArgs[2] = model.NewScalar(proofScalar)
+
+	finalResult, err := conn.RunFunc("scramClientFinal", finalArgs)
+	if err != nil {
+		return fmt.Errorf("scramClientFinal failed: %w", err)
+	}
+	serverSigBase64 := finalResult.(*model.Scalar).Value().(string)
+
+	mac = hmac.New(sha256.New, saltedPassword)
+	_, err = mac.Write([]byte("Server Key"))
+	if err != nil {
+		return fmt.Errorf("SCRAM login failed, HMAC calculation failed: %w", err)
+	}
+	serverKey := mac.Sum(nil)
+
+	mac = hmac.New(sha256.New, serverKey)
+	_, err = mac.Write([]byte(authMessage))
+	if err != nil {
+		return fmt.Errorf("SCRAM login failed, HMAC calculation failed: %w", err)
+	}
+	serverSig := mac.Sum(nil)
+
+	expectedSig := base64.StdEncoding.EncodeToString(serverSig)
+
+	if serverSigBase64 != "" && expectedSig != serverSigBase64 {
+		conn.Close()
+		return fmt.Errorf("invalid SCRAM server signature")
+	}
+
+	fmt.Println("SCRAM login succeeded")
+	return nil
+}

go.mod

@@ -9,8 +9,10 @@ require (
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/forgoer/openssl v1.6.0
 	github.com/kr/pretty v0.3.0 // indirect
 	github.com/shopspring/decimal v1.3.1
 	github.com/smartystreets/goconvey v1.7.2
+	golang.org/x/crypto v0.0.0-20220817201139-bc19a97f63c8
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 )

streaming/abstract_client.go

@@ -6,8 +6,8 @@ import (
 
 // AbstractClient is the client interface for streaming subscription.
 type AbstractClient interface {
-	activeCloseConnection(si *site) error
-	doReconnect(si *site) bool
+	activeCloseConnection(req *SubscribeRequest) error
+	doReconnect(req *SubscribeRequest) bool
 	getSubscriber() *subscriber
 	getConn() (net.Conn, bool)