Repository: snowflakedb/gosnowflake
Pull request #1128 (2024-05-11 02:00:03)

SNOW-1226624 Parameter to disable SAML URL check
### Description

[issue #942](https://github.com/snowflakedb/snowflake-sdks-drivers-issues-teamwork/issues/942)

### Checklist
- [X] Created tests which fail without the change (if possible)
- [X] Extended the README / documentation, if necessary


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1128?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
Attention: Patch coverage is `93.54839%` with `2 lines` in your changes are missing coverage. Please review.
> Project coverage is 75.42%. Comparing base [(`c355711`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/c355711dbd1f9ab10dfbfdcdd194656c23abb45d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) to head [(`e26cc2f`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1128?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).

| [Files](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1128?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Patch % | Lines |
|---|---|---|
| [dsn.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1128?src=pr&el=tree&filepath=dsn.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-ZHNuLmdv) | 85.71% | [2 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1128?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #1128      +/-   ##
==========================================
+ Coverage   75.40%   75.42%   +0.01%     
==========================================
  Files          53       53              
  Lines       11250    11266      +16     
==========================================
+ Hits         8483     8497      +14     
- Misses       2767     2769       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1128?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).


Patch:
diff --git a/auth.go b/auth.go
index 6d069a9..ebd3d23 100644
--- a/auth.go
+++ b/auth.go
@@ -428,7 +428,8 @@ func createRequestBody(sc *snowflakeConn, sessionParameters map[string]interface
 			sc.cfg.Application,
 			sc.cfg.Account,
 			sc.cfg.User,
-			sc.cfg.Password)
+			sc.cfg.Password,
+			sc.cfg.DisableSamlURLCheck)
 		if err != nil {
 			return nil, err
 		}
diff --git a/authokta.go b/authokta.go
index d3ea90f..a2dfed8 100644
--- a/authokta.go
+++ b/authokta.go
@@ -59,6 +59,7 @@ func authenticateBySAML(
 	account string,
 	user string,
 	password string,
+	disableSamlURLCheck ConfigBool,
 ) (samlResponse []byte, err error) {
 	logger.WithContext(ctx).Info("step 1: query GS to obtain IDP token and SSO url")
 	headers := make(map[string]string)
@@ -152,20 +153,22 @@ func authenticateBySAML(
 	if err != nil {
 		return nil, err
 	}
-	logger.WithContext(ctx).Info("step 5: validate post_back_url matches Snowflake URL")
-	tgtURL, err := postBackURL(bd)
-	if err != nil {
-		return nil, err
-	}
+	if disableSamlURLCheck == ConfigBoolFalse {
+		logger.WithContext(ctx).Info("step 5: validate post_back_url matches Snowflake URL")
+		tgtURL, err := postBackURL(bd)
+		if err != nil {
+			return nil, err
+		}
 
-	fullURL := sr.getURL()
-	logger.WithContext(ctx).Infof("tgtURL: %v, origURL: %v", tgtURL, fullURL)
-	if !isPrefixEqual(tgtURL, fullURL) {
-		return nil, &SnowflakeError{
-			Number:      ErrCodeSSOURLNotMatch,
-			SQLState:    SQLStateConnectionRejected,
-			Message:     errMsgSSOURLNotMatch,
-			MessageArgs: []interface{}{tgtURL, fullURL},
+		fullURL := sr.getURL()
+		logger.WithContext(ctx).Infof("tgtURL: %v, origURL: %v", tgtURL, fullURL)
+		if !isPrefixEqual(tgtURL, fullURL) {
+			return nil, &SnowflakeError{
+				Number:      ErrCodeSSOURLNotMatch,
+				SQLState:    SQLStateConnectionRejected,
+				Message:     errMsgSSOURLNotMatch,
+				MessageArgs: []interface{}{tgtURL, fullURL},
+			}
 		}
 	}
 	return bd, nil
diff --git a/doc.go b/doc.go
index 2529263..25f0eda 100644
--- a/doc.go
+++ b/doc.go
@@ -130,6 +130,8 @@ The following connection parameters are supported:
   - clientConfigFile: specifies the location of the client configuration json file.
     In this file you can configure Easy Logging feature.
 
+  - disableSamlURLCheck: disables the SAML URL check. Default value is false.
+
 All other parameters are interpreted as session parameters (https://docs.snowflake.com/en/sql-reference/parameters.html).
 For example, the TIMESTAMP_OUTPUT_FORMAT session parameter can be set by adding:
 
diff --git a/dsn.go b/dsn.go
index e2c541a..aefcee2 100644
--- a/dsn.go
+++ b/dsn.go
@@ -107,6 +107,8 @@ type Config struct {
 	ClientConfigFile string // File path to the client configuration json file
 
 	DisableConsoleLogin ConfigBool // Indicates whether console login should be disabled
+
+	DisableSamlURLCheck ConfigBool // Indicates whether the SAML URL check should be disabled
 }
 
 // Validate enables testing if config is correct.
@@ -267,6 +269,9 @@ func DSN(cfg *Config) (dsn string, err error) {
 	if cfg.DisableConsoleLogin != configBoolNotSet {
 		params.Add("disableConsoleLogin", strconv.FormatBool(cfg.DisableConsoleLogin != ConfigBoolFalse))
 	}
+	if cfg.DisableSamlURLCheck != configBoolNotSet {
+		params.Add("disableSamlURLCheck", strconv.FormatBool(cfg.DisableSamlURLCheck != ConfigBoolFalse))
+	}
 
 	dsn = fmt.Sprintf("%v:%v@%v:%v", url.QueryEscape(cfg.User), url.QueryEscape(cfg.Password), cfg.Host, cfg.Port)
 	if params.Encode() != "" {
@@ -770,6 +775,17 @@ func parseDSNParams(cfg *Config, params string) (err error) {
 			} else {
 				cfg.DisableConsoleLogin = ConfigBoolFalse
 			}
+		case "disableSamlURLCheck":
+			var vv bool
+			vv, err = strconv.ParseBool(value)
+			if err != nil {
+				return
+			}
+			if vv {
+				cfg.DisableSamlURLCheck = ConfigBoolTrue
+			} else {
+				cfg.DisableSamlURLCheck = ConfigBoolFalse
+			}
 		default:
 			if cfg.Params == nil {
 				cfg.Params = make(map[string]*string)


Test patch:
diff --git a/authokta_test.go b/authokta_test.go
index 56e1512..a1028e7 100644
--- a/authokta_test.go
+++ b/authokta_test.go
@@ -233,64 +233,95 @@ func TestUnitAuthenticateBySAML(t *testing.T) {
 		TokenAccessor:    getSimpleTokenAccessor(),
 	}
 	var err error
-	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
 	assertNotNilF(t, err, "should have failed at FuncPostAuthSAML.")
 	assertEqualE(t, err.Error(), "failed to get SAML response")
 
 	sr.FuncPostAuthSAML = postAuthSAMLAuthFail
-	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
 	assertNotNilF(t, err, "should have failed at FuncPostAuthSAML.")
 	assertEqualE(t, err.Error(), "strconv.Atoi: parsing \"\": invalid syntax")
 
 	sr.FuncPostAuthSAML = postAuthSAMLAuthFailWithCode
-	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
 	assertNotNilF(t, err, "should have failed at FuncPostAuthSAML.")
 	driverErr, ok := err.(*SnowflakeError)
 	assertTrueF(t, ok, "should be a SnowflakeError")
 	assertEqualE(t, driverErr.Number, ErrCodeIdpConnectionError)
 
 	sr.FuncPostAuthSAML = postAuthSAMLAuthSuccessButInvalidURL
-	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
 	assertNotNilF(t, err, "should have failed at FuncPostAuthSAML.")
 	driverErr, ok = err.(*SnowflakeError)
 	assertTrueF(t, ok, "should be a SnowflakeError")
 	assertEqualE(t, driverErr.Number, ErrCodeIdpConnectionError)
 
 	sr.FuncPostAuthSAML = postAuthSAMLAuthSuccessButInvalidTokenURL
-	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
 	assertNotNilF(t, err, "should have failed at FuncPostAuthSAML.")
 	assertEqualE(t, err.Error(), "failed to parse token URL. invalid!@url$%^")
 
 	sr.FuncPostAuthSAML = postAuthSAMLAuthSuccessButInvalidSSOURL
-	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
 	assertNotNilF(t, err, "should have failed at FuncPostAuthSAML.")
 	assertEqualE(t, err.Error(), "failed to parse SSO URL. invalid!@url$%^")
 
 	sr.FuncPostAuthSAML = postAuthSAMLAuthSuccess
 	sr.FuncPostAuthOKTA = postAuthOKTAError
-	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
 	assertNotNilF(t, err, "should have failed at FuncPostAuthOKTA.")
 	assertEqualE(t, err.Error(), "failed to get SAML response")
 
 	sr.FuncPostAuthOKTA = postAuthOKTASuccess
 	sr.FuncGetSSO = getSSOError
-	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
 	assertNotNilF(t, err, "should have failed at FuncGetSSO.")
 	assertEqualE(t, err.Error(), "failed to get SSO html")
 
 	sr.FuncGetSSO = getSSOSuccessButInvalidURL
-	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
 	assertNotNilF(t, err, "should have failed at FuncGetSSO.")
 	assertHasPrefixE(t, err.Error(), "failed to find action field in HTML response")
 
 	sr.FuncGetSSO = getSSOSuccess
-	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
 	assertNilF(t, err, "should have succeeded at FuncGetSSO.")
 
 	sr.FuncGetSSO = getSSOSuccessButWrongPrefixURL
-	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password)
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
 	assertNotNilF(t, err, "should have failed at FuncGetSSO.")
 	driverErr, ok = err.(*SnowflakeError)
 	assertTrueF(t, ok, "should be a SnowflakeError")
 	assertEqualE(t, driverErr.Number, ErrCodeSSOURLNotMatch)
 }
+
+func TestDisableSamlURLCheck(t *testing.T) {
+	authenticator := &url.URL{
+		Scheme: "https",
+		Host:   "abc.com",
+	}
+	application := "testapp"
+	account := "testaccount"
+	user := "u"
+	password := "p"
+	sr := &snowflakeRestful{
+		Protocol:         "https",
+		Host:             "abc.com",
+		Port:             443,
+		FuncPostAuthSAML: postAuthSAMLAuthSuccess,
+		FuncPostAuthOKTA: postAuthOKTASuccess,
+		FuncGetSSO:       getSSOSuccessButWrongPrefixURL,
+		TokenAccessor:    getSimpleTokenAccessor(),
+	}
+	var err error
+	// Test for disabled SAML URL check
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolTrue)
+	assertNilF(t, err, "SAML URL check should have disabled.")
+
+	// Test for enabled SAML URL check
+	_, err = authenticateBySAML(context.Background(), sr, authenticator, application, account, user, password, ConfigBoolFalse)
+	assertNotNilF(t, err, "should have failed at FuncGetSSO.")
+	driverErr, ok := err.(*SnowflakeError)
+	assertTrueF(t, ok, "should be a SnowflakeError")
+	assertEqualE(t, driverErr.Number, ErrCodeSSOURLNotMatch)
+}
diff --git a/dsn_test.go b/dsn_test.go
index dfa2a96..c9b032a 100644
--- a/dsn_test.go
+++ b/dsn_test.go
@@ -748,6 +748,40 @@ func TestParseDSN(t *testing.T) {
 			ocspMode: ocspModeFailOpen,
 			err:      nil,
 		},
+		{
+			dsn: "u:p@a.snowflake.local:9876?account=a&protocol=http&authenticator=EXTERNALBROWSER&disableSamlURLCheck=true",
+			config: &Config{
+				Account: "a", User: "u", Password: "p",
+				Authenticator: AuthTypeExternalBrowser,
+				Protocol:      "http", Host: "a.snowflake.local", Port: 9876,
+				OCSPFailOpen:              OCSPFailOpenTrue,
+				ValidateDefaultParameters: ConfigBoolTrue,
+				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
+				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
+				IncludeRetryReason:        ConfigBoolTrue,
+				DisableSamlURLCheck:       ConfigBoolTrue,
+			},
+			ocspMode: ocspModeFailOpen,
+			err:      nil,
+		},
+		{
+			dsn: "u:p@a.snowflake.local:9876?account=a&protocol=http&authenticator=EXTERNALBROWSER&disableSamlURLCheck=false",
+			config: &Config{
+				Account: "a", User: "u", Password: "p",
+				Authenticator: AuthTypeExternalBrowser,
+				Protocol:      "http", Host: "a.snowflake.local", Port: 9876,
+				OCSPFailOpen:              OCSPFailOpenTrue,
+				ValidateDefaultParameters: ConfigBoolTrue,
+				ClientTimeout:             defaultClientTimeout,
+				JWTClientTimeout:          defaultJWTClientTimeout,
+				ExternalBrowserTimeout:    defaultExternalBrowserTimeout,
+				IncludeRetryReason:        ConfigBoolTrue,
+				DisableSamlURLCheck:       ConfigBoolFalse,
+			},
+			ocspMode: ocspModeFailOpen,
+			err:      nil,
+		},
 	}
 
 	for _, at := range []AuthType{AuthTypeExternalBrowser, AuthTypeOAuth} {
@@ -910,6 +944,9 @@ func TestParseDSN(t *testing.T) {
 				if test.config.DisableConsoleLogin != cfg.DisableConsoleLogin {
 					t.Fatalf("%v: Failed to match DisableConsoleLogin. expected: %v, got: %v", i, test.config.DisableConsoleLogin, cfg.DisableConsoleLogin)
 				}
+				if test.config.DisableSamlURLCheck != cfg.DisableSamlURLCheck {
+					t.Fatalf("%v: Failed to match DisableSamlURLCheck. expected: %v, got: %v", i, test.config.DisableSamlURLCheck, cfg.DisableSamlURLCheck)
+				}
 				assertEqualF(t, cfg.ClientConfigFile, test.config.ClientConfigFile, "client config file")
 			case test.err != nil:
 				driverErrE, okE := test.err.(*SnowflakeError)
@@ -1379,6 +1416,26 @@ func TestDSN(t *testing.T) {
 			},
 			dsn: "u:p@a.b.c.snowflakecomputing.com:443?authenticator=externalbrowser&disableConsoleLogin=false&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
 		},
+		{
+			cfg: &Config{
+				User:                "u",
+				Password:            "p",
+				Account:             "a.b.c",
+				Authenticator:       AuthTypeExternalBrowser,
+				DisableSamlURLCheck: ConfigBoolTrue,
+			},
+			dsn: "u:p@a.b.c.snowflakecomputing.com:443?authenticator=externalbrowser&disableSamlURLCheck=true&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
+		},
+		{
+			cfg: &Config{
+				User:                "u",
+				Password:            "p",
+				Account:             "a.b.c",
+				Authenticator:       AuthTypeExternalBrowser,
+				DisableSamlURLCheck: ConfigBoolFalse,
+			},
+			dsn: "u:p@a.b.c.snowflakecomputing.com:443?authenticator=externalbrowser&disableSamlURLCheck=false&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
+		},
 	}
 	for _, test := range testcases {
 		t.Run(test.dsn, func(t *testing.T) {