Repository: snowflakedb/gosnowflake
Pull request #1032 (2024-01-17 06:35:20)

SNOW-916951: Retry okta auth with new token
### Description
issue: https://github.com/snowflakedb/snowflake-sdks-drivers-issues-teamwork/issues/655

The native Okta Authentication fails in retry due to the same token is used in an sso/saml retry.
This PR fixes the retry logic by generating a new token each time it retries.

### 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



codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1032?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`5c89d42`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/5c89d42a335ea41256e87ff71317ef6f8df6a086?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) 83.14% compared to head [(`946a1d4`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1032?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) 81.58%.


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


```diff
@@            Coverage Diff             @@
##           master    #1032      +/-   ##
==========================================
- Coverage   83.14%   81.58%   -1.57%     
==========================================
  Files          50       49       -1     
  Lines        9761     9757       -4     
==========================================
- Hits         8116     7960     -156     
- Misses       1645     1797     +152     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1032?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 c844b2e..cc839ca 100644
--- a/auth.go
+++ b/auth.go
@@ -418,6 +418,17 @@ func createRequestBody(sc *snowflakeConn, sessionParameters map[string]interface
 		requestMain.Authenticator = AuthTypeOAuth.String()
 		requestMain.Token = sc.cfg.Token
 	case AuthTypeOkta:
+		samlResponse, err := authenticateBySAML(
+			sc.ctx,
+			sc.rest,
+			sc.cfg.OktaURL,
+			sc.cfg.Application,
+			sc.cfg.Account,
+			sc.cfg.User,
+			sc.cfg.Password)
+		if err != nil {
+			return nil, err
+		}
 		requestMain.RawSAMLResponse = string(samlResponse)
 	case AuthTypeJwt:
 		requestMain.Authenticator = AuthTypeJwt.String()
@@ -535,19 +546,6 @@ func authenticateWithConfig(sc *snowflakeConn) error {
 				return err
 			}
 		}
-	case AuthTypeOkta:
-		samlResponse, err = authenticateBySAML(
-			sc.ctx,
-			sc.rest,
-			sc.cfg.OktaURL,
-			sc.cfg.Application,
-			sc.cfg.Account,
-			sc.cfg.User,
-			sc.cfg.Password)
-		if err != nil {
-			sc.cleanup()
-			return err
-		}
 	}
 	authData, err = authenticate(
 		sc.ctx,


Test patch:
diff --git a/auth_test.go b/auth_test.go
index 25e2d0d..0addaab 100644
--- a/auth_test.go
+++ b/auth_test.go
@@ -104,8 +104,11 @@ func postAuthSuccess(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *
 
 func postAuthCheckSAMLResponse(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {
 	var ar authRequest
-	jsonBody, _ := bodyCreator()
-	if err := json.Unmarshal(jsonBody, &ar); err != nil {
+	jsonBody, err := bodyCreator()
+	if err != nil {
+		return nil, err
+	}
+	if err = json.Unmarshal(jsonBody, &ar); err != nil {
 		return nil, err
 	}
 	if ar.Data.RawSAMLResponse == "" {
@@ -368,6 +371,51 @@ func postAuthCheckExternalBrowserFailed(_ context.Context, _ *snowflakeRestful,
 	}, nil
 }
 
+func postAuthOktaWithNewToken(_ context.Context, _ *snowflakeRestful, _ *http.Client, _ *url.Values, _ map[string]string, bodyCreator bodyCreatorType, _ time.Duration) (*authResponse, error) {
+	var ar authRequest
+
+	cfg := &Config{
+		Authenticator: AuthTypeOkta,
+	}
+
+	// Retry 3 times and success
+	client := &fakeHTTPClient{
+		cnt:        3,
+		success:    true,
+		statusCode: 429,
+	}
+
+	urlPtr, err := url.Parse("https://fakeaccountretrylogin.snowflakecomputing.com:443/login-request?request_guid=testguid")
+	if err != nil {
+		return &authResponse{}, err
+	}
+
+	body := func() ([]byte, error) {
+		jsonBody, _ := bodyCreator()
+		if err := json.Unmarshal(jsonBody, &ar); err != nil {
+			return nil, err
+		}
+		return jsonBody, err
+	}
+
+	_, err = newRetryHTTP(context.Background(), client, emptyRequest, urlPtr, make(map[string]string), 60*time.Second, 3, defaultTimeProvider, cfg).doPost().setBodyCreator(body).execute()
+	if err != nil {
+		return &authResponse{}, err
+	}
+
+	return &authResponse{
+		Success: true,
+		Data: authResponseMain{
+			Token:       "t",
+			MasterToken: "m",
+			MfaToken:    "mockedMfaToken",
+			SessionInfo: authResponseSessionInfo{
+				DatabaseName: "dbn",
+			},
+		},
+	}, nil
+}
+
 func getDefaultSnowflakeConn() *snowflakeConn {
 	sc := &snowflakeConn{
 		rest: &snowflakeRestful{
@@ -524,20 +572,24 @@ func TestUnitAuthenticate(t *testing.T) {
 func TestUnitAuthenticateSaml(t *testing.T) {
 	var err error
 	sr := &snowflakeRestful{
-		FuncPostAuth:  postAuthCheckSAMLResponse,
-		TokenAccessor: getSimpleTokenAccessor(),
+		Protocol:         "https",
+		Host:             "abc.com",
+		Port:             443,
+		FuncPostAuthSAML: postAuthSAMLAuthSuccess,
+		FuncPostAuthOKTA: postAuthOKTASuccess,
+		FuncGetSSO:       getSSOSuccess,
+		FuncPostAuth:     postAuthCheckSAMLResponse,
+		TokenAccessor:    getSimpleTokenAccessor(),
 	}
 	sc := getDefaultSnowflakeConn()
 	sc.cfg.Authenticator = AuthTypeOkta
 	sc.cfg.OktaURL = &url.URL{
 		Scheme: "https",
-		Host:   "blah.okta.com",
+		Host:   "abc.com",
 	}
 	sc.rest = sr
-	_, err = authenticate(context.Background(), sc, []byte("HTML data in bytes from"), []byte{})
-	if err != nil {
-		t.Fatalf("failed to run. err: %v", err)
-	}
+	_, err = authenticate(context.Background(), sc, []byte{}, []byte{})
+	assertNilF(t, err, "failed to run.")
 }
 
 // Unit test for OAuth.
@@ -665,7 +717,7 @@ func TestUnitAuthenticateWithConfigOkta(t *testing.T) {
 		FuncPostAuthSAML: postAuthSAMLAuthSuccess,
 		FuncPostAuthOKTA: postAuthOKTASuccess,
 		FuncGetSSO:       getSSOSuccess,
-		FuncPostAuth:     postAuthSuccess,
+		FuncPostAuth:     postAuthCheckSAMLResponse,
 		TokenAccessor:    getSimpleTokenAccessor(),
 	}
 	sc := getDefaultSnowflakeConn()
@@ -840,3 +892,36 @@ func TestDisableExternalBrowserCaching(t *testing.T) {
 		}
 	}
 }
+
+func TestOktaRetryWithNewToken(t *testing.T) {
+	expectedMasterToken := "m"
+	expectedToken := "t"
+	expectedMfaToken := "mockedMfaToken"
+	expectedDatabaseName := "dbn"
+
+	sr := &snowflakeRestful{
+		Protocol:         "https",
+		Host:             "abc.com",
+		Port:             443,
+		FuncPostAuthSAML: postAuthSAMLAuthSuccess,
+		FuncPostAuthOKTA: postAuthOKTASuccess,
+		FuncGetSSO:       getSSOSuccess,
+		FuncPostAuth:     postAuthOktaWithNewToken,
+		TokenAccessor:    getSimpleTokenAccessor(),
+	}
+	sc := getDefaultSnowflakeConn()
+	sc.cfg.Authenticator = AuthTypeOkta
+	sc.cfg.OktaURL = &url.URL{
+		Scheme: "https",
+		Host:   "abc.com",
+	}
+	sc.rest = sr
+	sc.ctx = context.Background()
+
+	authResponse, err := authenticate(context.Background(), sc, []byte{0x12, 0x34}, []byte{0x56, 0x78})
+	assertNilF(t, err, "should not have failed to run authenticate()")
+	assertEqualF(t, authResponse.MasterToken, expectedMasterToken)
+	assertEqualF(t, authResponse.Token, expectedToken)
+	assertEqualF(t, authResponse.MfaToken, expectedMfaToken)
+	assertEqualF(t, authResponse.SessionInfo.DatabaseName, expectedDatabaseName)
+}