Fixed block issue and tier issue

sdk/storage/azure-storage-blob-cryptography/pom.xml

@@ -98,6 +98,12 @@
       <version>1.0.8</version> <!-- {x-version-update;com.azure:azure-identity;dependency} -->
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-security-keyvault-keys</artifactId> <!-- {x-version-update;com.azure:azure-security-keyvault-keys;dependency} -->
+      <version>4.2.0-beta.5</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <profiles>

sdk/storage/azure-storage-blob-cryptography/src/main/java/com/azure/storage/blob/specialized/cryptography/EncryptedBlobAsyncClient.java

@@ -252,7 +252,7 @@ public Mono<Response<BlockBlobItem>> uploadWithResponse(Flux<ByteBuffer> data,
         AccessTier tier, BlobRequestConditions requestConditions) {
         return this.uploadWithResponse(new BlobParallelUploadOptions(data)
             .setParallelTransferOptions(parallelTransferOptions).setHeaders(headers).setMetadata(metadata)
-            .setTier(AccessTier.HOT).setRequestConditions(requestConditions));
+            .setTier(tier).setRequestConditions(requestConditions));
     }
 
     /**
@@ -443,58 +443,58 @@ Mono<EncryptedBlob> encryptBlob(Flux<ByteBuffer> plainTextFlux) throws InvalidKe
             keyWrappingMetadata.put(CryptographyConstants.AGENT_METADATA_KEY,
                 CryptographyConstants.AGENT_METADATA_VALUE);
 
-            return this.keyWrapper.wrapKey(keyWrapAlgorithm, aesKey.getEncoded())
-                .map(encryptedKey -> {
-                    WrappedKey wrappedKey = new WrappedKey(
-                        this.keyWrapper.getKeyId().block(), encryptedKey, keyWrapAlgorithm);
-
-                    // Build EncryptionData
-                    EncryptionData encryptionData = new EncryptionData()
-                        .setEncryptionMode(CryptographyConstants.ENCRYPTION_MODE)
-                        .setEncryptionAgent(
-                            new EncryptionAgent(CryptographyConstants.ENCRYPTION_PROTOCOL_V1,
-                                EncryptionAlgorithm.AES_CBC_256))
-                        .setKeyWrappingMetadata(keyWrappingMetadata)
-                        .setContentEncryptionIV(cipher.getIV())
-                        .setWrappedContentKey(wrappedKey);
-
-                    // Encrypt plain text with content encryption key
-                    Flux<ByteBuffer> encryptedTextFlux = plainTextFlux.map(plainTextBuffer -> {
-                        int outputSize = cipher.getOutputSize(plainTextBuffer.remaining());
+            return this.keyWrapper.getKeyId().flatMap(keyId ->
+                this.keyWrapper.wrapKey(keyWrapAlgorithm, aesKey.getEncoded())
+                    .map(encryptedKey -> {
+                        WrappedKey wrappedKey = new WrappedKey(keyId, encryptedKey, keyWrapAlgorithm);
+
+                        // Build EncryptionData
+                        EncryptionData encryptionData = new EncryptionData()
+                            .setEncryptionMode(CryptographyConstants.ENCRYPTION_MODE)
+                            .setEncryptionAgent(
+                                new EncryptionAgent(CryptographyConstants.ENCRYPTION_PROTOCOL_V1,
+                                    EncryptionAlgorithm.AES_CBC_256))
+                            .setKeyWrappingMetadata(keyWrappingMetadata)
+                            .setContentEncryptionIV(cipher.getIV())
+                            .setWrappedContentKey(wrappedKey);
+
+                        // Encrypt plain text with content encryption key
+                        Flux<ByteBuffer> encryptedTextFlux = plainTextFlux.map(plainTextBuffer -> {
+                            int outputSize = cipher.getOutputSize(plainTextBuffer.remaining());
 
                         /*
                         This should be the only place we allocate memory in encryptBlob(). Although there is an
                         overload that can encrypt in place that would save allocations, we do not want to overwrite
                         customer's memory, so we must allocate our own memory. If memory usage becomes unreasonable,
                         we should implement pooling.
                          */
-                        ByteBuffer encryptedTextBuffer = ByteBuffer.allocate(outputSize);
-
-                        int encryptedBytes;
-                        try {
-                            encryptedBytes = cipher.update(plainTextBuffer, encryptedTextBuffer);
-                        } catch (ShortBufferException e) {
-                            throw logger.logExceptionAsError(Exceptions.propagate(e));
-                        }
-                        encryptedTextBuffer.position(0);
-                        encryptedTextBuffer.limit(encryptedBytes);
-                        return encryptedTextBuffer;
-                    });
+                            ByteBuffer encryptedTextBuffer = ByteBuffer.allocate(outputSize);
+
+                            int encryptedBytes;
+                            try {
+                                encryptedBytes = cipher.update(plainTextBuffer, encryptedTextBuffer);
+                            } catch (ShortBufferException e) {
+                                throw logger.logExceptionAsError(Exceptions.propagate(e));
+                            }
+                            encryptedTextBuffer.position(0);
+                            encryptedTextBuffer.limit(encryptedBytes);
+                            return encryptedTextBuffer;
+                        });
 
                     /*
                     Defer() ensures the contained code is not executed until the Flux is subscribed to, in
                     other words, cipher.doFinal() will not be called until the plainTextFlux has completed
                     and therefore all other data has been encrypted.
                      */
-                    encryptedTextFlux = Flux.concat(encryptedTextFlux, Flux.defer(() -> {
-                        try {
-                            return Flux.just(ByteBuffer.wrap(cipher.doFinal()));
-                        } catch (GeneralSecurityException e) {
-                            throw logger.logExceptionAsError(Exceptions.propagate(e));
-                        }
+                        encryptedTextFlux = Flux.concat(encryptedTextFlux, Flux.defer(() -> {
+                            try {
+                                return Flux.just(ByteBuffer.wrap(cipher.doFinal()));
+                            } catch (GeneralSecurityException e) {
+                                throw logger.logExceptionAsError(Exceptions.propagate(e));
+                            }
+                        }));
+                        return new EncryptedBlob(encryptionData, encryptedTextFlux);
                     }));
-                    return new EncryptedBlob(encryptionData, encryptedTextFlux);
-                });
         } catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
             // These are hardcoded and guaranteed to work. There is no reason to propogate a checked exception.
             throw logger.logExceptionAsError(new RuntimeException(e));

sdk/storage/azure-storage-blob-cryptography/src/test/java/com/azure/storage/blob/specialized/cryptography/APISpec.groovy

@@ -19,6 +19,7 @@ import com.azure.core.test.utils.TestResourceNamer
 import com.azure.core.util.Configuration
 import com.azure.core.util.FluxUtil
 import com.azure.core.util.logging.ClientLogger
+import com.azure.security.keyvault.keys.cryptography.models.KeyWrapAlgorithm
 import com.azure.storage.blob.BlobAsyncClient
 import com.azure.storage.blob.BlobClient
 import com.azure.storage.blob.BlobServiceClientBuilder
@@ -88,7 +89,7 @@ class APISpec extends Specification {
     static def PREMIUM_STORAGE = "PREMIUM_STORAGE_"
 
     TestResourceNamer resourceNamer
-    private InterceptorManager interceptorManager
+    def InterceptorManager interceptorManager
     protected String testName
 
     // Fields used for conveniently creating blobs with data.
@@ -229,8 +230,10 @@ class APISpec extends Specification {
                                                          AsyncKeyEncryptionKeyResolver keyResolver,
                                                          StorageSharedKeyCredential credential, String endpoint,
                                                          HttpPipelinePolicy... policies) {
+
+        KeyWrapAlgorithm algorithm = key != null && key.getKeyId().block() == "local" ? KeyWrapAlgorithm.A256KW : KeyWrapAlgorithm.RSA_OAEP_256
         EncryptedBlobClientBuilder builder = new EncryptedBlobClientBuilder()
-            .key(key, "keyWrapAlgorithm")
+            .key(key, algorithm.toString())
             .keyResolver(keyResolver)
             .endpoint(endpoint)
             .httpClient(getHttpClient())
@@ -277,7 +280,7 @@ class APISpec extends Specification {
         generateResourceName(containerPrefix, entityNo++)
     }
 
-    private String generateResourceName(String prefix, int entityNo) {
+    def String generateResourceName(String prefix, int entityNo) {
         return resourceNamer.randomName(prefix + testName + entityNo, 63)
     }
 

sdk/storage/azure-storage-blob-cryptography/src/test/java/com/azure/storage/blob/specialized/cryptography/EncyptedBlockBlobAPITest.groovy

@@ -1201,10 +1201,8 @@ class EncyptedBlockBlobAPITest extends APISpec {
         Files.deleteIfExists(file.toPath())
 
         expect:
-        def bac = new EncryptedBlobClientBuilder()
-            .key(fakeKey, "keyWrapAlgorithm")
-            .pipeline(ebc.getHttpPipeline())
-            .endpoint(ebc.getBlobUrl())
+        def bac = getEncryptedClientBuilder(fakeKey, null, primaryCredential,
+            ebc.getBlobUrl().toString())
             .buildEncryptedBlobAsyncClient()
 
         /*

sdk/storage/azure-storage-blob-cryptography/src/test/java/com/azure/storage/blob/specialized/cryptography/KeyvaultKeyTest.groovy

@@ -0,0 +1,101 @@
+package com.azure.storage.blob.specialized.cryptography
+
+import com.azure.core.cryptography.AsyncKeyEncryptionKey
+import com.azure.core.test.TestMode
+import com.azure.core.util.Configuration
+import com.azure.identity.DefaultAzureCredentialBuilder
+import com.azure.identity.EnvironmentCredentialBuilder
+import com.azure.security.keyvault.keys.KeyClient
+import com.azure.security.keyvault.keys.KeyClientBuilder
+import com.azure.security.keyvault.keys.cryptography.KeyEncryptionKeyClientBuilder
+import com.azure.security.keyvault.keys.models.CreateRsaKeyOptions
+import com.azure.security.keyvault.keys.models.KeyVaultKey
+import com.azure.storage.blob.BlobContainerClient
+import com.azure.storage.common.implementation.Constants
+
+import java.time.OffsetDateTime
+
+class KeyvaultKeyTest extends APISpec {
+
+    BlobContainerClient cc
+    EncryptedBlobClient bec // encrypted client for download
+    KeyClient keyClient
+    String keyId
+
+    def setup() {
+        def keyVaultUrl = Configuration.getGlobalConfiguration().get("KEYVAULT_URL")
+
+        KeyClientBuilder builder = new KeyClientBuilder()
+
+        if (testMode != TestMode.PLAYBACK) {
+            if (testMode == TestMode.RECORD) {
+                builder.addPolicy(interceptorManager.getRecordPolicy())
+            }
+            // AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_CLIENT_SECRET
+            builder.credential(new EnvironmentCredentialBuilder().build())
+        } else {
+            builder.credential(new DefaultAzureCredentialBuilder().build())
+        }
+
+        keyClient = builder
+            .httpClient(getHttpClient())
+            .vaultUrl(keyVaultUrl)
+            .buildClient()
+
+        keyId = generateResourceName("keyId", entityNo++)
+
+        KeyVaultKey keyVaultKey = keyClient.createRsaKey(new CreateRsaKeyOptions(keyId)
+            .setExpiresOn(OffsetDateTime.now().plusYears(1))
+            .setKeySize(2048))
+
+        AsyncKeyEncryptionKey akek = new KeyEncryptionKeyClientBuilder()
+            .credential(new DefaultAzureCredentialBuilder().build())
+            .buildAsyncKeyEncryptionKey(keyVaultKey.getId())
+            .block()
+
+        cc = getServiceClientBuilder(primaryCredential,
+            String.format(defaultEndpointTemplate, primaryCredential.getAccountName()))
+            .buildClient()
+            .getBlobContainerClient(generateContainerName())
+        cc.create()
+
+        bec = getEncryptedClientBuilder(akek, null, primaryCredential,
+            cc.getBlobContainerUrl().toString())
+            .blobName(generateBlobName())
+            .buildEncryptedBlobClient()
+    }
+
+    def cleanup() {
+        keyClient.beginDeleteKey(keyId)
+    }
+
+    def "upload download"() {
+        setup:
+        def inputArray = getRandomByteArray(Constants.KB)
+        InputStream stream = new ByteArrayInputStream(inputArray)
+        def os = new ByteArrayOutputStream()
+
+        when:
+        bec.upload(stream, Constants.KB)
+        bec.download(os)
+
+        then:
+        inputArray == os.toByteArray()
+    }
+
+
+    def "encryption not a noop"() {
+        setup:
+        def inputArray = getRandomByteArray(Constants.KB)
+        InputStream stream = new ByteArrayInputStream(inputArray)
+        def os = new ByteArrayOutputStream()
+
+        when:
+        bec.upload(stream, Constants.KB)
+        cc.getBlobClient(bec.getBlobName()).download(os)
+
+        then:
+        inputArray != os.toByteArray()
+    }
+
+}

sdk/storage/azure-storage-blob-cryptography/src/test/java/com/azure/storage/blob/specialized/cryptography/LocalKeyTest.groovy

@@ -0,0 +1,70 @@
+package com.azure.storage.blob.specialized.cryptography
+
+import com.azure.core.cryptography.AsyncKeyEncryptionKey
+import com.azure.security.keyvault.keys.cryptography.LocalKeyEncryptionKeyClientBuilder
+import com.azure.security.keyvault.keys.models.JsonWebKey
+import com.azure.security.keyvault.keys.models.KeyOperation
+import com.azure.storage.blob.BlobContainerClient
+import com.azure.storage.common.implementation.Constants
+
+import javax.crypto.spec.SecretKeySpec
+
+class LocalKeyTest extends APISpec {
+
+    BlobContainerClient cc
+    EncryptedBlobClient bec // encrypted client for download
+
+    def setup() {
+
+        /* Insecurely generate a local key*/
+        def byteKey = getRandomByteArray(256)
+
+        JsonWebKey localKey = JsonWebKey.fromAes(new SecretKeySpec(byteKey, "AES"),
+            Arrays.asList(KeyOperation.WRAP_KEY, KeyOperation.UNWRAP_KEY))
+            .setId("local")
+        AsyncKeyEncryptionKey akek = new LocalKeyEncryptionKeyClientBuilder()
+            .buildAsyncKeyEncryptionKey(localKey)
+            .block();
+
+        cc = getServiceClientBuilder(primaryCredential,
+            String.format(defaultEndpointTemplate, primaryCredential.getAccountName()))
+            .buildClient()
+            .getBlobContainerClient(generateContainerName())
+        cc.create()
+
+        bec = getEncryptedClientBuilder(akek, null, primaryCredential,
+            cc.getBlobContainerUrl().toString())
+            .blobName(generateBlobName())
+            .buildEncryptedBlobClient()
+    }
+
+    def "upload download"() {
+        setup:
+        def inputArray = getRandomByteArray(Constants.KB)
+        InputStream stream = new ByteArrayInputStream(inputArray)
+        def os = new ByteArrayOutputStream()
+
+        when:
+        bec.upload(stream, Constants.KB)
+        bec.download(os)
+
+        then:
+        inputArray == os.toByteArray()
+    }
+
+
+    def "encryption not a noop"() {
+        setup:
+        def inputArray = getRandomByteArray(Constants.KB)
+        InputStream stream = new ByteArrayInputStream(inputArray)
+        def os = new ByteArrayOutputStream()
+
+        when:
+        bec.upload(stream, Constants.KB)
+        cc.getBlobClient(bec.getBlobName()).download(os)
+
+        then:
+        inputArray != os.toByteArray()
+    }
+
+}