[PECO-728] Add OAuth support

Signed-off-by: Levko Kravets <[email protected]>

Author: kravets-levko

lib/connection/auth/DatabricksOAuth/AuthorizationCode.ts

@@ -0,0 +1,145 @@
+import http, { Server, IncomingMessage, ServerResponse } from 'http';
+import { BaseClient, generators } from 'openid-client';
+import open from 'open';
+import IDBSQLLogger, { LogLevel } from '../../../contracts/IDBSQLLogger';
+
+export interface AuthorizationCodeOptions {
+  client: BaseClient;
+  ports: Array<number>;
+  logger?: IDBSQLLogger;
+}
+
+const scopeDelimiter = ' ';
+
+async function startServer(
+  host: string,
+  port: number,
+  requestHandler: (req: IncomingMessage, res: ServerResponse) => void,
+): Promise<Server> {
+  const server = http.createServer(requestHandler);
+
+  return new Promise((resolve, reject) => {
+    const errorListener = (error: Error) => {
+      server.off('error', errorListener);
+      reject(error);
+    };
+
+    server.on('error', errorListener);
+    server.listen(port, host, () => {
+      server.off('error', errorListener);
+      resolve(server);
+    });
+  });
+}
+
+async function stopServer(server: Server): Promise<void> {
+  if (!server.listening) {
+    return;
+  }
+
+  return new Promise((resolve, reject) => {
+    const errorListener = (error: Error) => {
+      server.off('error', errorListener);
+      reject(error);
+    };
+
+    server.on('error', errorListener);
+    server.close(() => {
+      server.off('error', errorListener);
+      resolve();
+    });
+  });
+}
+
+export interface AuthorizationCodeFetchResult {
+  code: string;
+  verifier: string;
+  redirectUri: string;
+}
+
+export default class AuthorizationCode {
+  private readonly client: BaseClient;
+  private readonly host: string = 'localhost';
+  private readonly ports: Array<number>;
+  private readonly logger?: IDBSQLLogger;
+
+  constructor(options: AuthorizationCodeOptions) {
+    this.client = options.client;
+    this.ports = options.ports;
+    this.logger = options.logger;
+  }
+
+  public async fetch(scopes: Array<string>): Promise<AuthorizationCodeFetchResult> {
+    const verifierString = generators.codeVerifier(32);
+    const challengeString = generators.codeChallenge(verifierString);
+    const state = generators.state(16);
+
+    let code: string | undefined = undefined;
+
+    const server = await this.startServer((req, res) => {
+      const params = this.client.callbackParams(req);
+      if (params.state === state) {
+        code = params.code;
+        res.writeHead(200);
+        res.end('You can close this tab');
+        server.stop();
+      } else {
+        res.writeHead(404);
+        res.end();
+      }
+    });
+
+    let redirectUri = `http://${server.host}:${server.port}/`;
+    const authUrl = this.client.authorizationUrl({
+      response_type: 'code',
+      response_mode: 'query',
+      scope: scopes.join(scopeDelimiter),
+      code_challenge: challengeString,
+      code_challenge_method: 'S256',
+      state,
+      redirect_uri: redirectUri,
+    });
+
+    await open(authUrl);
+    await server.stopped();
+
+    if (!code) {
+      throw new Error(`No path parameters were returned to the callback at ${redirectUri}`);
+    }
+
+    return { code, verifier: verifierString, redirectUri };
+  }
+
+  private async startServer(requestHandler: (req: IncomingMessage, res: ServerResponse) => void) {
+    for (const port of this.ports) {
+      const host = this.host;
+      try {
+        const server = await startServer(host, port, requestHandler);
+        this.logger?.log(LogLevel.info, `Listening for OAuth authorization callback at ${host}:${port}`);
+
+        let resolveStopped = () => {};
+        let rejectStopped = (reason?: any) => {};
+        const stoppedPromise = new Promise<void>((resolve, reject) => {
+          resolveStopped = resolve;
+          rejectStopped = reject;
+        });
+
+        return {
+          host,
+          port,
+          server,
+          stop: () => stopServer(server).then(resolveStopped).catch(rejectStopped),
+          stopped: () => stoppedPromise,
+        };
+      } catch (error) {
+        if (error instanceof Error && 'code' in error && error.code === 'EADDRINUSE') {
+          this.logger?.log(LogLevel.debug, `Failed to start server at ${host}:${port}: ${error.code}`);
+          continue; // try another port
+        }
+        throw error;
+      }
+    }
+
+    throw new Error('Failed to start server: all ports are in use');
+  }
+}

lib/connection/auth/DatabricksOAuth/OAuthManager.ts

@@ -0,0 +1,97 @@
+import { Issuer, BaseClient } from 'openid-client';
+import HiveDriverError from '../../../errors/HiveDriverError';
+import IDBSQLLogger, { LogLevel } from '../../../contracts/IDBSQLLogger';
+import OAuthToken from "./OAuthToken";
+import AuthorizationCode from "./AuthorizationCode";
+
+const oidcConfigPath = 'oidc/.well-known/oauth-authorization-server';
+
+export interface OAuthManagerOptions {
+  host: string;
+  callbackPorts: Array<number>;
+  clientId: string;
+  logger?: IDBSQLLogger;
+}
+
+export default class OAuthManager {
+  private readonly options: OAuthManagerOptions;
+  private readonly logger?: IDBSQLLogger;
+
+  private issuer?: Issuer;
+  private client?: BaseClient;
+
+  constructor(options: OAuthManagerOptions) {
+    this.options = options;
+    this.logger = options.logger;
+  }
+
+  private async getClient(): Promise<BaseClient> {
+    if (!this.issuer) {
+      const { host } = this.options;
+      const schema = host.startsWith('https://') ? '' : 'https://';
+      const trailingSlash = host.endsWith('/') ? '' : '/';
+      this.issuer = await Issuer.discover(`${schema}${host}${trailingSlash}${oidcConfigPath}`);
+    }
+
+    if (!this.client) {
+      this.client = new this.issuer.Client({
+        client_id: this.options.clientId,
+        token_endpoint_auth_method: 'none',
+      });
+    }
+
+    return this.client;
+  }
+
+  public async refreshAccessToken(token: OAuthToken): Promise<OAuthToken> {
+    try {
+      if (!token.hasExpired) {
+        // The access token is fine. Just return it.
+        return token;
+      }
+    } catch (error) {
+      this.logger?.log(LogLevel.error, `${error}`);
+      throw error;
+    }
+
+    if (!token.refreshToken) {
+      const message = `OAuth access token expired on ${token.expirationTime}.`;
+      this.logger?.log(LogLevel.error, message);
+      throw new HiveDriverError(message);
+    }
+
+    // Try to refresh using the refresh token
+    this.logger?.log(LogLevel.debug, `Attempting to refresh OAuth access token that expired on ${token.expirationTime}`);
+
+    const client = await this.getClient();
+    const { access_token, refresh_token } = await client.refresh(token.refreshToken);
+    if (!access_token || !refresh_token) {
+      throw new Error('Failed to refresh token: invalid response');
+    }
+    return new OAuthToken(access_token, refresh_token);
+  }
+
+  public async getToken(scopes: Array<string>): Promise<OAuthToken> {
+    const client = await this.getClient();
+    const authCode = new AuthorizationCode({
+      client,
+      ports: this.options.callbackPorts,
+      logger: this.logger,
+    });
+
+    const { code, verifier, redirectUri } = await authCode.fetch(scopes);
+
+    const { access_token, refresh_token } = await client.grant({
+      grant_type: 'authorization_code',
+      code,
+      code_verifier: verifier,
+      redirect_uri: redirectUri,
+    });
+
+    if (!access_token) {
+      throw new Error('Failed to fetch access token');
+    }
+
+    return new OAuthToken(access_token, refresh_token);
+  }
+}

lib/connection/auth/DatabricksOAuth/OAuthPersistence.ts

@@ -0,0 +1,10 @@
+import OAuthToken from "./OAuthToken";
+
+export default class OAuthPersistence {
+  public async persist(host: string, token: OAuthToken): Promise<void> {
+  }
+
+  public async read(host: string): Promise<OAuthToken | undefined> {
+    return undefined;
+  }
+}

lib/connection/auth/DatabricksOAuth/OAuthToken.ts

@@ -0,0 +1,36 @@
+export default class OAuthToken {
+  private readonly _accessToken: string;
+  private readonly _refreshToken?: string;
+  private _expirationTime?: number;
+
+  constructor(accessToken: string, refreshToken?: string) {
+    this._accessToken = accessToken;
+    this._refreshToken = refreshToken;
+  }
+
+  get accessToken(): string {
+    return this._accessToken;
+  }
+
+  get refreshToken(): string | undefined {
+    return this._refreshToken;
+  }
+
+  get expirationTime(): number {
+    if (this._expirationTime === undefined) {
+      const accessTokenPayload = Buffer.from(this._accessToken.split('.')[1], 'base64').toString('utf8');
+      const decoded = JSON.parse(accessTokenPayload);
+      this._expirationTime = Number(decoded['exp']);
+    }
+    return this._expirationTime;
+  }
+
+  get hasExpired(): boolean {
+    // This token has already been verified, and we are just parsing it.
+    // If it has been tampered with, it will be rejected on the server side.
+    // This avoids having to fetch the public key from the issuer and perform
+    // an unnecessary signature verification.
+    const now = Math.floor(Date.now() / 1000); // convert it to seconds
+    return this.expirationTime <= now;
+  }
+}

lib/connection/auth/DatabricksOAuth/index.ts

@@ -0,0 +1,68 @@
+import IAuthentication from '../../contracts/IAuthentication';
+import ITransport from '../../contracts/ITransport';
+import IDBSQLLogger from '../../../contracts/IDBSQLLogger';
+import { AuthOptions } from '../../types/AuthOptions';
+import OAuthPersistence from './OAuthPersistence';
+import OAuthManager from "./OAuthManager";
+
+interface DatabricksOAuthOptions extends AuthOptions {
+  host: string;
+  redirectPorts?: Array<number>;
+  clientId?: string;
+  scopes?: Array<string>;
+  logger?: IDBSQLLogger;
+  persistence?: OAuthPersistence;
+  headers?: object;
+}
+
+const defaultOAuthOptions = {
+  clientId: 'databricks-sql-python',
+  redirectPorts: [8020, 8021, 8022, 8023, 8024, 8025],
+  scopes: ['sql', 'offline_access'],
+} satisfies Partial<DatabricksOAuthOptions>;
+
+export default class DatabricksOAuth implements IAuthentication {
+  private readonly host: string;
+  private readonly redirectPorts: Array<number>;
+  private readonly clientId: string;
+  private readonly scopes: Array<string>;
+  private readonly logger?: IDBSQLLogger;
+  private readonly persistence?: OAuthPersistence;
+  private readonly headers?: object;
+
+  private readonly manager: OAuthManager;
+
+  constructor(options: DatabricksOAuthOptions) {
+    this.host = options.host;
+    this.redirectPorts = options.redirectPorts || defaultOAuthOptions.redirectPorts;
+    this.clientId = options.clientId || defaultOAuthOptions.clientId;
+    this.scopes = options.scopes || defaultOAuthOptions.scopes;
+    this.logger = options.logger;
+    this.persistence = options.persistence;
+    this.headers = options.headers;
+
+    this.manager = new OAuthManager({
+      host: this.host,
+      callbackPorts: this.redirectPorts,
+      clientId: this.clientId,
+      logger: this.logger,
+    });
+  }
+
+  async authenticate(transport: ITransport): Promise<ITransport> {
+    let token = await this.persistence?.read(this.host);
+    if (!token) {
+      token = await this.manager.getToken(this.scopes);
+    }
+
+    token = await this.manager.refreshAccessToken(token);
+    await this.persistence?.persist(this.host, token);
+
+    transport.setOptions('headers', {
+      ...this.headers,
+      Authorization: `Bearer ${token.accessToken}`,
+    });
+
+    return transport;
+  }
+}