SNOW-1825478- Support for Oauth PAT

snowflakedb · Mar 5, 2025 · 4c44ce4 · 4c44ce4
1 parent c097139
commit 4c44ce4
Show file tree

Hide file tree

Showing 12 changed files with 255 additions and 4 deletions.
diff --git a/lib/authentication/auth_oauth_pat.js b/lib/authentication/auth_oauth_pat.js
@@ -0,0 +1,29 @@
+const Util = require('../util');
+/**
+ * Creates an oauth PAT  authenticator.
+ *
+ * @param {String} token
+ * @param {String} password
+ *
+ * @returns {Object}
+ * @constructor
+ */
+function AuthOauthPAT(token, password) {
+  /**
+   * Update JSON body with token.
+   *
+   * @param {JSON} body
+   *
+   * @returns {null}
+   */
+  this.updateBody = function (body) {
+    if (Util.exists(token)) {
+      body['data']['TOKEN'] = token;
+    } else if (Util.exists(password)) {
+      body['data']['TOKEN'] = password;
+    }
+  };
+
+  this.authenticate = async function () {};
+}
+module.exports = AuthOauthPAT;
diff --git a/lib/authentication/authentication.js b/lib/authentication/authentication.js
@@ -2,6 +2,7 @@ const AuthDefault = require('./auth_default');
 const AuthWeb = require('./auth_web');
 const AuthKeypair = require('./auth_keypair');
 const AuthOauth = require('./auth_oauth');
+const AuthOauthPAT = require('./auth_oauth_pat');
 const AuthOkta = require('./auth_okta');
 const AuthIDToken = require('./auth_idtoken');
 const Logger = require('../logger');
@@ -58,7 +59,6 @@ exports.formAuthJSON = function formAuthJSON(
  */
 exports.getAuthenticator = function getAuthenticator(connectionConfig, httpClient) {
   const authType = connectionConfig.getAuthenticator();
-  const openExternalBrowserCallback = connectionConfig.openExternalBrowserCallback; // Important for SSO in the Snowflake VS Code extension
   let auth;
   if (authType === AuthenticationTypes.DEFAULT_AUTHENTICATOR || authType === AuthenticationTypes.USER_PWD_MFA_AUTHENTICATOR) {
     auth = new AuthDefault(connectionConfig);
@@ -70,8 +70,10 @@ exports.getAuthenticator = function getAuthenticator(connectionConfig, httpClien
     }
   } else if (authType === AuthenticationTypes.KEY_PAIR_AUTHENTICATOR) {
     auth = new AuthKeypair(connectionConfig);
-  } else if (authType === AuthenticationTypes.OAUTH_AUTHENTICATOR) {
+  } else if (authType === AuthenticationTypes.OAUTH_AUTHENTICATOR ) {
     auth = new AuthOauth(connectionConfig.getToken());
+  } else if (authType === AuthenticationTypes.PROGRAMMATIC_ACCESS_TOKEN ) {
+    auth = new AuthOauthPAT(connectionConfig.getToken(), connectionConfig.password);
   } else if (this.isOktaAuth(authType)) {
     auth = new AuthOkta(connectionConfig, httpClient);
   } else {

diff --git a/lib/authentication/authentication_types.js b/lib/authentication/authentication_types.js
@@ -6,6 +6,7 @@ const AuthenticationTypes =
   OAUTH_AUTHENTICATOR: 'OAUTH',
   USER_PWD_MFA_AUTHENTICATOR: 'USERNAME_PASSWORD_MFA',
   ID_TOKEN_AUTHENTICATOR: 'ID_TOKEN',
+  PROGRAMMATIC_ACCESS_TOKEN: 'PROGRAMMATIC_ACCESS_TOKEN'
 };
 
 module.exports = AuthenticationTypes;
diff --git a/lib/connection/connection_config.js b/lib/connection/connection_config.js
@@ -165,7 +165,8 @@ function ConnectionConfig(options, validateCredentials, qaMode, clientInfo) {
     // username is not required for oauth and external browser authenticators
     if (!Util.exists(options.authenticator) ||
       (options.authenticator.toUpperCase() !== AuthenticationTypes.OAUTH_AUTHENTICATOR &&
-        options.authenticator.toUpperCase() !== AuthenticationTypes.EXTERNAL_BROWSER_AUTHENTICATOR)) {
+        options.authenticator.toUpperCase() !== AuthenticationTypes.EXTERNAL_BROWSER_AUTHENTICATOR &&
+        options.authenticator.toUpperCase() !== AuthenticationTypes.PROGRAMMATIC_ACCESS_TOKEN)) {
       // check for missing username
       Errors.checkArgumentExists(Util.exists(options.username),
         ErrorCodes.ERR_CONN_CREATE_MISSING_USERNAME);
@@ -188,6 +189,22 @@ function ConnectionConfig(options, validateCredentials, qaMode, clientInfo) {
       Errors.checkArgumentValid(Util.isString(options.password),
         ErrorCodes.ERR_CONN_CREATE_INVALID_PASSWORD);
     }
+    if (!Util.exists(options.authenticator) ||
+      options.authenticator === AuthenticationTypes.PROGRAMMATIC_ACCESS_TOKEN) {
+      // PASSWORD or TOKEN is needed
+      Errors.checkArgumentExists(Util.exists(options.password) || Util.exists(options.token),
+        ErrorCodes.ERR_CONN_CREATE_MISSING_PASSWORD);
+
+      if (Util.exists(options.password)) {
+        // check for invalid password
+        Errors.checkArgumentValid(Util.isString(options.password),
+          ErrorCodes.ERR_CONN_CREATE_INVALID_PASSWORD);
+      }
+      if (Util.exists(options.token)) {
+        Errors.checkArgumentValid(Util.isString(options.token),
+          ErrorCodes.ERR_CONN_CREATE_INVALID_OAUTH_TOKEN);
+      }
+    }
 
     consolidateHostAndAccount(options);
   }

diff --git a/lib/constants/error_messages.js b/lib/constants/error_messages.js
@@ -80,6 +80,8 @@ exports[404053] = 'A host must be specified.';
 exports[404054] = 'Invalid host. The specified value must be a string.';
 exports[404055] = 'Invalid passcodeInPassword. The specified value must be a boolean';
 exports[404056] = 'Invalid passcode. The specified value must be a string';
+exports[404057] = 'A password or token must be specified.';
+
 
 // 405001
 exports[405001] = 'Invalid callback. The specified value must be a function.';

diff --git a/lib/errors.js b/lib/errors.js
@@ -84,6 +84,8 @@ codes.ERR_CONN_CREATE_MISSING_HOST = 404053;
 codes.ERR_CONN_CREATE_INVALID_HOST = 404054;
 codes.ERR_CONN_CREATE_INVALID_PASSCODE_IN_PASSWORD = 404055;
 codes.ERR_CONN_CREATE_INVALID_PASSCODE = 404056;
+codes.ERR_CONN_CREATE_MISSING_PASSWORD_AND_TOKEN = 404057;
+
 
 // 405001
 codes.ERR_CONN_CONNECT_INVALID_CALLBACK = 405001;

diff --git a/lib/util.js b/lib/util.js
@@ -4,6 +4,7 @@ const os = require('os');
 const Logger = require('./logger');
 const fs = require('fs');
 const Errors = require('./errors');
+const net = require('net');
 
 /**
  * Note: A simple wrapper around util.inherits() for now, but this might change
@@ -765,6 +766,18 @@ exports.isWindows = function () {
   return os.platform() === 'win32';
 };
 
+
+exports.getFreePort = async function () {
+  return new Promise(res => {
+    const srv = net.createServer();
+    srv.listen(0, () => {
+      const port = srv.address().port;
+      srv.close(() => res(port));
+    });
+  });
+};
+
+
 /**
 * Left strip the specified character from a string.
 *

diff --git a/test/authentication/connectionParameters.js b/test/authentication/connectionParameters.js
@@ -78,6 +78,17 @@ const keypairEncryptedPrivateKeyPath =
       authenticator: 'SNOWFLAKE_JWT'
     };
 
+const oauthPATOnWiremock =
+  {
+    ...baseParameters,
+    accessUrl: null,
+    username: 'MOCK_USERNAME',
+    account: 'MOCK_ACCOUNT_NAME',
+    host: 'localhost',
+    protocol: 'http',
+    authenticator: 'PROGRAMMATIC_ACCESS_TOKEN',
+  };
+
 exports.externalBrowser = externalBrowser;
 exports.okta = okta;
 exports.oauth = oauth;
@@ -93,3 +104,4 @@ exports.snowflakeAuthTestOauthClientSecret = snowflakeAuthTestOauthClientSecret;
 exports.snowflakeAuthTestOauthUrl = snowflakeAuthTestOauthUrl;
 exports.snowflakeAuthTestPrivateKeyPath = snowflakeAuthTestPrivateKeyPath;
 exports.snowflakeAuthTestInvalidPrivateKeyPath = snowflakeAuthTestInvalidPrivateKeyPath;
+exports.oauthPATOnWiremock = oauthPATOnWiremock;
diff --git a/test/integration/wiremock/testOauthPat.js b/test/integration/wiremock/testOauthPat.js
@@ -0,0 +1,51 @@
+const net = require('net');
+const connParameters = require('../../authentication/connectionParameters');
+const AuthTest = require('../../authentication/authTestsBaseClass');
+const { runWireMockAsync, addWireMockMappingsFromFile, } = require('../../wiremockRunner');
+const os = require('os');
+const { getFreePort } = require('../../../lib/util');
+
+if (os.platform !== 'win32')  {
+  describe('Oauth PAT authentication', function () {
+    let port;
+    let authTest;
+    let wireMock;
+    before(async () => {
+      port = await getFreePort();
+      wireMock = await runWireMockAsync(port);
+    });
+    beforeEach(async () => {
+      authTest = new AuthTest();
+    });
+    afterEach(async () => {
+      wireMock.scenarios.resetAllScenarios();
+    });
+    after(async () => {
+      await wireMock.global.shutdown();
+    });
+
+    it('Successful flow scenario PAT as token', async function () {
+      await addWireMockMappingsFromFile(wireMock, 'wiremock/mappings/pat/successful_flow.json');
+      const connectionOption = { ...connParameters.oauthPATOnWiremock, token: 'MOCK_TOKEN', port: port };
+      authTest.createConnection(connectionOption);
+      await authTest.connectAsync();
+      authTest.verifyNoErrorWasThrown();
+    });
+
+    it('Successful flow scenario PAT as password', async function () {
+      await addWireMockMappingsFromFile(wireMock, 'wiremock/mappings/pat/successful_flow.json');
+      const connectionOption = { ...connParameters.oauthPATOnWiremock, password: 'MOCK_TOKEN', port: port };
+      authTest.createConnection(connectionOption);
+      await authTest.connectAsync();
+      authTest.verifyNoErrorWasThrown();
+    });
+
+    it('Invalid token', async function () {
+      await addWireMockMappingsFromFile(wireMock, 'wiremock/mappings/pat/invalid_pat_token.json');
+      const connectionOption = { ...connParameters.oauthPATOnWiremock, token: 'INVALID_TOKEN', port: port };
+      authTest.createConnection(connectionOption);
+      await authTest.connectAsync();
+      authTest.verifyErrorWasThrown('Programmatic access token is invalid.');
+    });
+  });
+}
diff --git a/test/wiremockRunner.js b/test/wiremockRunner.js
@@ -9,7 +9,7 @@ async function runWireMockAsync(port) {
   const waitingWireMockPromise =  new Promise( (resolve, reject) => {
     try {
       exec(`npx wiremock --enable-browser-proxying --proxy-pass-through  false --port ${port} `);
-      const wireMock = new WireMockRestClient(`http://localhost:${port}`, { logLevel: 'debug' });
+      const wireMock = new WireMockRestClient(`http://localhost:${port}`, { logLevel: 'trace' });
       const readyWireMock =  waitForWiremockStarted(wireMock);
       resolve(readyWireMock);
     } catch (err) {

diff --git a/wiremock/mappings/pat/invalid_pat_token.json b/wiremock/mappings/pat/invalid_pat_token.json
@@ -0,0 +1,48 @@
+{
+  "mappings": [
+    {
+      "scenarioName": "Invalid PAT authentication flow",
+      "requiredScenarioState": "Started",
+      "newScenarioState": "Authenticated",
+      "request": {
+        "urlPathPattern": "/session/v1/login-request.*",
+        "method": "POST",
+        "headers": {
+          "accept": {
+            "equalTo": "application/json"
+          }
+        },
+        "bodyPatterns": [
+          {
+            "equalToJson": {
+              "data": {
+                "ACCOUNT_NAME": "MOCK_ACCOUNT_NAME",
+                "TOKEN": "INVALID_TOKEN",
+                "LOGIN_NAME": "MOCK_USERNAME",
+                "AUTHENTICATOR": "PROGRAMMATIC_ACCESS_TOKEN"
+              }
+            },
+            "ignoreExtraElements": true
+          }
+        ]
+      },
+      "response": {
+        "status": 200,
+        "headers": {
+          "Content-Type": "application/json"
+        },
+        "jsonBody": {
+          "data": {
+            "nextAction": "RETRY_LOGIN",
+            "authnMethod": "PAT",
+            "signInOptions": {}
+          },
+          "code": "394400",
+          "message": "Programmatic access token is invalid.",
+          "success": false,
+          "headers": null
+        }
+      }
+    }
+  ]
+}
diff --git a/wiremock/mappings/pat/successful_flow.json b/wiremock/mappings/pat/successful_flow.json
@@ -0,0 +1,74 @@
+{
+  "mappings": [
+    {
+      "scenarioName": "Successful PAT authentication flow",
+      "requiredScenarioState": "Started",
+      "newScenarioState": "Authenticated",
+      "request": {
+        "urlPathPattern": "/session/v1/login-request.*",
+        "method": "POST",
+        "headers": {
+          "accept": {
+            "equalTo": "application/json"
+          }
+        },
+        "bodyPatterns": [
+          {
+            "equalToJson": {
+              "data": {
+                "ACCOUNT_NAME": "MOCK_ACCOUNT_NAME",
+                "TOKEN": "MOCK_TOKEN",
+                "LOGIN_NAME": "MOCK_USERNAME",
+                "AUTHENTICATOR": "PROGRAMMATIC_ACCESS_TOKEN"
+              }
+            },
+            "ignoreExtraElements": true
+          }
+        ]
+      },
+      "response": {
+        "status": 200,
+        "headers": {
+          "Content-Type": "application/json"
+        },
+        "jsonBody": {
+          "data": {
+            "masterToken": "master token",
+            "token": "session token",
+            "validityInSeconds": 3600,
+            "masterValidityInSeconds": 14400,
+            "displayUserName": "OAUTH_TEST_AUTH_CODE",
+            "serverVersion": "8.48.0 b2024121104444034239f05",
+            "firstLogin": false,
+            "remMeToken": null,
+            "remMeValidityInSeconds": 0,
+            "healthCheckInterval": 45,
+            "newClientForUpgrade": "3.12.3",
+            "sessionId": 1172562260498,
+            "parameters": [
+              {
+                "name": "CLIENT_PREFETCH_THREADS",
+                "value": 4
+              }
+            ],
+            "sessionInfo": {
+              "databaseName": "TEST_DHEYMAN",
+              "schemaName": "TEST_JDBC",
+              "warehouseName": "TEST_XSMALL",
+              "roleName": "ANALYST"
+            },
+            "idToken": null,
+            "idTokenValidityInSeconds": 0,
+            "responseData": null,
+            "mfaToken": null,
+            "mfaTokenValidityInSeconds": 0
+          },
+          "code": null,
+          "message": null,
+          "success": true
+        }
+      }
+    }
+  ]
+}
+