Support for Long-Running-Operation-options extension handling (#442)

* Support for Long-Running-Operation-options extension handling

* ensure MSICredentials::getToken honor provided tokenAudience

azure-client-authentication/src/main/java/com/microsoft/azure/credentials/MSICredentials.java

@@ -60,7 +60,7 @@ public MSICredentials() {
      */
     public MSICredentials(AzureEnvironment environment) {
         super(environment, null /** retrieving MSI token does not require tenant **/);
-        this.resource = environment.resourceManagerEndpoint();
+        this.resource = environment.managementEndpoint();
         this.tokenSource = MSITokenSource.IMDS_ENDPOINT;
     }
 
@@ -74,7 +74,7 @@ public MSICredentials(AzureEnvironment environment) {
     @Deprecated()
     public MSICredentials(AzureEnvironment environment, int msiPort) {
         super(environment, null /** retrieving MSI token does not require tenant **/);
-        this.resource = environment.resourceManagerEndpoint();
+        this.resource = environment.managementEndpoint();
         this.msiPort = msiPort;
         this.tokenSource = MSITokenSource.MSI_EXTENSION;
     }
@@ -126,17 +126,17 @@ public MSICredentials withIdentityId(String identityId) {
 
 
     @Override
-    public String getToken(String resource) throws IOException {
+    public String getToken(String tokenAudience) throws IOException {
         if (this.tokenSource == MSITokenSource.MSI_EXTENSION) {
-            return this.getTokenFromMSIExtension();
+            return this.getTokenFromMSIExtension(tokenAudience == null ? this.resource : tokenAudience);
         } else {
-            return this.getTokenFromIMDSEndpoint();
+            return this.getTokenFromIMDSEndpoint(tokenAudience == null ? this.resource : tokenAudience);
         }
     }
 
-    private String getTokenFromMSIExtension() throws IOException {
+    private String getTokenFromMSIExtension(String tokenAudience) throws IOException {
         URL url = new URL(String.format("http://localhost:%d/oauth2/token", this.msiPort));
-        String postData = String.format("resource=%s", this.resource);
+        String postData = String.format("resource=%s", tokenAudience);
         if (this.objectId != null) {
             postData += String.format("&object_id=%s", this.objectId);
         } else if (this.clientId != null) {
@@ -177,21 +177,21 @@ private String getTokenFromMSIExtension() throws IOException {
         }
     }
 
-    private String getTokenFromIMDSEndpoint() {
-        MSIToken token = cache.get(resource);
+    private String getTokenFromIMDSEndpoint(String tokenAudience) {
+        MSIToken token = cache.get(tokenAudience);
         if (token != null && !token.isExpired()) {
             return token.accessToken();
         }
         lock.lock();
         try {
-            token = cache.get(resource);
+            token = cache.get(tokenAudience);
             if (token != null && !token.isExpired()) {
                 return token.accessToken();
             }
             try {
-                token = retrieveTokenFromIDMSWithRetry();
+                token = retrieveTokenFromIDMSWithRetry(tokenAudience);
                 if (token != null) {
-                    cache.put(resource, token);
+                    cache.put(tokenAudience, token);
                 }
             } catch (IOException exception) {
                 throw new RuntimeException(exception);
@@ -202,7 +202,7 @@ private String getTokenFromIMDSEndpoint() {
         }
     }
 
-    private MSIToken retrieveTokenFromIDMSWithRetry() throws IOException {
+    private MSIToken retrieveTokenFromIDMSWithRetry(String tokenAudience) throws IOException {
         StringBuilder payload = new StringBuilder();
         final int imdsUpgradeTimeInMs = 70 * 1000;
 
@@ -214,7 +214,7 @@ private MSIToken retrieveTokenFromIDMSWithRetry() throws IOException {
             payload.append("&");
             payload.append("resource");
             payload.append("=");
-            payload.append(URLEncoder.encode(this.resource, "UTF-8"));
+            payload.append(URLEncoder.encode(tokenAudience, "UTF-8"));
             if (this.objectId != null) {
                 payload.append("&");
                 payload.append("object_id");

azure-client-runtime/src/main/java/com/microsoft/azure/AzureClient.java

@@ -162,7 +162,7 @@ public PollingState<T> call(Response<ResponseBody> response) {
                     throw  exception;
                 }
                 try {
-                    final PollingState<T> pollingState = PollingState.create(response, longRunningOperationRetryTimeout(), resourceType, restClient().serializerAdapter());
+                    final PollingState<T> pollingState = PollingState.create(response, LongRunningOperationOptions.DEFAULT, longRunningOperationRetryTimeout(), resourceType, restClient().serializerAdapter());
                     pollingState.withPollingUrlFromResponse(response);
                     pollingState.withPollingRetryTimeoutFromResponse(response);
                     pollingState.withPutOrPatchResourceUri(response.raw().request().url().toString());
@@ -337,7 +337,22 @@ public <T, THeader> ServiceResponseWithHeaders<T, THeader> getPostOrDeleteResult
      * @return          the task describing the asynchronous polling.
      */
     public <T> Observable<ServiceResponse<T>> getPostOrDeleteResultAsync(Observable<Response<ResponseBody>> observable, final Type resourceType) {
-        return this.<T>beginPostOrDeleteAsync(observable, resourceType)
+        return this.<T>getPostOrDeleteResultAsync(observable, LongRunningOperationOptions.DEFAULT, resourceType);
+    }
+
+    /**
+     * Handles an initial response from a POST or DELETE operation response by polling
+     * the status of the operation asynchronously, calling the user provided callback
+     * when the operation terminates.
+     *
+     * @param observable  the initial response from the POST or DELETE operation.
+     * @param lroOptions long running operation options.
+     * @param <T>       the return type of the caller.
+     * @param resourceType the java.lang.reflect.Type of the resource.
+     * @return          the task describing the asynchronous polling.
+     */
+    public <T> Observable<ServiceResponse<T>> getPostOrDeleteResultAsync(Observable<Response<ResponseBody>> observable, final LongRunningOperationOptions lroOptions, final Type resourceType) {
+        return this.<T>beginPostOrDeleteAsync(observable, lroOptions, resourceType)
                 .toObservable()
                 .flatMap(new Func1<PollingState<T>, Observable<PollingState<T>>>() {
                     @Override
@@ -359,12 +374,13 @@ public ServiceResponse<T> call(PollingState<T> pollingState) {
      * when subscribed to it, the deferred action will be performed and emits the polling state containing information
      * to track the progress of the action.
      *
-     * @param observable an observable representing a deferred PUT or PATCH operation.
+     * @param observable an observable representing a deferred POST or DELETE operation.
+     * @param lroOptions long running operation options.
      * @param resourceType the java.lang.reflect.Type of the resource.
      * @param <T> the type of the resource
      * @return the observable of which a subscription will lead POST or DELETE action.
      */
-    public <T> Single<PollingState<T>> beginPostOrDeleteAsync(Observable<Response<ResponseBody>> observable, final Type resourceType) {
+    public <T> Single<PollingState<T>> beginPostOrDeleteAsync(Observable<Response<ResponseBody>> observable, final LongRunningOperationOptions lroOptions, final Type resourceType) {
         return observable.map(new Func1<Response<ResponseBody>, PollingState<T>>() {
             @Override
             public PollingState<T> call(Response<ResponseBody> response) {
@@ -373,7 +389,7 @@ public PollingState<T> call(Response<ResponseBody> response) {
                     throw  exception;
                 }
                 try {
-                    final PollingState<T> pollingState = PollingState.create(response, longRunningOperationRetryTimeout(), resourceType, restClient().serializerAdapter());
+                    final PollingState<T> pollingState = PollingState.create(response, lroOptions, longRunningOperationRetryTimeout(), resourceType, restClient().serializerAdapter());
                     pollingState.withPollingUrlFromResponse(response);
                     pollingState.withPollingRetryTimeoutFromResponse(response);
                     return pollingState;

azure-client-runtime/src/main/java/com/microsoft/azure/LongRunningFinalState.java

@@ -0,0 +1,29 @@
+/**
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT License. See License.txt in the project root for
+ * license information.
+ */
+
+package com.microsoft.azure;
+
+/**
+ * Describes how to retrieve the final state of a long running operation.
+ */
+public enum LongRunningFinalState {
+    /**
+     * Indicate that no specific action required to retrieve the final state.
+     */
+    DEFAULT,
+    /**
+     * Indicate that use azure async operation uri to retrieve the final state.
+     */
+    AZURE_ASYNC_OPERATION,
+    /**
+     * Indicate that use location uri to retrieve the final state.
+     */
+    LOCATION,
+    /**
+     * Indicate that use original uri to retrieve the final state.
+     */
+    ORIGINAL_URI
+}

azure-client-runtime/src/main/java/com/microsoft/azure/LongRunningOperationOptions.java

@@ -0,0 +1,40 @@
+/**
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT License. See License.txt in the project root for
+ * license information.
+ */
+
+package com.microsoft.azure;
+
+/**
+ * Type representing LRO meta-data present in the x-ms-long-running-operation-options autorest extension.
+ */
+public final class LongRunningOperationOptions {
+    /**
+     * Default instance of this type.
+     */
+    public static final LongRunningOperationOptions DEFAULT = new LongRunningOperationOptions().withFinalStateVia(LongRunningFinalState.DEFAULT);
+
+    /**
+     * Describes how to retrieve the final state of the LRO.
+     */
+    private LongRunningFinalState finalStateVia;
+
+    /**
+     * @return indicates how to retrieve the final state of LRO.
+     */
+    public LongRunningFinalState finalStateVia() {
+        return this.finalStateVia;
+    }
+
+    /**
+     * Sets LongRunningFinalState value.
+     *
+     * @param finalStateVia indicates how to retrieve the final state of LRO.
+     * @return LongRunningOperationOptions
+     */
+    public LongRunningOperationOptions withFinalStateVia(LongRunningFinalState finalStateVia) {
+        this.finalStateVia = finalStateVia;
+        return this;
+    }
+}

azure-client-runtime/src/main/java/com/microsoft/azure/PollingState.java

@@ -51,7 +51,8 @@ public class PollingState<T> {
     private String putOrPatchResourceUri;
     /** The logging context. **/
     private String loggingContext;
-
+    /** indicate how to retrieve the final state of LRO. **/
+    private LongRunningFinalState finalStateVia;
 
     // Non-serializable properties
     //
@@ -87,21 +88,23 @@ public class PollingState<T> {
      * Creates a polling state.
      *
      * @param response the response from Retrofit REST call that initiate the long running operation.
+     * @param lroOptions long running operation options.
      * @param defaultRetryTimeout the long running operation retry timeout.
      * @param resourceType the type of the resource the long running operation returns
      * @param serializerAdapter the adapter for the Jackson object mapper
      * @param <T> the result type
      * @return the polling state
      * @throws IOException thrown by deserialization
      */
-    public static <T> PollingState<T> create(Response<ResponseBody> response, int defaultRetryTimeout, Type resourceType, SerializerAdapter<?> serializerAdapter) throws IOException {
+    public static <T> PollingState<T> create(Response<ResponseBody> response, LongRunningOperationOptions lroOptions, int defaultRetryTimeout, Type resourceType, SerializerAdapter<?> serializerAdapter) throws IOException {
         PollingState<T> pollingState = new PollingState<>();
         pollingState.initialHttpMethod = response.raw().request().method();
         pollingState.defaultRetryTimeout = defaultRetryTimeout;
         pollingState.withResponse(response);
         pollingState.resourceType = resourceType;
         pollingState.serializerAdapter = serializerAdapter;
         pollingState.loggingContext = response.raw().request().header(LOGGING_HEADER);
+        pollingState.finalStateVia = lroOptions.finalStateVia();
 
         String responseContent = null;
         PollingResource resource = null;
@@ -172,6 +175,7 @@ public static <ResultT> PollingState<ResultT> createFromPollingState(PollingStat
         pollingState.defaultRetryTimeout = other.defaultRetryTimeout;
         pollingState.retryTimeout = other.retryTimeout;
         pollingState.loggingContext = other.loggingContext;
+        pollingState.finalStateVia = other.finalStateVia;
         return pollingState;
     }
 
@@ -353,11 +357,17 @@ boolean isStatusSucceeded() {
     }
 
     boolean resourcePending() {
-        return statusCode() != 204
+        boolean resourcePending = statusCode() != 204
                 && isStatusSucceeded()
                 && resource() == null
                 && resourceType() != Void.class
                 && locationHeaderLink() != null;
+        if (resourcePending) {
+            // Keep current behaviour for backward-compact
+            return true;
+        } else {
+            return this.finalStateVia() == LongRunningFinalState.LOCATION;
+        }
     }
 
     /**
@@ -452,6 +462,30 @@ Type resourceType() {
         return resourceType;
     }
 
+    /**
+     * @return describes how to retrieve the final result of long running operation.
+     */
+    LongRunningFinalState finalStateVia() {
+        if (this.initialHttpMethod().equalsIgnoreCase("POST") && resourceType() != Void.class) {
+            // FinalStateVia is supported only for POST LRO at the moment.
+            //
+            if (this.locationHeaderLink() != null && this.azureAsyncOperationHeaderLink() != null) {
+                // Consider final-state-via option only if both headers are provided on the wire otherwise
+                // there is nothing to disambiguate.
+                //
+                if (this.finalStateVia == LongRunningFinalState.LOCATION) {
+                    // A POST LRO can be tracked only using Location or AsyncOperation Header.
+                    // If AsyncOperationHeader is present then anyway polling will be performed using
+                    // it and there is no point in making one additional call to retrieve state using
+                    // async operation uri anyway. Hence we consider only LongRunningFinalState.LOCATION.
+                    //
+                    return LongRunningFinalState.LOCATION;
+                }
+            }
+        }
+        return LongRunningFinalState.DEFAULT;
+    }
+
     /**
      * Sets resource type.
      *