@@ -26,19 +26,26 @@ import { IamAuthUtils, TokenInfo } from "../utils/iam_auth_utils";
2626import { ClientWrapper } from "../client_wrapper" ;
2727import { RegionUtils } from "../utils/region_utils" ;
2828import { CanReleaseResources } from "../can_release_resources" ;
29+ import { RdsUrlType } from "../utils/rds_url_type" ;
30+ import { RdsUtils } from "../utils/rds_utils" ;
31+ import { GDBRegionUtils } from "../utils/gdb_region_utils" ;
2932
3033export class IamAuthenticationPlugin extends AbstractConnectionPlugin implements CanReleaseResources {
3134 private static readonly SUBSCRIBED_METHODS = new Set < string > ( [ "connect" , "forceConnect" ] ) ;
3235 protected static readonly tokenCache = new Map < string , TokenInfo > ( ) ;
3336 private readonly telemetryFactory ;
3437 private readonly fetchTokenCounter ;
35- private pluginService : PluginService ;
38+ private readonly pluginService : PluginService ;
39+ private readonly rdsUtils : RdsUtils = new RdsUtils ( ) ;
40+ protected regionUtils : RegionUtils ;
41+ protected readonly iamAuthUtils : IamAuthUtils ;
3642
37- constructor ( pluginService : PluginService ) {
43+ constructor ( pluginService : PluginService , iamAuthUtils : IamAuthUtils = new IamAuthUtils ( ) ) {
3844 super ( ) ;
3945 this . pluginService = pluginService ;
4046 this . telemetryFactory = this . pluginService . getTelemetryFactory ( ) ;
4147 this . fetchTokenCounter = this . telemetryFactory . createCounter ( "iam.fetchTokenCount" ) ;
48+ this . iamAuthUtils = iamAuthUtils ;
4249 }
4350
4451 getSubscribedMethods ( ) : Set < string > {
@@ -74,14 +81,22 @@ export class IamAuthenticationPlugin extends AbstractConnectionPlugin implements
7481 throw new AwsWrapperError ( `${ WrapperProperties . USER . name } is null or empty` ) ;
7582 }
7683
77- const host = IamAuthUtils . getIamHost ( props , hostInfo ) ;
78- const region : string = RegionUtils . getRegion ( props . get ( WrapperProperties . IAM_REGION . name ) , host ) ;
79- const port = IamAuthUtils . getIamPort ( props , hostInfo , this . pluginService . getCurrentClient ( ) . defaultPort ) ;
84+ const host = this . iamAuthUtils . getIamHost ( props , hostInfo ) ;
85+ const port = this . iamAuthUtils . getIamPort ( props , hostInfo , this . pluginService . getCurrentClient ( ) . defaultPort ) ;
86+
87+ const type : RdsUrlType = this . rdsUtils . identifyRdsType ( host . host ) ;
88+ this . regionUtils = type == RdsUrlType . RDS_GLOBAL_WRITER_CLUSTER ? new GDBRegionUtils ( ) : new RegionUtils ( ) ;
89+ const region : string | null = await this . regionUtils . getRegion ( WrapperProperties . IAM_REGION . name , host , props ) ;
90+
91+ if ( ! region ) {
92+ throw new AwsWrapperError ( Messages . get ( "SamlAuthPlugin.unableToDetermineRegion" , WrapperProperties . IAM_REGION . name ) ) ;
93+ }
94+
8095 const tokenExpirationSec = WrapperProperties . IAM_TOKEN_EXPIRATION . get ( props ) ;
8196 if ( tokenExpirationSec < 0 ) {
8297 throw new AwsWrapperError ( Messages . get ( "AuthenticationToken.tokenExpirationLessThanZero" ) ) ;
8398 }
84- const cacheKey : string = IamAuthUtils . getCacheKey ( port , user , host , region ) ;
99+ const cacheKey : string = this . iamAuthUtils . getCacheKey ( port , user , host . host , region ) ;
85100
86101 const tokenInfo = IamAuthenticationPlugin . tokenCache . get ( cacheKey ) ;
87102 const isCachedToken : boolean = tokenInfo !== undefined && ! tokenInfo . isExpired ( ) ;
@@ -91,8 +106,8 @@ export class IamAuthenticationPlugin extends AbstractConnectionPlugin implements
91106 WrapperProperties . PASSWORD . set ( props , tokenInfo . token ) ;
92107 } else {
93108 const tokenExpiry : number = Date . now ( ) + tokenExpirationSec * 1000 ;
94- const token = await IamAuthUtils . generateAuthenticationToken (
95- host ,
109+ const token = await this . iamAuthUtils . generateAuthenticationToken (
110+ host . host ,
96111 port ,
97112 region ,
98113 user ,
@@ -118,8 +133,8 @@ export class IamAuthenticationPlugin extends AbstractConnectionPlugin implements
118133 // Try to generate a new token and try to connect again
119134
120135 const tokenExpiry : number = Date . now ( ) + tokenExpirationSec * 1000 ;
121- const token = await IamAuthUtils . generateAuthenticationToken (
122- host ,
136+ const token = await this . iamAuthUtils . generateAuthenticationToken (
137+ host . host ,
123138 port ,
124139 region ,
125140 user ,
0 commit comments