Repository: snowflakedb/gosnowflake
Pull request #1462 (2025-07-10 09:42:19)

SNOW-2113973: Add support for opt-in single use refresh token in OAuth flow
### Description

SNOW-2113973 This PR introduces support for opt-in change of refresh token type returned from Snowflake IdP during OAuth authentication. With the flag enabled, the refresh token will be single-use instead of time-valid.

### Checklist
- [ ] Created tests which fail without the change (if possible)
- [ ] Extended the README / documentation, if necessary


sfc-gh-pbulawa: Add a description of the PR please.
codecov[bot]: ### :x: 2 Tests Failed:
| Tests completed | Failed | Passed | Skipped |
|---|---|---|---|
| 2758 | 2 | 2756 | 69 |
<details><summary>View the full list of 2 :snowflake: flaky tests</summary>

> <pre><code class="language-python">github.com/snowflakedb/gosnowflake::TestOCSP/https://s3-us-west-2.amazonaws.com/sfc-snowsql-updates/?prefix=1.1/windows_x86_64_true#01</code></pre>
> **Flake rate in main:** 6.67% (Passed 28 times, Failed 2 times)
> <details><summary>Stack Traces | 20.9s run time</summary>
> 
> > <pre><code class="language-python">Failed</code></pre>
> 
> </details>


> <pre><code class="language-python">github.com/snowflakedb/gosnowflake::TestOCSP</code></pre>
> **Flake rate in main:** 6.67% (Passed 28 times, Failed 2 times)
> <details><summary>Stack Traces | 29s run time</summary>
> 
> > <pre><code class="language-python">Failed</code></pre>
> 
> </details>

</details>

To view more test analytics, go to the [Test Analytics Dashboard](https://app.codecov.io/gh/snowflakedb/gosnowflake/tests/SNOW-2113973-token-rotation)
<sub>📋 Got 3 mins? [Take this short survey](https://forms.gle/BpocVj23nhr2Y45G7) to help us improve Test Analytics.</sub>

Patch:
diff --git a/auth_oauth.go b/auth_oauth.go
index ea69673..7c62951 100644
--- a/auth_oauth.go
+++ b/auth_oauth.go
@@ -8,8 +8,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"golang.org/x/oauth2"
-	"golang.org/x/oauth2/clientcredentials"
 	"html"
 	"io"
 	"net"
@@ -18,6 +16,9 @@ import (
 	"strconv"
 	"strings"
 	"time"
+
+	"golang.org/x/oauth2"
+	"golang.org/x/oauth2/clientcredentials"
 )
 
 const (
@@ -207,7 +208,11 @@ func (oauthClient *oauthClient) exchangeAccessToken(codeReq *http.Request, state
 	}
 
 	code := queryParams.Get("code")
-	token, err := oauth2cfg.Exchange(oauthClient.ctx, code, oauth2.VerifierOption(codeVerifier))
+	opts := []oauth2.AuthCodeOption{oauth2.VerifierOption(codeVerifier)}
+	if oauthClient.cfg.EnableSingleUseRefreshTokens {
+		opts = append(opts, oauth2.SetAuthURLParam("enable_single_use_refresh_tokens", "true"))
+	}
+	token, err := oauth2cfg.Exchange(oauthClient.ctx, code, opts...)
 	if err != nil {
 		responseBodyChan <- err.Error()
 		return nil, err
diff --git a/dsn.go b/dsn.go
index 7da1292..5f3144a 100644
--- a/dsn.go
+++ b/dsn.go
@@ -52,12 +52,13 @@ type Config struct {
 	Role      string // Role
 	Region    string // Region
 
-	OauthClientID         string // Client id for OAuth2 external IdP
-	OauthClientSecret     string // Client secret for OAuth2 external IdP
-	OauthAuthorizationURL string // Authorization URL of Auth2 external IdP
-	OauthTokenRequestURL  string // Token request URL of Auth2 external IdP
-	OauthRedirectURI      string // Redirect URI registered in IdP. The default is http://127.0.0.1:<random port>/
-	OauthScope            string // Comma separated list of scopes. If empty it is derived from role.
+	OauthClientID                string // Client id for OAuth2 external IdP
+	OauthClientSecret            string // Client secret for OAuth2 external IdP
+	OauthAuthorizationURL        string // Authorization URL of Auth2 external IdP
+	OauthTokenRequestURL         string // Token request URL of Auth2 external IdP
+	OauthRedirectURI             string // Redirect URI registered in IdP. The default is http://127.0.0.1:<random port>/
+	OauthScope                   string // Comma separated list of scopes. If empty it is derived from role.
+	EnableSingleUseRefreshTokens bool   // Enables single use refresh tokens for Snowflake IdP
 
 	// ValidateDefaultParameters disable the validation checks for Database, Schema, Warehouse and Role
 	// at the time a connection is established
@@ -217,6 +218,9 @@ func DSN(cfg *Config) (dsn string, err error) {
 	if cfg.OauthScope != "" {
 		params.Add("oauthScope", cfg.OauthScope)
 	}
+	if cfg.EnableSingleUseRefreshTokens {
+		params.Add("enableSingleUseRefreshTokens", strconv.FormatBool(cfg.EnableSingleUseRefreshTokens))
+	}
 	if cfg.WorkloadIdentityProvider != "" {
 		params.Add("workloadIdentityProvider", cfg.WorkloadIdentityProvider)
 	}
@@ -763,6 +767,13 @@ func parseDSNParams(cfg *Config, params string) (err error) {
 			cfg.OauthRedirectURI = value
 		case "oauthScope":
 			cfg.OauthScope = value
+		case "enableSingleUseRefreshTokens":
+			var vv bool
+			vv, err = strconv.ParseBool(value)
+			if err != nil {
+				return
+			}
+			cfg.EnableSingleUseRefreshTokens = vv
 		case "passcodeInPassword":
 			var vv bool
 			vv, err = strconv.ParseBool(value)


Test patch:
diff --git a/auth_oauth_test.go b/auth_oauth_test.go
index 08f2e1a..aa831ca 100644
--- a/auth_oauth_test.go
+++ b/auth_oauth_test.go
@@ -4,12 +4,13 @@ import (
 	"context"
 	"database/sql"
 	"errors"
-	"golang.org/x/oauth2"
 	"io"
 	"net/http"
 	"sync"
 	"testing"
 	"time"
+
+	"golang.org/x/oauth2"
 )
 
 func TestUnitOAuthAuthorizationCode(t *testing.T) {
@@ -283,6 +284,26 @@ func TestAuthorizationCodeFlow(t *testing.T) {
 		runSmokeQuery(t, db)
 	})
 
+	t.Run("successful flow with single-use refresh token enabled", func(t *testing.T) {
+		wiremock.registerMappings(t,
+			newWiremockMapping("oauth2/authorization_code/successful_flow_with_single_use_refresh_token.json"),
+			newWiremockMapping("oauth2/login_request.json"),
+			newWiremockMapping("select1.json"))
+		cfg := wiremock.connectionConfig()
+		cfg.Role = "ANALYST"
+		cfg.Authenticator = AuthTypeOAuthAuthorizationCode
+		cfg.OauthRedirectURI = "http://localhost:1234/snowflake/oauth-redirect"
+		cfg.Transporter = roundTripper
+		cfg.EnableSingleUseRefreshTokens = true
+		oauthAccessTokenSpec := newOAuthAccessTokenSpec(cfg.OauthTokenRequestURL, cfg.User)
+		oauthRefreshTokenSpec := newOAuthRefreshTokenSpec(cfg.OauthTokenRequestURL, cfg.User)
+		credentialsStorage.deleteCredential(oauthAccessTokenSpec)
+		credentialsStorage.deleteCredential(oauthRefreshTokenSpec)
+		connector := NewConnector(SnowflakeDriver{}, *cfg)
+		db := sql.OpenDB(connector)
+		runSmokeQuery(t, db)
+	})
+
 	t.Run("should use cached access token", func(t *testing.T) {
 		roundTripper.reset()
 		wiremock.registerMappings(t,
diff --git a/dsn_test.go b/dsn_test.go
index 3f2d829..d38d54f 100644
--- a/dsn_test.go
+++ b/dsn_test.go
@@ -262,6 +262,24 @@ func TestParseDSN(t *testing.T) {
 			ocspMode: ocspModeFailOpen,
 			err:      nil,
 		},
+		{
+			dsn: "user:pass@account?oauthRedirectUri=http:%2F%2Flocalhost:8001%2Fsome-path&oauthClientId=testClientId&oauthClientSecret=testClientSecret&oauthAuthorizationUrl=http:%2F%2Fsomehost.com&oauthTokenRequestUrl=https:%2F%2Fsomehost2.com%2Fsomepath&oauthScope=test+scope&enableSingleUseRefreshTokens=true",
+			config: &Config{
+				Account: "account", User: "user", Password: "pass",
+				Protocol: "https", Host: "account.snowflakecomputing.com", Port: 443,
+				OauthClientID: "testClientId", OauthClientSecret: "testClientSecret", OauthAuthorizationURL: "http://somehost.com", OauthTokenRequestURL: "https://somehost2.com/somepath", OauthRedirectURI: "http://localhost:8001/some-path", OauthScope: "test scope",
+				EnableSingleUseRefreshTokens: true,
+				OCSPFailOpen:                 OCSPFailOpenTrue,
+				ValidateDefaultParameters:    ConfigBoolTrue,
+				ClientTimeout:                defaultClientTimeout,
+				JWTClientTimeout:             defaultJWTClientTimeout,
+				ExternalBrowserTimeout:       defaultExternalBrowserTimeout,
+				CloudStorageTimeout:          defaultCloudStorageTimeout,
+				IncludeRetryReason:           ConfigBoolTrue,
+			},
+			ocspMode: ocspModeFailOpen,
+			err:      nil,
+		},
 		{
 			dsn: "user:pass@host:123/db/schema?account=ac&protocol=http",
 			config: &Config{
@@ -1350,6 +1368,9 @@ func TestParseDSN(t *testing.T) {
 				if test.config.OauthScope != cfg.OauthScope {
 					t.Fatalf("%v: Failed to match OauthScope. expected: %v, got: %v", i, test.config.OauthScope, cfg.OauthScope)
 				}
+				if test.config.EnableSingleUseRefreshTokens != cfg.EnableSingleUseRefreshTokens {
+					t.Fatalf("%v: Failed to match EnableSingleUseRefreshTokens. expected: %v, got: %v", i, test.config.OauthScope, cfg.OauthScope)
+				}
 				assertEqualE(t, cfg.Token, test.config.Token, "token")
 				assertEqualE(t, cfg.ClientConfigFile, test.config.ClientConfigFile, "client config file")
 			case test.err != nil:
@@ -1519,6 +1540,22 @@ func TestDSN(t *testing.T) {
 			},
 			dsn: "u:p@a.r.snowflakecomputing.com:443?oauthAuthorizationUrl=http%3A%2F%2Fsomehost.com&oauthClientId=testClientId&oauthClientSecret=testClientSecret&oauthRedirectUri=http%3A%2F%2Flocalhost%3A8001%2Fsome-path&oauthScope=test+scope&oauthTokenRequestUrl=https%3A%2F%2Fsomehost2.com%2Fsomepath&ocspFailOpen=true&region=r&validateDefaultParameters=true",
 		},
+		{
+			cfg: &Config{
+				User:                         "u",
+				Password:                     "p",
+				Account:                      "a",
+				Region:                       "r",
+				OauthClientID:                "testClientId",
+				OauthClientSecret:            "testClientSecret",
+				OauthAuthorizationURL:        "http://somehost.com",
+				OauthTokenRequestURL:         "https://somehost2.com/somepath",
+				OauthRedirectURI:             "http://localhost:8001/some-path",
+				OauthScope:                   "test scope",
+				EnableSingleUseRefreshTokens: true,
+			},
+			dsn: "u:p@a.r.snowflakecomputing.com:443?enableSingleUseRefreshTokens=true&oauthAuthorizationUrl=http%3A%2F%2Fsomehost.com&oauthClientId=testClientId&oauthClientSecret=testClientSecret&oauthRedirectUri=http%3A%2F%2Flocalhost%3A8001%2Fsome-path&oauthScope=test+scope&oauthTokenRequestUrl=https%3A%2F%2Fsomehost2.com%2Fsomepath&ocspFailOpen=true&region=r&validateDefaultParameters=true",
+		},
 		{
 			cfg: &Config{
 				User:                   "u",
diff --git a/test_data/wiremock/mappings/oauth2/authorization_code/successful_flow_with_single_use_refresh_token.json b/test_data/wiremock/mappings/oauth2/authorization_code/successful_flow_with_single_use_refresh_token.json
new file mode 100644
index 0000000..734f4b7
--- /dev/null
+++ b/test_data/wiremock/mappings/oauth2/authorization_code/successful_flow_with_single_use_refresh_token.json
@@ -0,0 +1,82 @@
+{
+  "mappings": [
+    {
+      "request": {
+        "urlPathPattern": "/oauth/authorize",
+        "queryParameters": {
+          "response_type": {
+            "equalTo": "code"
+          },
+          "scope": {
+            "equalTo": "session:role:ANALYST"
+          },
+          "code_challenge_method": {
+            "equalTo": "S256"
+          },
+          "redirect_uri": {
+            "equalTo": "http://localhost:1234/snowflake/oauth-redirect"
+          },
+          "code_challenge": {
+            "matches": "JZpN_-zfNduuWm-zUo-D-m7vMw_pgUGv8wGDGqBR8PM"
+          },
+          "state": {
+            "matches": "testState|invalidState"
+          },
+          "client_id": {
+            "equalTo": "testClientId"
+          }
+        },
+        "method": "GET"
+      },
+      "response": {
+        "status": 302,
+        "headers": {
+          "Location": "http://localhost:1234/snowflake/oauth-redirect?code=testCode&state=testState"
+        }
+      }
+    },
+    {
+      "request": {
+        "urlPathPattern": "/oauth/token",
+        "method": "POST",
+        "headers": {
+          "Content-Type": {
+            "contains": "application/x-www-form-urlencoded"
+          },
+          "Authorization": {
+            "equalTo": "Basic dGVzdENsaWVudElkOnRlc3RDbGllbnRTZWNyZXQ="
+          }
+        },
+        "formParameters": {
+          "grant_type": {
+            "equalTo": "authorization_code"
+          },
+          "code_verifier": {
+            "matches": "testCodeVerifier"
+          },
+          "code": {
+            "equalTo": "testCode"
+          },
+          "redirect_uri": {
+            "equalTo": "http://localhost:1234/snowflake/oauth-redirect"
+          },
+          "enable_single_use_refresh_tokens": {
+            "equalTo": "true"
+          }
+        }
+      },
+      "response": {
+        "status": 200,
+        "jsonBody": {
+          "access_token": "access-token-123",
+          "token_type": "Bearer",
+          "username": "test-user",
+          "scope": "refresh_token session:role:ANALYST",
+          "expires_in": 600,
+          "refresh_token_expires_in": 86399,
+          "idpInitiated": false
+        }
+      }
+    }
+  ]
+}
\ No newline at end of file