Skip to content

Commit 9a286ab

Browse files
committed
Fix failing tests
1 parent 7150965 commit 9a286ab

File tree

1 file changed

+11
-4
lines changed

1 file changed

+11
-4
lines changed

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AadInstanceDiscoveryProvider.java

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ class AadInstanceDiscoveryProvider {
3131
private static final String DEFAULT_API_VERSION = "2020-06-01";
3232
private static final String IMDS_ENDPOINT = "https://169.254.169.254/metadata/instance/compute/location?" + DEFAULT_API_VERSION + "&format=text";
3333

34+
private static final int IMDS_TIMEOUT = 2;
35+
private static final TimeUnit IMDS_TIMEOUT_UNIT = TimeUnit.SECONDS;
3436
static final TreeSet<String> TRUSTED_HOSTS_SET = new TreeSet<>(String.CASE_INSENSITIVE_ORDER);
3537
static final TreeSet<String> TRUSTED_SOVEREIGN_HOSTS_SET = new TreeSet<>(String.CASE_INSENSITIVE_ORDER);
3638

@@ -66,13 +68,19 @@ static InstanceDiscoveryMetadataEntry getMetadataEntry(URL authorityUrl,
6668

6769
try {
6870
log.info("Starting call to IMDS endpoint.");
69-
host = future.get(2, TimeUnit.SECONDS);
71+
host = future.get(IMDS_TIMEOUT, IMDS_TIMEOUT_UNIT);
7072
} catch (TimeoutException ex) {
7173
log.info("Cancelled call to IMDS endpoint after waiting for 2 seconds");
7274
future.cancel(true);
75+
if (msalRequest.application().azureRegion() != null) {
76+
host = getRegionalizedHost(authorityUrl.getHost(), msalRequest.application().azureRegion());
77+
}
7378
} catch (Exception ex) {
7479
// handle other exceptions
7580
log.info("Exception while calling IMDS endpoint" + ex.getMessage());
81+
if (msalRequest.application().azureRegion() != null) {
82+
host = getRegionalizedHost(authorityUrl.getHost(), msalRequest.application().azureRegion());
83+
}
7684
} finally {
7785
executor.shutdownNow();
7886
}
@@ -113,7 +121,7 @@ private static String performRegionalDiscovery(URL authorityUrl, MsalRequest msa
113121
&& null != detectedRegion) {
114122
msalRequest.application().azureRegion = detectedRegion;
115123
}
116-
cacheRegionInstanceMetadata(authorityUrl.getHost(), msalRequest.application().azureRegion());
124+
cacheRegionInstanceMetadata(host, authorityUrl.getHost());
117125
serviceBundle.getServerSideTelemetry().getCurrentRequest().regionOutcome(
118126
determineRegionOutcome(detectedRegion, msalRequest.application().azureRegion(), msalRequest.application().autoDetectRegion()));
119127
}
@@ -184,11 +192,10 @@ private static boolean shouldUseRegionalEndpoint(MsalRequest msalRequest){
184192
return false;
185193
}
186194

187-
static void cacheRegionInstanceMetadata(String host, String region) {
195+
static void cacheRegionInstanceMetadata(String regionalHost, String host) {
188196

189197
Set<String> aliases = new HashSet<>();
190198
aliases.add(host);
191-
String regionalHost = getRegionalizedHost(host, region);
192199

193200
cache.putIfAbsent(regionalHost, InstanceDiscoveryMetadataEntry.builder().
194201
preferredCache(host).

0 commit comments

Comments
 (0)