Add Azure storage tests

Signed-off-by: Daniel Weiße <dw@edgeless.systems>

go.mod

@@ -53,13 +53,14 @@ require (
 	github.com/Azure/go-autorest/autorest/azure/auth v0.5.11
 	github.com/Azure/go-autorest/autorest/date v0.3.0
 	github.com/Azure/go-autorest/autorest/to v0.4.0
-	github.com/aws/aws-sdk-go-v2 v1.15.0
+	github.com/aws/aws-sdk-go-v2 v1.16.1
-	github.com/aws/aws-sdk-go-v2/config v1.15.0
+	github.com/aws/aws-sdk-go-v2/config v1.15.2
-	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.0
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.2
+	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.11.2
 	github.com/aws/aws-sdk-go-v2/service/ec2 v1.32.0
 	github.com/aws/aws-sdk-go-v2/service/kms v1.16.0
-	github.com/aws/aws-sdk-go-v2/service/s3 v1.26.0
+	github.com/aws/aws-sdk-go-v2/service/s3 v1.26.2
-	github.com/aws/smithy-go v1.11.1
+	github.com/aws/smithy-go v1.11.2
 	github.com/coreos/go-systemd/v22 v22.3.2
 	github.com/docker/docker v20.10.13+incompatible
 	github.com/docker/go-connections v0.4.0
@@ -116,17 +117,17 @@ require (
 	github.com/Microsoft/hcsshim v0.9.2 // indirect
 	github.com/PuerkitoBio/purell v1.1.1 // indirect
 	github.com/PuerkitoBio/urlesc v0.0.0-20170810143723-de5bf2ad4578 // indirect
-	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.4.0 // indirect
+	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.4.1 // indirect
-	github.com/aws/aws-sdk-go-v2/credentials v1.10.0 // indirect
+	github.com/aws/aws-sdk-go-v2/credentials v1.11.1 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.6 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.8 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.0 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.2 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/ini v1.3.7 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/ini v1.3.9 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.9.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.9.1 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.1.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.1.2 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.2 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.13.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.13.2 // indirect
-	github.com/aws/aws-sdk-go-v2/service/sso v1.11.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sso v1.11.2 // indirect
-	github.com/aws/aws-sdk-go-v2/service/sts v1.16.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sts v1.16.2 // indirect
 	github.com/benbjohnson/clock v1.3.0 // indirect
 	github.com/containerd/cgroups v1.0.3 // indirect
 	github.com/containerd/containerd v1.6.0 // indirect

go.sum

@@ -262,42 +262,49 @@ github.com/aws/aws-sdk-go v1.36.29/go.mod h1:hcU610XS61/+aQV88ixoOzUoG7v3b31pl2z
 github.com/aws/aws-sdk-go v1.37.0/go.mod h1:hcU610XS61/+aQV88ixoOzUoG7v3b31pl2zKMmprdro=
 github.com/aws/aws-sdk-go v1.38.49/go.mod h1:hcU610XS61/+aQV88ixoOzUoG7v3b31pl2zKMmprdro=
 github.com/aws/aws-sdk-go-v2 v0.18.0/go.mod h1:JWVYvqSMppoMJC0x5wdwiImzgXTI9FuZwxzkQq9wy+g=
-github.com/aws/aws-sdk-go-v2 v1.15.0 h1:f9kWLNfyCzCB43eupDAk3/XgJ2EpgktiySD6leqs0js=
 github.com/aws/aws-sdk-go-v2 v1.15.0/go.mod h1:lJYcuZZEHWNIb6ugJjbQY1fykdoobWbOS7kJYb4APoI=
-github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.4.0 h1:J/tiyHbl07LL4/1i0rFrW5pbLMvo7M6JrekBUNpLeT4=
+github.com/aws/aws-sdk-go-v2 v1.16.1 h1:udzee98w8H6ikRgtFdVN9JzzYEbi/quFfSvduZETJIU=
-github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.4.0/go.mod h1:ohZjRmiToJ4NybwWTGOCbzlUQU8dxSHxYKzuX7k5l6Y=
+github.com/aws/aws-sdk-go-v2 v1.16.1/go.mod h1:ytwTPBG6fXTZLxxeeCCWj2/EMYp/xDUgX+OET6TLNNU=
-github.com/aws/aws-sdk-go-v2/config v1.15.0 h1:cibCYF2c2uq0lsbu0Ggbg8RuGeiHCmXwUlTMS77CiK4=
+github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.4.1 h1:SdK4Ppk5IzLs64ZMvr6MrSficMtjY2oS0WOORXTlxwU=
-github.com/aws/aws-sdk-go-v2/config v1.15.0/go.mod h1:NccaLq2Z9doMmeQXHQRrt2rm+2FbkrcPvfdbCaQn5hY=
+github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.4.1/go.mod h1:n8Bs1ElDD2wJ9kCRTczA83gYbBmjSwZp3umc6zF4EeM=
-github.com/aws/aws-sdk-go-v2/credentials v1.10.0 h1:M/FFpf2w31F7xqJqJLgiM0mFpLOtBvwZggORr6QCpo8=
+github.com/aws/aws-sdk-go-v2/config v1.15.2 h1:4oGcm1yqqtTc2Z8YpwehwjSiBA3TR0iZbFCgNlXcVFQ=
-github.com/aws/aws-sdk-go-v2/credentials v1.10.0/go.mod h1:HWJMr4ut5X+Lt/7epc7I6Llg5QIcoFHKAeIzw32t6EE=
+github.com/aws/aws-sdk-go-v2/config v1.15.2/go.mod h1:S1p1xf7DGVp0srNq0BakyxfirOldPQeDVlx7+fllyok=
-github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.0 h1:gUlb+I7NwDtqJUIRcFYDiheYa97PdVHG/5Iz+SwdoHE=
+github.com/aws/aws-sdk-go-v2/credentials v1.11.1 h1:uR323+M7ca3v2GKXbFSwWbNA3kLjjFzaalL6W4rpB9s=
-github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.0/go.mod h1:prX26x9rmLwkEE1VVCelQOQgRN9sOVIssgowIJ270SE=
+github.com/aws/aws-sdk-go-v2/credentials v1.11.1/go.mod h1:pYrHWfKUoWTmbr+xTf6ZoWeyyvLAQ5BPT3aL+nKlTpE=
-github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.6 h1:xiGjGVQsem2cxoIX61uRGy+Jux2s9C/kKbTrWLdrU54=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.2 h1:+AULPOLHEDjH2TcNKpixl4gt26hFOdlUuuisZUBFczA=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.12.2/go.mod h1:jmsqNRVo2XlUTNXG/NF7hM7o2gd2jhfg8vdJ135d4XA=
+github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.11.2 h1:PTFTblDWY/HdQ6ix5+to1uLARgLLuYbzKGLQnIdE5Us=
+github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.11.2/go.mod h1:j4OwU2Gb7yaQaidJRpdlIRYX93jBCWhVYIgMlPjf89o=
 github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.6/go.mod h1:SSPEdf9spsFgJyhjrXvawfpyzrXHBCUe+2eQ1CjC1Ak=
-github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.0 h1:bt3zw79tm209glISdMRCIVRCwvSDXxgAxh5KWe2qHkY=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.8 h1:CDaO90VZVBAL1sK87S5oSPIrp7yZqORv1hPIi2UsTMk=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.8/go.mod h1:LnTQMTqbKsbtt+UI5+wPsB7jedW+2ZgozoPG8k6cMxg=
 github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.0/go.mod h1:viTrxhAuejD+LszDahzAE2x40YjYWhMqzHxv2ZiWaME=
-github.com/aws/aws-sdk-go-v2/internal/ini v1.3.7 h1:QOMEP8jnO8sm0SX/4G7dbaIq2eEP2wcWEsF0jzrXLJc=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.2 h1:XXR3cdOcKRCTZf6ctcqpMf+go1BdzTm6+T9Ul5zxcMI=
-github.com/aws/aws-sdk-go-v2/internal/ini v1.3.7/go.mod h1:P5sjYYf2nc5dE6cZIzEMsVtq6XeLD7c4rM+kQJPrByA=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.2/go.mod h1:1x4ZP3Z8odssdhuLI+/1Tqw6Pt/VAaP4Tr8EUxHvPXE=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.3.9 h1:8umg6LSQ/b0+ZTq+Ro8K7VLGVwd7kiYQtIACpf2N/Yo=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.3.9/go.mod h1:kASRBzoVW4I8KUmGCjsowAqVor9QU9DuTUABVducrTY=
 github.com/aws/aws-sdk-go-v2/service/ec2 v1.32.0 h1:0Vbs1G2zV7uvBhMj7o/igTzAg1/roh4ksgIr5oRKFIo=
 github.com/aws/aws-sdk-go-v2/service/ec2 v1.32.0/go.mod h1:Z8942YP2VgLQpgPCx06iXCrOt7mxxCe0dESCm9FFhgs=
-github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.9.0 h1:uhb7moM7VjqIEpWzTpCvceLDSwrWpaleXm39OnVjuLE=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.9.1 h1:T4pFel53bkHjL2mMo+4DKE6r6AuoZnM0fg7k1/ratr4=
-github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.9.0/go.mod h1:pA2St3Pu2Ldy6fBPY45Azoh1WBG4oS7eIKOd4XN7Meg=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.9.1/go.mod h1:GeUru+8VzrTXV/83XyMJ80KpH8xO89VPoUileyNQ+tc=
-github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.1.0 h1:IhiVUezzcKlszx6wXSDQYDjEn/bIO6Mc73uNQ1YfTmA=
+github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.1.2 h1:VoMBHtQZygRs8mcQNDrfmn09vFH2ccjf79nGJ0xuUfo=
-github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.1.0/go.mod h1:kLKc4lo+XKlMhENIpKbp7dCePpyUqUG1PqGIAXoxwNE=
+github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.1.2/go.mod h1:2Fzbfwkx7z4yue1Lz6KDSKG84UpOcUKFl3VAtSF/gcg=
-github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.0 h1:YQ3fTXACo7xeAqg0NiqcCmBOXJruUfh+4+O2qxF2EjQ=
 github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.0/go.mod h1:R31ot6BgESRCIoxwfKtIHzZMo/vsZn2un81g9BJ4nmo=
-github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.13.0 h1:i+7ve93k5G0S2xWBu60CKtmzU5RjBj9g7fcSypQNLR0=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.2 h1:RrN7V0r8+lUUKZM4OAoCOIZqjPLZPOl6wuwMd2QIryI=
-github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.13.0/go.mod h1:L8EoTDLnnN2zL7MQPhyfCbmiZqEs8Cw7+1d9RlLXT5s=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.9.2/go.mod h1:7hwSi01X5Yj9H0qLQljrn8OSdLwwSym1aQCfGn1tDQQ=
+github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.13.2 h1:yxr9h06slG9fdVmO3CpBVuFVD73AeUHLmBxhCr3T3+E=
+github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.13.2/go.mod h1:rvV/Jr4T8H3kMMw/9fFQw9kxqb70YKihA0oWuUFd3K8=
 github.com/aws/aws-sdk-go-v2/service/kms v1.16.0 h1:C33c+TSGU85CXcGi+WGv6Tc8o4QHTtM1cWQNQiTrp3k=
 github.com/aws/aws-sdk-go-v2/service/kms v1.16.0/go.mod h1:tNTRFAwvy+Nu4jjsxsyYmsv8R8Q2eouijsLUh/3CWsI=
-github.com/aws/aws-sdk-go-v2/service/s3 v1.26.0 h1:6IdBZVY8zod9umkwWrtbH2opcM00eKEmIfZKGUg5ywI=
+github.com/aws/aws-sdk-go-v2/service/s3 v1.26.2 h1:Op/A+5+D1K0bmwH3BStYbp/7iod9Rdfm9898A0qYxLc=
-github.com/aws/aws-sdk-go-v2/service/s3 v1.26.0/go.mod h1:WJzrjAFxq82Hl42oh8HuvwpugTgxmoiJBBX8SLwVs74=
+github.com/aws/aws-sdk-go-v2/service/s3 v1.26.2/go.mod h1:Ao1W746VIMdV1WhEkjeVa5JzlaE1JkxJ46facHX9kzs=
-github.com/aws/aws-sdk-go-v2/service/sso v1.11.0 h1:gZLEXLH6NiU8Y52nRhK1jA+9oz7LZzBK242fi/ziXa4=
+github.com/aws/aws-sdk-go-v2/service/sso v1.11.2 h1:8fVz1c9B/63w7O0kxbrCTT69iV4DgXnFumarPCZ3Cns=
-github.com/aws/aws-sdk-go-v2/service/sso v1.11.0/go.mod h1:d1WcT0OjggjQCAdOkph8ijkr5sUwk1IH/VenOn7W1PU=
+github.com/aws/aws-sdk-go-v2/service/sso v1.11.2/go.mod h1:GdCj3+FzI3D5tauOzz8n3YjN70XvgZz82PVVtJXmDds=
-github.com/aws/aws-sdk-go-v2/service/sts v1.16.0 h1:0+X/rJ2+DTBKWbUsn7WtF0JvNk/fRf928vkFsXkbbZs=
+github.com/aws/aws-sdk-go-v2/service/sts v1.16.2 h1:qgK5htfKByTiPxS/diZ/mTCfDwGAVuyjRdqu6VoCh80=
-github.com/aws/aws-sdk-go-v2/service/sts v1.16.0/go.mod h1:+8k4H2ASUZZXmjx/s3DFLo9tGBb44lkz3XcgfypJY7s=
+github.com/aws/aws-sdk-go-v2/service/sts v1.16.2/go.mod h1:RoMljzynmRe3jyOsRgqIMTzyhpAv6XNxu549M1X4Mdo=
-github.com/aws/smithy-go v1.11.1 h1:IQ+lPZVkSM3FRtyaDox41R8YS6iwPMYIreejOgPW49g=
 github.com/aws/smithy-go v1.11.1/go.mod h1:3xHYmszWVx2c0kIwQeEVf9uSm4fYZt67FBJnwub1bgM=
+github.com/aws/smithy-go v1.11.2 h1:eG/N+CcUMAvsdffgMvjMKwfyDzIkjM6pfxMJ8Mzc6mE=
+github.com/aws/smithy-go v1.11.2/go.mod h1:3xHYmszWVx2c0kIwQeEVf9uSm4fYZt67FBJnwub1bgM=
 github.com/aybabtme/rgbterm v0.0.0-20170906152045-cc83f3b3ce59/go.mod h1:q/89r3U2H7sSsE2t6Kca0lfwTK8JdoNGS/yzM/4iH5I=
 github.com/benbjohnson/clock v1.0.3/go.mod h1:bGMdMPoPVvcYyt1gHDf4J2KE153Yf9BuiUKYMaxlTDM=
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
@@ -1006,8 +1013,6 @@ github.com/mailru/easyjson v0.7.6/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJ
 github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
 github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
 github.com/marstr/guid v1.1.0/go.mod h1:74gB1z2wpxxInTG6yaqA7KrtM0NZ+RbrcqDvYHefzho=
-github.com/martinjungblut/go-cryptsetup v0.0.0-20220306213448-685e4930d722 h1:vfx+2bYxFA1H0g1uTbEjJUFqPPyhzCOZvBCIvM+8aZM=
-github.com/martinjungblut/go-cryptsetup v0.0.0-20220306213448-685e4930d722/go.mod h1:gZoZ0+POlM1ge/VUxWpMmZVNPzzMJ7l436CgkQ5+qzU=
 github.com/martinjungblut/go-cryptsetup v0.0.0-20220317181052-e70d6b615049 h1:RhjbYE5voarNcN87XH0A4RWEPcW5exQ+w4WYPKgqT1I=
 github.com/martinjungblut/go-cryptsetup v0.0.0-20220317181052-e70d6b615049/go.mod h1:gZoZ0+POlM1ge/VUxWpMmZVNPzzMJ7l436CgkQ5+qzU=
 github.com/mattn/go-colorable v0.0.9/go.mod h1:9vuHe8Xs5qXnSaW/c/ABM9alt+Vo+STaOChaDxuIBZU=

kms/storage/azurestorage.go

@@ -6,23 +6,42 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
+	"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
 	"github.com/edgelesssys/constellation/kms/config"
 )
 
+type azureContainerAPI interface {
+	Create(ctx context.Context, options *azblob.CreateContainerOptions) (azblob.ContainerCreateResponse, error)
+	NewBlockBlobClient(blobName string) azureBlobAPI
+}
+
+type azureBlobAPI interface {
+	DownloadBlobToWriterAt(ctx context.Context, offset int64, count int64, writer io.WriterAt, o azblob.HighLevelDownloadFromBlobOptions) error
+	Upload(ctx context.Context, body io.ReadSeekCloser, options *azblob.UploadBlockBlobOptions) (azblob.BlockBlobUploadResponse, error)
+}
+
+type wrappedAzureClient struct {
+	azblob.ContainerClient
+}
+
+func (c wrappedAzureClient) NewBlockBlobClient(blobName string) azureBlobAPI {
+	return c.ContainerClient.NewBlockBlobClient(blobName)
+}
+
 // AzureStorage is an implementation of the Storage interface, storing keys in the Azure Blob Store.
 type AzureStorage struct {
-	client azblob.ContainerClient
+	newClient        func(ctx context.Context, connectionString, containerName string, opts *azblob.ClientOptions) (azureContainerAPI, error)
-	opts   *AzureOpts
+	connectionString string
+	containerName    string
+	opts             *AzureOpts
 }
 
 // AzureOpts are additional options to be used when interacting with the Azure API.
 type AzureOpts struct {
-	download *azblob.DownloadBlobOptions
+	upload  *azblob.UploadBlockBlobOptions
-	upload   *azblob.UploadBlockBlobOptions
+	service *azblob.ClientOptions
-	service  *azblob.ClientOptions
 }
 
 // NewAzureStorage initializes a storage client using Azure's Blob Storage: https://azure.microsoft.com/en-us/services/storage/blobs/
@@ -34,28 +53,40 @@ func NewAzureStorage(ctx context.Context, connectionString, containerName string
 	if opts == nil {
 		opts = &AzureOpts{}
 	}
-	service, err := azblob.NewServiceClientFromConnectionString(connectionString, opts.service)
+
-	if err != nil {
+	s := &AzureStorage{
-		return nil, fmt.Errorf("creating storage client from connection string: %w", err)
+		newClient:        azureContainerClientFactory,
+		connectionString: connectionString,
+		containerName:    containerName,
+		opts:             opts,
 	}
-	client := service.NewContainerClient(containerName)
 
 	// Try to create a new storage container, continue if it already exists
-	_, err = client.Create(ctx, &azblob.CreateContainerOptions{
+	if err := s.createContainerOrContinue(ctx); err != nil {
-		Metadata: config.StorageTags,
+		return nil, err
-	})
-	if (err != nil) && !strings.Contains(err.Error(), string(azblob.StorageErrorCodeContainerAlreadyExists)) {
-		return nil, fmt.Errorf("creating storage container: %w", err)
 	}
 
-	return &AzureStorage{client: client, opts: opts}, nil
+	return s, nil
 }
 
 // Get returns a DEK from from Azure Blob Storage by key ID.
 func (s *AzureStorage) Get(ctx context.Context, keyID string) ([]byte, error) {
-	client := s.client.NewBlockBlobClient(keyID)
+	client, err := s.newBlobClient(ctx, keyID)
-	res, err := client.Download(ctx, s.opts.download)
 	if err != nil {
+		return nil, err
+	}
+
+	// the Azure SDK requires an io.WriterAt, the AWS SDK provides a utility function to create one from a byte slice
+	keyBuffer := manager.NewWriteAtBuffer([]byte{})
+
+	opts := azblob.HighLevelDownloadFromBlobOptions{
+		RetryReaderOptionsPerBlock: azblob.RetryReaderOptions{
+			MaxRetryRequests:       5,
+			TreatEarlyCloseAsError: true,
+		},
+	}
+
+	if err := client.DownloadBlobToWriterAt(ctx, 0, 0, keyBuffer, opts); err != nil {
 		var storeErr *azblob.StorageError
 		if errors.As(err, &storeErr) && (storeErr.ErrorCode == azblob.StorageErrorCodeBlobNotFound) {
 			return nil, ErrDEKUnset
@@ -63,36 +94,63 @@ func (s *AzureStorage) Get(ctx context.Context, keyID string) ([]byte, error) {
 		return nil, fmt.Errorf("downloading DEK from storage: %w", err)
 	}
 
-	key := &bytes.Buffer{}
+	return keyBuffer.Bytes(), nil
-	reader := res.Body(&azblob.RetryReaderOptions{MaxRetryRequests: 5, TreatEarlyCloseAsError: true})
-	defer reader.Close()
-	_, err = key.ReadFrom(reader)
-	if err != nil {
-		return nil, fmt.Errorf("downloading DEK from storage: %w", err)
-	}
-
-	return key.Bytes(), nil
 }
 
 // Put saves a DEK to Azure Blob Storage by key ID.
 func (s *AzureStorage) Put(ctx context.Context, keyID string, encDEK []byte) error {
-	client := s.client.NewBlockBlobClient(keyID)
+	client, err := s.newBlobClient(ctx, keyID)
-	if _, err := client.Upload(ctx, newNopCloser(bytes.NewReader(encDEK)), s.opts.upload); err != nil {
+	if err != nil {
+		return err
+	}
+
+	if _, err := client.Upload(ctx, readSeekNopCloser{bytes.NewReader(encDEK)}, s.opts.upload); err != nil {
 		return fmt.Errorf("uploading DEK to storage: %w", err)
 	}
 	return nil
 }
 
-// nopCloser is a wrapper for io.ReadSeeker implementing the Close method. This is required by the Azure SDK.
+// createContainerOrContinue creates a new storage container if necessary, or continues if it already exists.
-type nopCloser struct {
+func (s *AzureStorage) createContainerOrContinue(ctx context.Context) error {
+	client, err := s.newClient(ctx, s.connectionString, s.containerName, s.opts.service)
+	if err != nil {
+		return err
+	}
+
+	var storeErr *azblob.StorageError
+	_, err = client.Create(ctx, &azblob.CreateContainerOptions{
+		Metadata: config.StorageTags,
+	})
+	if (err == nil) || (errors.As(err, &storeErr) && (storeErr.ErrorCode == azblob.StorageErrorCodeContainerAlreadyExists)) {
+		return nil
+	}
+
+	return fmt.Errorf("creating storage container: %w", err)
+}
+
+// newBlobClient is a convenience function to create BlockBlobClients.
+func (s *AzureStorage) newBlobClient(ctx context.Context, blobName string) (azureBlobAPI, error) {
+	c, err := s.newClient(ctx, s.connectionString, s.containerName, s.opts.service)
+	if err != nil {
+		return nil, err
+	}
+	return c.NewBlockBlobClient(blobName), nil
+}
+
+func azureContainerClientFactory(ctx context.Context, connectionString, containerName string, opts *azblob.ClientOptions) (azureContainerAPI, error) {
+	service, err := azblob.NewServiceClientFromConnectionString(connectionString, opts)
+	if err != nil {
+		return nil, fmt.Errorf("creating storage client from connection string: %w", err)
+	}
+
+	return wrappedAzureClient{service.NewContainerClient(containerName)}, nil
+}
+
+// readSeekNopCloser is a wrapper for io.ReadSeeker implementing the Close method. This is required by the Azure SDK.
+type readSeekNopCloser struct {
 	io.ReadSeeker
 }
 
-func (n nopCloser) Close() error {
+func (n readSeekNopCloser) Close() error {
 	return nil
 }
-
-// newNopCloser returns a ReadSeekCloser with a no-op close method wrapping the provided io.ReadSeeker.
-func newNopCloser(rs io.ReadSeeker) io.ReadSeekCloser {
-	return nopCloser{rs}
-}

kms/storage/azurestorage_test.go

@@ -0,0 +1,213 @@
+package storage
+
+import (
+	"context"
+	"errors"
+	"io"
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
+	"github.com/stretchr/testify/assert"
+)
+
+type stubAzureContainerAPI struct {
+	newClientErr error
+	createErr    error
+	createCalled *bool
+	blockBlobAPI stubAzureBlockBlobAPI
+}
+
+func newStubClientFactory(stub stubAzureContainerAPI) func(ctx context.Context, connectionString, containerName string, opts *azblob.ClientOptions) (azureContainerAPI, error) {
+	return func(ctx context.Context, connectionString, containerName string, opts *azblob.ClientOptions) (azureContainerAPI, error) {
+		return stub, stub.newClientErr
+	}
+}
+
+func (s stubAzureContainerAPI) Create(ctx context.Context, options *azblob.CreateContainerOptions) (azblob.ContainerCreateResponse, error) {
+	*s.createCalled = true
+	return azblob.ContainerCreateResponse{}, s.createErr
+}
+
+func (s stubAzureContainerAPI) NewBlockBlobClient(blobName string) azureBlobAPI {
+	return s.blockBlobAPI
+}
+
+type stubAzureBlockBlobAPI struct {
+	downloadBlobToWriterAtErr  error
+	downloadBlobToWriterOutput []byte
+	uploadErr                  error
+	uploadData                 chan []byte
+}
+
+func (s stubAzureBlockBlobAPI) DownloadBlobToWriterAt(ctx context.Context, offset int64, count int64, writer io.WriterAt, o azblob.HighLevelDownloadFromBlobOptions) error {
+	if _, err := writer.WriteAt(s.downloadBlobToWriterOutput, 0); err != nil {
+		panic(err)
+	}
+	return s.downloadBlobToWriterAtErr
+}
+
+func (s stubAzureBlockBlobAPI) Upload(ctx context.Context, body io.ReadSeekCloser, options *azblob.UploadBlockBlobOptions) (azblob.BlockBlobUploadResponse, error) {
+	res, err := io.ReadAll(body)
+	if err != nil {
+		panic(err)
+	}
+	s.uploadData <- res
+	return azblob.BlockBlobUploadResponse{}, s.uploadErr
+}
+
+func TestAzureGet(t *testing.T) {
+	someErr := errors.New("error")
+
+	testCases := map[string]struct {
+		client      stubAzureContainerAPI
+		unsetError  bool
+		errExpected bool
+	}{
+		"success": {
+			client: stubAzureContainerAPI{
+				blockBlobAPI: stubAzureBlockBlobAPI{downloadBlobToWriterOutput: []byte("test-data")},
+			},
+		},
+		"creating client fails": {
+			client:      stubAzureContainerAPI{newClientErr: someErr},
+			errExpected: true,
+		},
+		"DownloadBlobToBuffer fails": {
+			client: stubAzureContainerAPI{
+				blockBlobAPI: stubAzureBlockBlobAPI{downloadBlobToWriterAtErr: someErr},
+			},
+			errExpected: true,
+		},
+		"BlobNotFound error": {
+			client: stubAzureContainerAPI{
+				blockBlobAPI: stubAzureBlockBlobAPI{
+					downloadBlobToWriterAtErr: &azblob.StorageError{
+						ErrorCode: azblob.StorageErrorCodeBlobNotFound,
+					},
+				},
+			},
+			unsetError:  true,
+			errExpected: true,
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			client := &AzureStorage{
+				newClient:        newStubClientFactory(tc.client),
+				connectionString: "test",
+				containerName:    "test",
+				opts:             &AzureOpts{},
+			}
+
+			out, err := client.Get(context.Background(), "test-key")
+			if tc.errExpected {
+				assert.Error(err)
+
+				if tc.unsetError {
+					assert.ErrorIs(err, ErrDEKUnset)
+				} else {
+					assert.False(errors.Is(err, ErrDEKUnset))
+				}
+
+			} else {
+				assert.NoError(err)
+				assert.Equal(tc.client.blockBlobAPI.downloadBlobToWriterOutput, out)
+			}
+		})
+	}
+}
+
+func TestAzurePut(t *testing.T) {
+	someErr := errors.New("error")
+
+	testCases := map[string]struct {
+		client      stubAzureContainerAPI
+		errExpected bool
+	}{
+		"success": {
+			client: stubAzureContainerAPI{},
+		},
+		"creating client fails": {
+			client:      stubAzureContainerAPI{newClientErr: someErr},
+			errExpected: true,
+		},
+		"Upload fails": {
+			client: stubAzureContainerAPI{
+				blockBlobAPI: stubAzureBlockBlobAPI{uploadErr: someErr},
+			},
+			errExpected: true,
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			testData := []byte{0x1, 0x2, 0x3}
+			tc.client.blockBlobAPI.uploadData = make(chan []byte, len(testData))
+
+			client := &AzureStorage{
+				newClient:        newStubClientFactory(tc.client),
+				connectionString: "test",
+				containerName:    "test",
+				opts:             &AzureOpts{},
+			}
+
+			err := client.Put(context.Background(), "test-key", testData)
+			if tc.errExpected {
+				assert.Error(err)
+			} else {
+				assert.NoError(err)
+				assert.Equal(testData, <-tc.client.blockBlobAPI.uploadData)
+			}
+		})
+	}
+}
+
+func TestCreateContainerOrContinue(t *testing.T) {
+	someErr := errors.New("error")
+	testCases := map[string]struct {
+		client      stubAzureContainerAPI
+		errExpected bool
+	}{
+		"success": {
+			client: stubAzureContainerAPI{},
+		},
+		"container already exists": {
+			client: stubAzureContainerAPI{createErr: &azblob.StorageError{ErrorCode: azblob.StorageErrorCodeContainerAlreadyExists}},
+		},
+		"creating client fails": {
+			client:      stubAzureContainerAPI{newClientErr: someErr},
+			errExpected: true,
+		},
+		"Create fails": {
+			client:      stubAzureContainerAPI{createErr: someErr},
+			errExpected: true,
+		},
+	}
+
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			assert := assert.New(t)
+
+			tc.client.createCalled = new(bool)
+			client := &AzureStorage{
+				newClient:        newStubClientFactory(tc.client),
+				connectionString: "test",
+				containerName:    "test",
+				opts:             &AzureOpts{},
+			}
+
+			err := client.createContainerOrContinue(context.Background())
+			if tc.errExpected {
+				assert.Error(err)
+			} else {
+				assert.NoError(err)
+				assert.True(*tc.client.createCalled)
+			}
+		})
+	}
+}