@@ -31,6 +31,8 @@ class AadInstanceDiscoveryProvider {
31
31
private static final String DEFAULT_API_VERSION = "2020-06-01" ;
32
32
private static final String IMDS_ENDPOINT = "https://169.254.169.254/metadata/instance/compute/location?" + DEFAULT_API_VERSION + "&format=text" ;
33
33
34
+ private static final int IMDS_TIMEOUT = 2 ;
35
+ private static final TimeUnit IMDS_TIMEOUT_UNIT = TimeUnit .SECONDS ;
34
36
static final TreeSet <String > TRUSTED_HOSTS_SET = new TreeSet <>(String .CASE_INSENSITIVE_ORDER );
35
37
static final TreeSet <String > TRUSTED_SOVEREIGN_HOSTS_SET = new TreeSet <>(String .CASE_INSENSITIVE_ORDER );
36
38
@@ -66,13 +68,19 @@ static InstanceDiscoveryMetadataEntry getMetadataEntry(URL authorityUrl,
66
68
67
69
try {
68
70
log .info ("Starting call to IMDS endpoint." );
69
- host = future .get (2 , TimeUnit . SECONDS );
71
+ host = future .get (IMDS_TIMEOUT , IMDS_TIMEOUT_UNIT );
70
72
} catch (TimeoutException ex ) {
71
73
log .info ("Cancelled call to IMDS endpoint after waiting for 2 seconds" );
72
74
future .cancel (true );
75
+ if (msalRequest .application ().azureRegion () != null ) {
76
+ host = getRegionalizedHost (authorityUrl .getHost (), msalRequest .application ().azureRegion ());
77
+ }
73
78
} catch (Exception ex ) {
74
79
// handle other exceptions
75
80
log .info ("Exception while calling IMDS endpoint" + ex .getMessage ());
81
+ if (msalRequest .application ().azureRegion () != null ) {
82
+ host = getRegionalizedHost (authorityUrl .getHost (), msalRequest .application ().azureRegion ());
83
+ }
76
84
} finally {
77
85
executor .shutdownNow ();
78
86
}
@@ -113,7 +121,7 @@ private static String performRegionalDiscovery(URL authorityUrl, MsalRequest msa
113
121
&& null != detectedRegion ) {
114
122
msalRequest .application ().azureRegion = detectedRegion ;
115
123
}
116
- cacheRegionInstanceMetadata (authorityUrl . getHost (), msalRequest . application (). azureRegion ());
124
+ cacheRegionInstanceMetadata (host , authorityUrl . getHost ());
117
125
serviceBundle .getServerSideTelemetry ().getCurrentRequest ().regionOutcome (
118
126
determineRegionOutcome (detectedRegion , msalRequest .application ().azureRegion (), msalRequest .application ().autoDetectRegion ()));
119
127
}
@@ -184,11 +192,10 @@ private static boolean shouldUseRegionalEndpoint(MsalRequest msalRequest){
184
192
return false ;
185
193
}
186
194
187
- static void cacheRegionInstanceMetadata (String host , String region ) {
195
+ static void cacheRegionInstanceMetadata (String regionalHost , String host ) {
188
196
189
197
Set <String > aliases = new HashSet <>();
190
198
aliases .add (host );
191
- String regionalHost = getRegionalizedHost (host , region );
192
199
193
200
cache .putIfAbsent (regionalHost , InstanceDiscoveryMetadataEntry .builder ().
194
201
preferredCache (host ).
0 commit comments