Merge pull request Azure#11 from xingwu1/master

Workaround the autorest bugs and Fix auth bug

pom.xml

@@ -70,6 +70,12 @@
       <version>4.0.0</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>commons-io</groupId>
+      <artifactId>commons-io</artifactId>
+      <version>2.4</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
   <build>
     <plugins>
@@ -157,20 +163,20 @@
        <groupId>com.googlecode.addjars-maven-plugin</groupId>
        <artifactId>addjars-maven-plugin</artifactId>
        <version>1.0.5</version>
-          <executions>
-             <execution>
-                <goals>
-                  <goal>add-jars</goal>
-                </goals>
-                <configuration>
-                  <resources>
-                    <resource>
-                      <directory>../extlib</directory>
-                    </resource>
-                 </resources>
-              </configuration>
-           </execution>
-        </executions>
+       <executions>
+         <execution>
+           <goals>
+             <goal>add-jars</goal>
+           </goals>
+           <configuration>
+             <resources>
+               <resource>
+                 <directory>../extlib</directory>
+               </resource>
+             </resources>
+           </configuration>
+         </execution>
+       </executions>
       </plugin>
 
     </plugins>

src/main/java/com/microsoft/azure/batch/FileOperations.java

@@ -20,9 +20,10 @@
 import com.microsoft.azure.batch.protocol.models.FileProperties;
 import com.microsoft.azure.batch.protocol.models.NodeFile;
 import com.microsoft.rest.ServiceResponseWithHeaders;
+import rx.exceptions.Exceptions;
+import rx.functions.Func1;
 
-import java.io.IOException;
-import java.io.InputStream;
+import java.io.*;
 import java.util.Collection;
 import java.util.List;
 
@@ -284,7 +285,27 @@ public InputStream getFileFromTask(String jobId, String taskId, String fileName,
         BehaviorManager bhMgr = new BehaviorManager(this.customBehaviors(), additionalBehaviors);
         bhMgr.applyRequestBehaviors(options);
 
-        return this._parentBatchClient.protocolLayer().files().getFromTask(jobId, taskId, fileName, options);
+        // Duplicate the stream due to AutoRest issue https://github.com/Azure/autorest/issues/1385
+        ByteArrayOutputStream output = this._parentBatchClient.protocolLayer().files().getFromTaskAsync(jobId, taskId, fileName, options)
+                .map(new Func1<InputStream, ByteArrayOutputStream>() {
+            @Override
+            public ByteArrayOutputStream call(InputStream input) {
+                ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+                byte[] data = new byte[16384];
+                int nRead;
+                try {
+                    while ((nRead = input.read(data, 0, data.length)) != -1) {
+                        buffer.write(data, 0, nRead);
+                    }
+                    buffer.flush();
+                    return buffer;
+                } catch (IOException e) {
+                    throw Exceptions.propagate(e);
+                }
+            }
+        }).toBlocking().single();
+
+        return new ByteArrayInputStream(output.toByteArray());
     }
 
     /**
@@ -317,7 +338,27 @@ public InputStream getFileFromComputeNode(String poolId, String nodeId, String f
         BehaviorManager bhMgr = new BehaviorManager(this.customBehaviors(), additionalBehaviors);
         bhMgr.applyRequestBehaviors(options);
 
-        return this._parentBatchClient.protocolLayer().files().getFromComputeNode(poolId, nodeId, fileName, options);
+        // Duplicate the stream due to AutoRest issue https://github.com/Azure/autorest/issues/1385
+        ByteArrayOutputStream output = this._parentBatchClient.protocolLayer().files().getFromComputeNodeAsync(poolId, nodeId, fileName, options)
+                .map(new Func1<InputStream, ByteArrayOutputStream>() {
+            @Override
+            public ByteArrayOutputStream call(InputStream input) {
+                ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+                byte[] data = new byte[16384];
+                int nRead;
+                try {
+                    while ((nRead = input.read(data, 0, data.length)) != -1) {
+                        buffer.write(data, 0, nRead);
+                    }
+                    buffer.flush();
+                    return buffer;
+                } catch (IOException e) {
+                    throw Exceptions.propagate(e);
+                }
+            }
+        }).toBlocking().single();
+
+        return new ByteArrayInputStream(output.toByteArray());
     }
 
     /**

src/main/java/com/microsoft/azure/batch/auth/BatchSharedKeyCredentialsInterceptor.java

@@ -18,6 +18,8 @@
 import javax.crypto.spec.SecretKeySpec;
 import java.io.IOException;
 import java.net.URLDecoder;
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Locale;
@@ -31,6 +33,8 @@ class BatchSharedKeyCredentialsInterceptor implements Interceptor {
 
     private final BatchSharedKeyCredentials credentials;
 
+    private Mac hmacSha256;
+
     /**
      * Constructor for BatchSharedKeyCredentialsInterceptor
      *
@@ -62,20 +66,26 @@ private String headerValue(Request request, String headerName) {
         return headerValue;
     }
 
-    private String sign(String accessKey, String stringToSign) {
+    private synchronized String sign(String stringToSign) {
         try {
             // Encoding the Signature
             // Signature=Base64(HMAC-SHA256(UTF8(StringToSign)))
-            Mac hmac = Mac.getInstance("hmacSHA256");
-            hmac.init(new SecretKeySpec(Base64.decodeBase64(accessKey),
-                    "hmacSHA256"));
-            byte[] digest = hmac.doFinal(stringToSign.getBytes("UTF-8"));
+            byte[] digest = getHmac256().doFinal(stringToSign.getBytes("UTF-8"));
             return Base64.encodeBase64String(digest);
         } catch (Exception e) {
             throw new IllegalArgumentException("accessKey", e);
         }
     }
 
+    private synchronized Mac getHmac256() throws NoSuchAlgorithmException, InvalidKeyException {
+        if (this.hmacSha256 == null) {
+            // Initializes the HMAC-SHA256 Mac and SecretKey.
+            this.hmacSha256 = Mac.getInstance("HmacSHA256");
+            this.hmacSha256.init(new SecretKeySpec(Base64.decodeBase64(this.credentials.keyValue()), "HmacSHA256"));
+        }
+        return this.hmacSha256;
+    }
+
     private Request signHeader(Request request) throws IOException {
 
         Request.Builder builder = request.newBuilder();
@@ -141,7 +151,7 @@ private Request signHeader(Request request) throws IOException {
 
         signature = signature + "/"
                 + credentials.accountName().toLowerCase() + "/"
-                + request.url().uri().getPath().replaceAll("^[/]+", "");
+                + request.url().uri().getRawPath().replaceAll("^[/]+", "");
         // We temporary change client side auth code generator to bypass server
         // bug 4092533
         signature = signature.replace("%5C", "/").replace("%2F", "/");
@@ -156,22 +166,18 @@ private Request signHeader(Request request) throws IOException {
                         .toLowerCase(Locale.US);
                 queryComponents.put(
                         key,
-                        key
-                                + ":"
-                                + URLDecoder.decode(pair.substring(idx + 1),
-                                        "UTF-8"));
+                        key + ":" + URLDecoder.decode(pair.substring(idx + 1),"UTF-8"));
             }
 
             for (Map.Entry<String, String> entry : queryComponents.entrySet()) {
                 signature = signature + "\n" + entry.getValue();
             }
         }
-        String signedSignature = sign(credentials.keyValue(), signature);
+        String signedSignature = sign(signature);
         String authorization = "SharedKey " + credentials.accountName()
                 + ":" + signedSignature;
         builder.header("Authorization", authorization);
 
         return builder.build();
     }
-
 }

src/main/java/com/microsoft/azure/batch/protocol/models/FileGetPropertiesFromComputeNodeHeaders.java

@@ -43,13 +43,13 @@ public class FileGetPropertiesFromComputeNodeHeaders {
      * particular, you can pass the ETag to one of the If-Modified-Since,
      * If-Unmodified-Since, If-Match or If-None-Match headers.
      */
-    @JsonProperty(value = "ETag")
+    @JsonProperty(value = "etag")
     private String eTag;
 
     /**
      * The time at which the resource was last modified.
      */
-    @JsonProperty(value = "Last-Modified")
+    @JsonProperty(value = "last-modified")
     private DateTimeRfc1123 lastModified;
 
     /**
@@ -79,13 +79,13 @@ public class FileGetPropertiesFromComputeNodeHeaders {
     /**
      * The content type of the file.
      */
-    @JsonProperty(value = "Content-Type")
+    @JsonProperty(value = "content-type")
     private String contentType;
 
     /**
      * The length of the file.
      */
-    @JsonProperty(value = "Content-Length")
+    @JsonProperty(value = "content-length")
     private Long contentLength;
 
     /**

src/main/java/com/microsoft/azure/batch/protocol/models/FileGetPropertiesFromTaskHeaders.java

@@ -43,13 +43,13 @@ public class FileGetPropertiesFromTaskHeaders {
      * particular, you can pass the ETag to one of the If-Modified-Since,
      * If-Unmodified-Since, If-Match or If-None-Match headers.
      */
-    @JsonProperty(value = "ETag")
+    @JsonProperty(value = "etag")
     private String eTag;
 
     /**
      * The time at which the resource was last modified.
      */
-    @JsonProperty(value = "Last-Modified")
+    @JsonProperty(value = "last-modified")
     private DateTimeRfc1123 lastModified;
 
     /**
@@ -79,13 +79,13 @@ public class FileGetPropertiesFromTaskHeaders {
     /**
      * The content type of the file.
      */
-    @JsonProperty(value = "Content-Type")
+    @JsonProperty(value = "content-type")
     private String contentType;
 
     /**
      * The length of the file.
      */
-    @JsonProperty(value = "Content-Length")
+    @JsonProperty(value = "content-length")
     private Long contentLength;
 
     /**

src/test/java/com/microsoft/azure/batch/BatchTestBase.java

@@ -17,6 +17,7 @@
 import java.io.IOException;
 import java.net.URISyntaxException;
 import java.security.InvalidKeyException;
+import java.time.Duration;
 import java.util.*;
 
 import com.microsoft.azure.storage.CloudStorageAccount;
@@ -162,4 +163,44 @@ static String uploadFileToCloud(CloudBlobContainer container, String fileName, S
         String sas = blob.generateSharedAccessSignature(policy, null);
         return blob.getUri() + "?" + sas;
     }
+
+    /**
+     * Wait all tasks under a specified job to be completed
+     * @param client batch client instance
+     * @param jobId job id
+     * @param expiryTime the waiting period
+     * @return if task completed in time, return true, otherwise, return false
+     * @throws BatchErrorException
+     * @throws IOException
+     * @throws InterruptedException
+     */
+    static boolean waitForTasksToComplete(BatchClient client, String jobId, Duration expiryTime) throws BatchErrorException, IOException, InterruptedException {
+        long startTime = System.currentTimeMillis();
+        long elapsedTime = 0L;
+
+        while (elapsedTime < expiryTime.toMillis()) {
+            List<CloudTask> taskCollection = client.taskOperations().listTasks(jobId, new DetailLevel.Builder().withSelectClause("id, state").build());
+
+            boolean allComplete = true;
+            for (CloudTask task : taskCollection) {
+                if (task.state() != TaskState.COMPLETED) {
+                    allComplete = false;
+                    break;
+                }
+            }
+
+            if (allComplete) {
+                // All tasks completed
+                return true;
+            }
+
+            // Check again after 10 seconds
+            Thread.sleep(10 * 1000);
+            elapsedTime = (new Date()).getTime() - startTime;
+        }
+
+        // Timeout, return false
+        return false;
+    }
+
 }