Repository: snowflakedb/gosnowflake
Pull request #705 (2023-01-23 17:35:20)

Fix connection caching for MFA and external browser authenticators
### Description
https://github.com/snowflakedb/snowflake-sdks-drivers-issues-teamwork/issues/119

Connection caching was added in version 1.6.15 as part of #661 but is not working as expected. 

The PR added CLIENT_REQUEST_MFA_TOKEN/CLIENT_STORE_TEMPORARY_CREDENTIAL properties in the authentication request so the server would send back MFA/ID token. The implementation was sending the session parameter value as string "true" instead of boolean true.

This PR is to change the type of CLIENT_REQUEST_MFA_TOKEN/CLIENT_STORE_TEMPORARY_CREDENTIAL from string to boolean so that it is recognized by the server, and returns the token so it can be cached.

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [ ] Created tests which fail without the change (if possible)
- [x] All tests passing
- [ ] Extended the README / documentation, if necessary


sfc-gh-ext-simba-lb: @sfc-gh-igarish the go driver supports writing to local cache file the same as ODBC and JDBC: https://github.com/snowflakedb/gosnowflake/blob/master/secure_storage_manager.go#L239

For Linux it's disabled by default the same as JDBC and ODBC. We use the CLIENT_STORE_TEMPORARY_CREDENTIAL and CLIENT_REQUEST_MFA_TOKEN session parameters to decide whether to cache the token or not. For Windows and Mac OS these properties will be set to true by default if the authenticator is external browser or MFA, but for Linux they are disabled by default. 

i.e JDBC: https://github.com/snowflakedb/snowflake-jdbc/blob/master/src/main/java/net/snowflake/client/core/SessionUtil.java#L261
sfc-gh-igarish: Can we make sure this feature work on all supported platforms? i.e. Windows, Linux, and Mac.
sfc-gh-ext-simba-lb: @sfc-gh-igarish Added the Linux option to the test and confirmed the feature on Windows, Linux and OSX.

Patch:
diff --git a/auth.go b/auth.go
index 36679f3..71bfe4b 100644
--- a/auth.go
+++ b/auth.go
@@ -301,6 +301,12 @@ func authenticate(
 	}
 
 	sessionParameters[sessionClientValidateDefaultParameters] = sc.cfg.ValidateDefaultParameters != ConfigBoolFalse
+	if sc.cfg.ClientRequestMfaToken == ConfigBoolTrue {
+		sessionParameters[clientRequestMfaToken] = true
+	}
+	if sc.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {
+		sessionParameters[clientStoreTemporaryCredential] = true
+	}
 	requestMain := authRequestData{
 		ClientAppID:       clientType,
 		ClientAppVersion:  SnowflakeGoDriverVersion,
@@ -314,6 +320,7 @@ func authenticate(
 		if sc.cfg.IDToken != "" {
 			requestMain.Authenticator = idTokenAuthenticator
 			requestMain.Token = sc.cfg.IDToken
+			requestMain.LoginName = sc.cfg.User
 		} else {
 			requestMain.ProofKey = string(proofKey)
 			requestMain.Token = string(samlResponse)
@@ -401,10 +408,10 @@ func authenticate(
 	if !respd.Success {
 		logger.Errorln("Authentication FAILED")
 		sc.rest.TokenAccessor.SetTokens("", "", -1)
-		if sessionParameters[clientRequestMfaToken] == "true" {
+		if sessionParameters[clientRequestMfaToken] == true {
 			deleteCredential(sc, mfaToken)
 		}
-		if sessionParameters[clientStoreTemporaryCredential] == "true" {
+		if sessionParameters[clientStoreTemporaryCredential] == true {
 			deleteCredential(sc, idToken)
 		}
 		code, err := strconv.Atoi(respd.Code)
@@ -420,11 +427,13 @@ func authenticate(
 	}
 	logger.Info("Authentication SUCCESS")
 	sc.rest.TokenAccessor.SetTokens(respd.Data.Token, respd.Data.MasterToken, respd.Data.SessionID)
-	if sc.isClientRequestMfaToken() {
-		setCredential(sc, mfaToken, respd.Data.MfaToken)
+	if sessionParameters[clientRequestMfaToken] == true {
+		token := respd.Data.MfaToken
+		setCredential(sc, mfaToken, token)
 	}
-	if sc.isClientStoreTemporaryCredential() {
-		setCredential(sc, idToken, respd.Data.IDToken)
+	if sessionParameters[clientStoreTemporaryCredential] == true {
+		token := respd.Data.IDToken
+		setCredential(sc, idToken, token)
 	}
 	return &respd.Data, nil
 }
@@ -466,21 +475,20 @@ func authenticateWithConfig(sc *snowflakeConn) error {
 	var err error
 	//var consentCacheIdToken = true
 
-	paramBoolValue := "true"
 	if sc.cfg.Authenticator == AuthTypeExternalBrowser {
 		if runtime.GOOS == "windows" || runtime.GOOS == "darwin" {
-			sc.cfg.Params[clientStoreTemporaryCredential] = &paramBoolValue
+			sc.cfg.ClientStoreTemporaryCredential = ConfigBoolTrue
 		}
-		if sc.isClientStoreTemporaryCredential() {
+		if sc.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {
 			fillCachedIDToken(sc)
 		}
 	}
 
 	if sc.cfg.Authenticator == AuthTypeUsernamePasswordMFA {
 		if runtime.GOOS == "windows" || runtime.GOOS == "darwin" {
-			sc.cfg.Params[clientRequestMfaToken] = &paramBoolValue
+			sc.cfg.ClientRequestMfaToken = ConfigBoolTrue
 		}
-		if sc.isClientRequestMfaToken() {
+		if sc.cfg.ClientRequestMfaToken == ConfigBoolTrue {
 			fillCachedMfaToken(sc)
 		}
 	}
diff --git a/connection_util.go b/connection_util.go
index c963ed3..10801c8 100644
--- a/connection_util.go
+++ b/connection_util.go
@@ -21,22 +21,6 @@ func (sc *snowflakeConn) isClientSessionKeepAliveEnabled() bool {
 	return strings.Compare(*v, "true") == 0
 }
 
-func (sc *snowflakeConn) isClientStoreTemporaryCredential() bool {
-	v, ok := sc.cfg.Params[clientStoreTemporaryCredential]
-	if !ok {
-		return false
-	}
-	return strings.Compare(*v, "true") == 0
-}
-
-func (sc *snowflakeConn) isClientRequestMfaToken() bool {
-	v, ok := sc.cfg.Params[clientRequestMfaToken]
-	if !ok {
-		return false
-	}
-	return strings.Compare(*v, "true") == 0
-}
-
 func (sc *snowflakeConn) startHeartBeat() {
 	if !sc.isClientSessionKeepAliveEnabled() {
 		return
diff --git a/dsn.go b/dsn.go
index c872888..dc950e1 100644
--- a/dsn.go
+++ b/dsn.go
@@ -83,8 +83,10 @@ type Config struct {
 
 	Tracing string // sets logging level
 
-	MfaToken string // Internally used to cache the MFA token
-	IDToken  string // Internally used to cache the Id Token for external browser
+	MfaToken                       string     // Internally used to cache the MFA token
+	IDToken                        string     // Internally used to cache the Id Token for external browser
+	ClientRequestMfaToken          ConfigBool // When true the MFA token is cached in the credential manager. True by default in Windows/OSX. False for Linux.
+	ClientStoreTemporaryCredential ConfigBool // When true the ID token is cached in the credential manager. True by default in Windows/OSX. False for Linux.
 }
 
 // ocspMode returns the OCSP mode in string INSECURE, FAIL_OPEN, FAIL_CLOSED


Test patch:
diff --git a/auth_test.go b/auth_test.go
index 8ea27c7..c1048bb 100644
--- a/auth_test.go
+++ b/auth_test.go
@@ -6,10 +6,13 @@ import (
 	"context"
 	"crypto/rand"
 	"crypto/rsa"
+	"database/sql"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"net/url"
+	"os"
+	"runtime"
 	"testing"
 	"time"
 
@@ -232,6 +235,9 @@ func postAuthCheckUsernamePasswordMfa(_ context.Context, _ *snowflakeRestful, _
 		return nil, err
 	}
 
+	if ar.Data.SessionParameters["CLIENT_REQUEST_MFA_TOKEN"] != true {
+		return nil, fmt.Errorf("expected client_request_mfa_token to be true but was %v", ar.Data.SessionParameters["CLIENT_REQUEST_MFA_TOKEN"])
+	}
 	return &authResponse{
 		Success: true,
 		Data: authResponseMain{
@@ -245,6 +251,28 @@ func postAuthCheckUsernamePasswordMfa(_ context.Context, _ *snowflakeRestful, _
 	}, nil
 }
 
+func postAuthCheckExternalBrowser(_ context.Context, _ *snowflakeRestful, _ *url.Values, _ map[string]string, jsonBody []byte, _ time.Duration) (*authResponse, error) {
+	var ar authRequest
+	if err := json.Unmarshal(jsonBody, &ar); err != nil {
+		return nil, err
+	}
+
+	if ar.Data.SessionParameters["CLIENT_STORE_TEMPORARY_CREDENTIAL"] != true {
+		return nil, fmt.Errorf("expected client_store_temporary_credential to be true but was %v", ar.Data.SessionParameters["CLIENT_STORE_TEMPORARY_CREDENTIAL"])
+	}
+	return &authResponse{
+		Success: true,
+		Data: authResponseMain{
+			Token:       "t",
+			MasterToken: "m",
+			IDToken:     "mockedIDToken",
+			SessionInfo: authResponseSessionInfo{
+				DatabaseName: "dbn",
+			},
+		},
+	}, nil
+}
+
 func getDefaultSnowflakeConn() *snowflakeConn {
 	cfg := Config{
 		Account:            "a",
@@ -497,11 +525,82 @@ func TestUnitAuthenticateUsernamePasswordMfa(t *testing.T) {
 	}
 	sc := getDefaultSnowflakeConn()
 	sc.cfg.Authenticator = AuthTypeUsernamePasswordMFA
-	requestMfaToken := "true"
-	sc.cfg.Params[clientRequestMfaToken] = &requestMfaToken
+	sc.cfg.ClientRequestMfaToken = ConfigBoolTrue
 	sc.rest = sr
 	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
 	if err != nil {
 		t.Fatalf("failed to run. err: %v", err)
 	}
 }
+
+func TestUnitAuthenticateExternalBrowser(t *testing.T) {
+	var err error
+	sr := &snowflakeRestful{
+		FuncPostAuth:  postAuthCheckExternalBrowser,
+		TokenAccessor: getSimpleTokenAccessor(),
+	}
+	sc := getDefaultSnowflakeConn()
+	sc.cfg.Authenticator = AuthTypeExternalBrowser
+	sc.cfg.ClientStoreTemporaryCredential = ConfigBoolTrue
+	sc.rest = sr
+	_, err = authenticate(context.TODO(), sc, []byte{}, []byte{})
+	if err != nil {
+		t.Fatalf("failed to run. err: %v", err)
+	}
+}
+
+// To run this test you need to set environment variables in parameters.json to a user with MFA authentication enabled
+// Set any other snowflake_test variables needed for database, schema, role for this user
+func TestUsernamePasswordMfaCaching(t *testing.T) {
+	t.Skip("manual test for MFA token caching")
+
+	config, err := ParseDSN(dsn)
+	if err != nil {
+		t.Fatal("Failed to parse dsn")
+	}
+	// connect with MFA authentication
+	user := os.Getenv("SNOWFLAKE_TEST_MFA_USER")
+	password := os.Getenv("SNOWFLAKE_TEST_MFA_PASSWORD")
+	config.User = user
+	config.Password = password
+	config.Authenticator = AuthTypeUsernamePasswordMFA
+	if runtime.GOOS == "linux" {
+		config.ClientRequestMfaToken = ConfigBoolTrue
+	}
+	connector := NewConnector(SnowflakeDriver{}, *config)
+	db := sql.OpenDB(connector)
+	for i := 0; i < 3; i++ {
+		// should only be prompted to authenticate first time around.
+		_, err := db.Query("select current_user()")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+// To run this test you need to set SNOWFLAKE_TEST_EXT_BROWSER_USER environment variable to an external browser user
+// Set any other snowflake_test variables needed for database, schema, role for this user
+func TestExternalBrowserCaching(t *testing.T) {
+	t.Skip("manual test for external browser token caching")
+
+	config, err := ParseDSN(dsn)
+	if err != nil {
+		t.Fatal("Failed to parse dsn")
+	}
+	// connect with external browser authentication
+	user := os.Getenv("SNOWFLAKE_TEST_EXT_BROWSER_USER")
+	config.User = user
+	config.Authenticator = AuthTypeExternalBrowser
+	if runtime.GOOS == "linux" {
+		config.ClientStoreTemporaryCredential = ConfigBoolTrue
+	}
+	connector := NewConnector(SnowflakeDriver{}, *config)
+	db := sql.OpenDB(connector)
+	for i := 0; i < 3; i++ {
+		// should only be prompted to authenticate first time around.
+		_, err := db.Query("select current_user()")
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+}