fix(FIR-46131): Atomic authentication mechanism and token expiry (#142)

Author: ptiurin

src/auth/index.ts

@@ -6,15 +6,26 @@ import {
   ServiceAccountAuth,
   UsernamePasswordAuth
 } from "../types";
-import { TokenKey, inMemoryCache, noneCache } from "../common/tokenCache";
-import ReadWriteLock from "rwlock";
+import {
+  TokenKey,
+  TokenRecord,
+  inMemoryCache,
+  noneCache,
+  rwLock
+} from "../common/tokenCache";
 
 type Login = {
   access_token: string;
   token_type: string;
   expires_in: number;
 };
 
+type TokenInfo = {
+  access_token: string;
+  // seconds until expiration
+  expires_in: number;
+};
+
 const AUTH_AUDIENCE = "https://api.firebolt.io";
 const AUTH_GRANT_TYPE = "client_credentials";
 
@@ -23,7 +34,9 @@ export class Authenticator {
   options: ConnectionOptions;
 
   accessToken?: string;
-  private readonly rwlock = new ReadWriteLock();
+  // Expiration time is half way to the actual expiration time
+  // to allow for some buffer time before the token expires
+  tokenExpiryTimestampMs?: number;
 
   constructor(context: Context, options: ConnectionOptions) {
     context.httpClient.authenticator = this;
@@ -62,23 +75,48 @@ export class Authenticator {
   }
 
   private setToken(token: string, expiresIn: number) {
+    // Set expiration to half of the expiresIn value
+    // to allow for some buffer time before the token expires
+    const tokenExpiryTimestampMs = Date.now() + (expiresIn * 1000) / 2;
     this.accessToken = token;
+    this.tokenExpiryTimestampMs = tokenExpiryTimestampMs;
+    // Update cache
     const key = this.getCacheKey();
-    key && this.getCache().set(key, { token, expiration: expiresIn });
+    key &&
+      this.getCache().set(key, {
+        token,
+        tokenExpiryTimestampMs
+      });
   }
 
-  private getCachedToken(): string | undefined {
+  private getCachedTokenRecord(): TokenRecord | undefined {
     const key = this.getCacheKey();
-    return key ? this.getCache().get(key)?.token : undefined;
-  }
-
-  getHeaders() {
-    if (this.accessToken) {
+    if (!key) return undefined;
+
+    const cachedTokenRecord = this.getCache().get(key);
+    // Check if token exists and is not expired
+    if (
+      cachedTokenRecord &&
+      Date.now() < cachedTokenRecord.tokenExpiryTimestampMs
+    ) {
       return {
-        Authorization: `Bearer ${this.accessToken}`
+        token: cachedTokenRecord.token,
+        tokenExpiryTimestampMs: cachedTokenRecord.tokenExpiryTimestampMs
       };
     }
-    return {};
+
+    return undefined;
+  }
+
+  async getToken(): Promise<string | undefined> {
+    if (
+      (this.tokenExpiryTimestampMs &&
+        this.tokenExpiryTimestampMs < Date.now()) ||
+      !this.accessToken
+    ) {
+      await this.authenticate();
+    }
+    return this.accessToken;
   }
 
   private static getAuthEndpoint(apiEndpoint: string) {
@@ -94,7 +132,9 @@ export class Authenticator {
     return myURL.toString();
   }
 
-  private async authenticateUsernamePassword(auth: UsernamePasswordAuth) {
+  private async authenticateUsernamePassword(
+    auth: UsernamePasswordAuth
+  ): Promise<TokenInfo> {
     const { httpClient, apiEndpoint } = this.context;
     const { username, password } = auth;
     const url = `${apiEndpoint}/${USERNAME_PASSWORD_LOGIN}`;
@@ -103,19 +143,21 @@ export class Authenticator {
       password
     });
 
-    this.accessToken = undefined;
-
+    // Expiration is in seconds
     const { access_token, expires_in } = await httpClient
       .request<Login>("POST", url, {
         body,
-        retry: false
+        retry: false,
+        noAuth: true
       })
       .ready();
 
-    this.setToken(access_token, expires_in);
+    return { access_token, expires_in };
   }
 
-  private async authenticateServiceAccount(auth: ServiceAccountAuth) {
+  private async authenticateServiceAccount(
+    auth: ServiceAccountAuth
+  ): Promise<TokenInfo> {
     const { httpClient, apiEndpoint } = this.context;
     const { client_id, client_secret } = auth;
 
@@ -128,19 +170,19 @@ export class Authenticator {
     });
     const url = `${authEndpoint}${SERVICE_ACCOUNT_LOGIN}`;
 
-    this.accessToken = undefined;
-
+    // Expiration is in seconds
     const { access_token, expires_in } = await httpClient
       .request<Login>("POST", url, {
         retry: false,
         headers: {
           "Content-Type": "application/x-www-form-urlencoded"
         },
-        body: params
+        body: params,
+        noAuth: true
       })
       .ready();
 
-    this.setToken(access_token, expires_in);
+    return { access_token, expires_in };
   }
 
   isUsernamePassword() {
@@ -161,22 +203,45 @@ export class Authenticator {
 
   async authenticate(): Promise<void> {
     // Try to get token from cache using read lock
-    const cachedToken = await this.tryGetCachedToken();
+    const cachedToken = await this.tryGetCachedTokenRecord();
     if (cachedToken) {
-      this.accessToken = cachedToken;
+      this.accessToken = cachedToken.token;
+      this.tokenExpiryTimestampMs = cachedToken.tokenExpiryTimestampMs;
       return;
     }
-
     // No cached token, acquire write lock and authenticate
     await this.acquireWriteLockAndAuthenticate();
   }
 
-  private async tryGetCachedToken(): Promise<string | undefined> {
+  async reAuthenticate(): Promise<void> {
+    // Acquire write lock, clear cache and authenticate
     return new Promise((resolve, reject) => {
-      this.rwlock.readLock(releaseReadLock => {
+      rwLock.writeLock(async releaseWriteLock => {
         try {
-          const cachedToken = this.getCachedToken();
-          resolve(cachedToken);
+          // Clear the cache under write lock
+          const key = this.getCacheKey();
+          key && this.getCache().clear(key);
+
+          // Perform authentication directly rather than calling acquireWriteLockAndAuthenticate
+          // since we already have the write lock
+          await this.performAuthentication();
+
+          resolve();
+        } catch (error) {
+          reject(error instanceof Error ? error : new Error(String(error)));
+        } finally {
+          releaseWriteLock();
+        }
+      });
+    });
+  }
+
+  private async tryGetCachedTokenRecord(): Promise<TokenRecord | undefined> {
+    return new Promise((resolve, reject) => {
+      rwLock.readLock(releaseReadLock => {
+        try {
+          const cachedTokenRecord = this.getCachedTokenRecord();
+          resolve(cachedTokenRecord);
         } catch (error) {
           reject(error instanceof Error ? error : new Error(String(error)));
         } finally {
@@ -188,15 +253,16 @@ export class Authenticator {
 
   private async acquireWriteLockAndAuthenticate(): Promise<void> {
     return new Promise((resolve, reject) => {
-      this.rwlock.writeLock(async releaseWriteLock => {
+      rwLock.writeLock(async releaseWriteLock => {
         try {
           // Double-check cache in case another thread authenticated while waiting
-          const cachedToken = this.getCachedToken();
-          if (cachedToken) {
-            this.accessToken = cachedToken;
+          const cachedTokenInfo = this.getCachedTokenRecord();
+          if (cachedTokenInfo) {
+            this.accessToken = cachedTokenInfo.token;
+            this.tokenExpiryTimestampMs =
+              cachedTokenInfo.tokenExpiryTimestampMs;
             return resolve();
           }
-
           await this.performAuthentication();
 
           resolve();
@@ -212,14 +278,20 @@ export class Authenticator {
   private async performAuthentication(): Promise<void> {
     const options = this.options.auth || this.options;
 
-    if (this.isUsernamePassword()) {
-      return this.authenticateUsernamePassword(options as UsernamePasswordAuth);
-    }
+    let auth: TokenInfo;
 
-    if (this.isServiceAccount()) {
-      return this.authenticateServiceAccount(options as ServiceAccountAuth);
+    if (this.isUsernamePassword()) {
+      auth = await this.authenticateUsernamePassword(
+        options as UsernamePasswordAuth
+      );
+    } else if (this.isServiceAccount()) {
+      auth = await this.authenticateServiceAccount(
+        options as ServiceAccountAuth
+      );
+    } else {
+      throw new Error("Please provide valid auth credentials");
     }
 
-    throw new Error("Please provide valid auth credentials");
+    this.setToken(auth.access_token, auth.expires_in);
   }
 }

src/common/tokenCache.ts

@@ -1,3 +1,4 @@
+import ReadWriteLock from "rwlock";
 import { AccountInfo } from "../connection/connection_v1";
 
 export type TokenKey = {
@@ -8,7 +9,7 @@ export type TokenKey = {
 
 export type TokenRecord = {
   token: string;
-  expiration: number;
+  tokenExpiryTimestampMs: number;
 };
 
 export type AccountKey = {
@@ -62,7 +63,6 @@ export class InMemoryCacheStorage<KeyType, RecordType> {
     const lookup = this.makeLookupString(key);
     const record = this.storage[lookup];
     if (!this.isValidRecord(record)) {
-      this.clear(key);
       return null;
     }
     return record;
@@ -84,11 +84,12 @@ export class InMemoryTokenCacheStorage extends InMemoryCacheStorage<
   TokenRecord
 > {
   protected isValidRecord(record: TokenRecord | undefined): boolean {
-    return typeof record != "undefined" && Date.now() < record.expiration;
+    return (
+      typeof record != "undefined" && Date.now() < record.tokenExpiryTimestampMs
+    );
   }
 
   protected modifyRecord(record: TokenRecord): TokenRecord {
-    record.expiration = Date.now() + record.expiration;
     return record;
   }
 }
@@ -116,3 +117,4 @@ export class InMemoryCache implements Cache {
 
 export const inMemoryCache = new InMemoryCache();
 export const noneCache = new NoneCache();
+export const rwLock = new ReadWriteLock();