fix authentication mechanism

Author: ptiurin

src/auth/index.ts

@@ -6,15 +6,25 @@ import {
   ServiceAccountAuth,
   UsernamePasswordAuth
 } from "../types";
-import { TokenKey, inMemoryCache, noneCache } from "../common/tokenCache";
-import ReadWriteLock from "rwlock";
+import {
+  TokenKey,
+  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 +33,7 @@ export class Authenticator {
   options: ConnectionOptions;
 
   accessToken?: string;
-  private readonly rwlock = new ReadWriteLock();
+  tokenExpirationTime?: number;
 
   constructor(context: Context, options: ConnectionOptions) {
     context.httpClient.authenticator = this;
@@ -62,23 +72,40 @@ 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 expirationTimeMs = Date.now() + (expiresIn * 1000) / 2;
     this.accessToken = token;
+    this.tokenExpirationTime = expirationTimeMs;
+    // Update cache
     const key = this.getCacheKey();
-    key && this.getCache().set(key, { token, expiration: expiresIn });
+    key &&
+      this.getCache().set(key, { token, tokenExpiryTimeMs: expirationTimeMs });
   }
 
-  private getCachedToken(): string | undefined {
+  private getCachedTokenInfo():
+    | { token: string; tokenExpiryTimeMs: number }
+    | undefined {
     const key = this.getCacheKey();
-    return key ? this.getCache().get(key)?.token : undefined;
-  }
+    if (!key) return undefined;
 
-  getHeaders() {
-    if (this.accessToken) {
+    const cachedTokenInfo = this.getCache().get(key);
+    // Check if token exists and is not expired
+    if (cachedTokenInfo && Date.now() < cachedTokenInfo.tokenExpiryTimeMs) {
       return {
-        Authorization: `Bearer ${this.accessToken}`
+        token: cachedTokenInfo.token,
+        tokenExpiryTimeMs: cachedTokenInfo.tokenExpiryTimeMs
       };
     }
-    return {};
+
+    return undefined;
+  }
+
+  async getToken(): Promise<string | undefined> {
+    if (this.tokenExpirationTime && this.tokenExpirationTime < Date.now()) {
+      await this.authenticate();
+    }
+    return this.accessToken;
   }
 
   private static getAuthEndpoint(apiEndpoint: string) {
@@ -94,7 +121,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 +132,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,
+        auth: false
       })
       .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 +159,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,
+        auth: false
       })
       .ready();
 
-    this.setToken(access_token, expires_in);
+    return { access_token, expires_in };
   }
 
   isUsernamePassword() {
@@ -163,19 +194,21 @@ export class Authenticator {
     // Try to get token from cache using read lock
     const cachedToken = await this.tryGetCachedToken();
     if (cachedToken) {
-      this.accessToken = cachedToken;
+      this.accessToken = cachedToken.token;
+      this.tokenExpirationTime = cachedToken.tokenExpiryTimeMs;
       return;
     }
-
     // No cached token, acquire write lock and authenticate
     await this.acquireWriteLockAndAuthenticate();
   }
 
-  private async tryGetCachedToken(): Promise<string | undefined> {
+  private async tryGetCachedToken(): Promise<
+    { token: string; tokenExpiryTimeMs: number } | undefined
+  > {
     return new Promise((resolve, reject) => {
-      this.rwlock.readLock(releaseReadLock => {
+      rwLock.readLock(releaseReadLock => {
         try {
-          const cachedToken = this.getCachedToken();
+          const cachedToken = this.getCachedTokenInfo();
           resolve(cachedToken);
         } catch (error) {
           reject(error instanceof Error ? error : new Error(String(error)));
@@ -188,15 +221,15 @@ 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.getCachedTokenInfo();
+          if (cachedTokenInfo) {
+            this.accessToken = cachedTokenInfo.token;
+            this.tokenExpirationTime = cachedTokenInfo.tokenExpiryTimeMs;
             return resolve();
           }
-
           await this.performAuthentication();
 
           resolve();
@@ -212,14 +245,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;
+  tokenExpiryTimeMs: 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.tokenExpiryTimeMs
+    );
   }
 
   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();

src/http/node.ts

@@ -14,6 +14,7 @@ type RequestOptions = {
   body?: string;
   raw?: boolean;
   retry?: boolean;
+  auth?: boolean;
 };
 
 type ErrorResponse = {
@@ -84,8 +85,17 @@ export class NodeHttpClient {
     };
 
     const makeRequest = async () => {
-      if (this.authenticator) {
-        const authHeaders = await this.authenticator.getHeaders();
+      // No authentication for auth request
+      if (options?.auth !== false) {
+        const token = await this.authenticator.getToken();
+        if (!token) {
+          throw new AuthenticationError({
+            message: "Failed to get the access token when making a request."
+          });
+        }
+        const authHeaders = {
+          Authorization: `Bearer ${token}`
+        };
         Object.assign(headers, authHeaders);
       }
 
@@ -107,6 +117,9 @@ export class NodeHttpClient {
 
       if (response.status === 401 && retry) {
         try {
+          console.warn(
+            "Access token expired (401), refreshing access token and retrying request"
+          );
           this.authenticator.clearCache();
           await this.authenticator.authenticate();
         } catch (error) {

test/unit/http.test.ts

@@ -31,6 +31,9 @@ const authHandlerV1 = rest.post(
   }
 );
 
+// Increase Jest timeout for tests that might involve timeouts
+jest.setTimeout(10000);
+
 describe.each([
   ["v1", authHandlerV1, { username: "user", password: "fake_password" }],
   ["v2", authHandlerV2, { client_id: "user", client_secret: "fake_password" }]