feat: enhance static credentials authentication with token refresh interval and improved token handling

Author: polRk

src/credentials/static-credentials-auth-service.ts

@@ -1,30 +1,54 @@
-import {Ydb} from "ydb-sdk-proto";
+import { Ydb } from "ydb-sdk-proto";
 import AuthServiceResult = Ydb.Auth.LoginResult;
-import {ISslCredentials} from "../utils/ssl-credentials";
-import {GrpcService, withTimeout} from "../utils";
-import {retryable} from "../retries_obsoleted";
-import {DateTime} from "luxon";
-import {getOperationPayload} from "../utils/process-ydb-operation-result";
+import { ISslCredentials } from "../utils/ssl-credentials";
+import { GrpcService, withTimeout } from "../utils";
+import { retryable } from "../retries_obsoleted";
+import { DateTime } from "luxon";
+import { getOperationPayload } from "../utils/process-ydb-operation-result";
 import * as grpc from "@grpc/grpc-js";
-import {addCredentialsToMetadata} from "./add-credentials-to-metadata";
+import { addCredentialsToMetadata } from "./add-credentials-to-metadata";
 
-import {IAuthService} from "./i-auth-service";
-import {HasLogger} from "../logger/has-logger";
-import {Logger} from "../logger/simple-logger";
-import {getDefaultLogger} from "../logger/get-default-logger";
+import { IAuthService } from "./i-auth-service";
+import { HasLogger } from "../logger/has-logger";
+import { Logger } from "../logger/simple-logger";
+import { getDefaultLogger } from "../logger/get-default-logger";
 
+/**
+ * Static credentials token.
+ */
+export type StaticCredentialsToken = {
+    value: string
+    aud: string[]
+    exp: number
+    iat: number
+    sub: string
+}
+
+/**
+ * Interface for options used in static credentials authentication.
+ */
 interface StaticCredentialsAuthOptions {
-    /** Custom ssl sertificates. If you use it in driver, you must use it here too */
+    /** Custom SSL certificates. If you use it in driver, you must use it here too */
     sslCredentials?: ISslCredentials;
+
     /**
      * Timeout for token request in milliseconds
      * @default 10 * 1000
      */
     tokenRequestTimeout?: number;
-    /** Expiration time for token in milliseconds
+
+    /**
+     * Expiration time for token in milliseconds
+     * @deprecated Use tokenRefreshInterval instead
      * @default 6 * 60 * 60 * 1000
      */
-    tokenExpirationTimeout?: number
+    tokenExpirationTimeout?: number;
+
+    /**
+     * Time interval in milliseconds after which the token will be refreshed.
+     * When specified, token refresh is based on this timer rather than the token's exp field.
+     */
+    tokenRefreshInterval?: number;
 }
 
 class StaticCredentialsGrpcService extends GrpcService<Ydb.Auth.V1.AuthService> implements HasLogger {
@@ -44,16 +68,19 @@ class StaticCredentialsGrpcService extends GrpcService<Ydb.Auth.V1.AuthService>
 
 export class StaticCredentialsAuthService implements IAuthService {
     private readonly tokenRequestTimeout = 10 * 1000;
-    private readonly tokenExpirationTimeout = 6 * 60 * 60 * 1000;
-    private tokenTimestamp: DateTime | null;
-    private token: string = '';
-    private tokenUpdatePromise: Promise<any> | null = null;
-    private user: string;
-    private password: string;
-    private endpoint: string;
-    private sslCredentials: ISslCredentials | undefined;
+    private readonly tokenRefreshInterval: number | null = null;
+
+    private readonly user: string;
+    private readonly password: string;
+    private readonly endpoint: string;
+    private readonly sslCredentials: ISslCredentials | undefined;
+
     public readonly logger: Logger;
 
+    private token: StaticCredentialsToken | null = null;
+    // Mutex
+    private promise: Promise<grpc.Metadata> | null = null;
+
     constructor(
         user: string,
         password: string,
@@ -74,25 +101,37 @@ export class StaticCredentialsAuthService implements IAuthService {
         loggerOrOptions?: Logger | StaticCredentialsAuthOptions,
         options?: StaticCredentialsAuthOptions
     ) {
-        this.tokenTimestamp = null;
         this.user = user;
         this.password = password;
         this.endpoint = endpoint;
         this.sslCredentials = options?.sslCredentials;
+
         if (typeof loggerOrOptions === 'object' && loggerOrOptions !== null && 'error' in loggerOrOptions) {
             this.logger = loggerOrOptions as Logger;
         } else {
             options = loggerOrOptions;
             this.logger = getDefaultLogger();
         }
+
         if (options?.tokenRequestTimeout) this.tokenRequestTimeout = options.tokenRequestTimeout;
-        if (options?.tokenExpirationTimeout) this.tokenExpirationTimeout = options.tokenExpirationTimeout;
-    }
+        if (options?.tokenExpirationTimeout) this.tokenRefreshInterval = options.tokenExpirationTimeout;
+        if (options?.tokenRefreshInterval) this.tokenRefreshInterval = options.tokenRefreshInterval;
 
-    private get expired() {
-        return !this.tokenTimestamp || (
-            DateTime.utc().diff(this.tokenTimestamp).valueOf() > this.tokenExpirationTimeout
-        );
+        if (this.tokenRefreshInterval) {
+            let timer = setInterval(() => {
+                if (this.promise) {
+                    return
+                }
+
+                this.promise = this.updateToken()
+                    .then(token => addCredentialsToMetadata(token.value))
+                    .finally(() => {
+                        this.promise = null;
+                    })
+            }, this.tokenRefreshInterval);
+
+            timer.unref()
+        }
     }
 
     private async sendTokenRequest(): Promise<AuthServiceResult> {
@@ -101,34 +140,52 @@ export class StaticCredentialsAuthService implements IAuthService {
             this.sslCredentials,
             this.logger,
         );
+
         const tokenPromise = runtimeAuthService.login({
             user: this.user,
             password: this.password,
         });
+
         const response = await withTimeout(tokenPromise, this.tokenRequestTimeout);
         const result = AuthServiceResult.decode(getOperationPayload(response));
         runtimeAuthService.destroy();
+
         return result;
     }
 
-    private async updateToken() {
-        const {token} = await this.sendTokenRequest();
-        if (token) {
-            this.token = token;
-            this.tokenTimestamp = DateTime.utc();
-        } else {
+    private async updateToken(): Promise<StaticCredentialsToken> {
+        const { token } = await this.sendTokenRequest();
+        if (!token) {
             throw new Error('Received empty token from static credentials!');
         }
+
+        // Parse the JWT token to extract expiration time
+        const [, payload] = token.split('.');
+        const decodedPayload = JSON.parse(Buffer.from(payload, 'base64').toString());
+
+        this.token = {
+            value: token,
+            ...decodedPayload
+        };
+
+        return this.token!
     }
 
     public async getAuthMetadata(): Promise<grpc.Metadata> {
-        if (this.expired || this.tokenUpdatePromise) {
-            if (!this.tokenUpdatePromise) {
-                this.tokenUpdatePromise = this.updateToken();
-            }
-            await this.tokenUpdatePromise;
-            this.tokenUpdatePromise = null;
+        if (this.token && this.token.exp > Date.now() / 1000) {
+            return addCredentialsToMetadata(this.token.value)
         }
-        return addCredentialsToMetadata(this.token);
+
+        if (this.promise) {
+            return this.promise;
+        }
+
+        this.promise = this.updateToken()
+            .then(token => addCredentialsToMetadata(token.value))
+            .finally(() => {
+                this.promise = null;
+            })
+
+        return this.promise
     }
 }