package azure
import (
"bufio"
"bytes"
"context"
"crypto/md5"
"errors"
"fmt"
"io"
"net/url"
"os"
"strings"
"sync"
"github.com/Azure/azure-storage-blob-go/azblob"
)
// Credentials contains credentials to connect to your account
// It uses so called "Client credentials", see the official documentation for more information:
// https://docs.microsoft.com/en-us/azure/go/azure-sdk-go-authorization#available-authentication-types-and-methods
type Credentials struct {
StorageAccount string
StorageAccessKey string
}
// ImageMetadata contains information needed to store the image in a proper place.
// In case of Azure cloud storage this includes container name and blob name.
type ImageMetadata struct {
ContainerName string
ImageName string
}
// UploadImage takes the metadata and credentials required to upload the image specified by `fileName`
// It can speed up the upload by using goroutines. The number of parallel goroutines is bounded by
// the `threads` argument.
func UploadImage(credentials Credentials, metadata ImageMetadata, fileName string, threads int) error {
// Azure cannot create an image from a storage blob without .vhd extension
if !strings.HasSuffix(metadata.ImageName, ".vhd") {
metadata.ImageName = metadata.ImageName + ".vhd"
}
// Create a default request pipeline using your storage account name and account key.
credential, err := azblob.NewSharedKeyCredential(credentials.StorageAccount, credentials.StorageAccessKey)
if err != nil {
return fmt.Errorf("cannot create azure credentials: %v", err)
}
p := azblob.NewPipeline(credential, azblob.PipelineOptions{})
// get storage account blob service URL endpoint.
URL, _ := url.Parse(fmt.Sprintf("https://%s.blob.core.windows.net/%s", credentials.StorageAccount, metadata.ContainerName))
// Create a ContainerURL object that wraps the container URL and a request
// pipeline to make requests.
containerURL := azblob.NewContainerURL(*URL, p)
// Create the container, use a never-expiring context
ctx := context.Background()
// Open the image file for reading
imageFile, err := os.Open(fileName)
if err != nil {
return fmt.Errorf("cannot open the image: %v", err)
}
defer imageFile.Close()
// Stat image to get the file size
stat, err := imageFile.Stat()
if err != nil {
return fmt.Errorf("cannot stat the image: %v", err)
}
// Hash the imageFile
imageFileHash := md5.New()
if _, err := io.Copy(imageFileHash, imageFile); err != nil {
return fmt.Errorf("cannot create md5 of the image: %v", err)
}
// Move the cursor back to the start of the imageFile
if _, err := imageFile.Seek(0, 0); err != nil {
return fmt.Errorf("cannot seek the image: %v", err)
}
// Create page blob URL. Page blob is required for VM images
blobURL := newPageBlobURL(containerURL, metadata.ImageName)
_, err = blobURL.Create(ctx, stat.Size(), 0, azblob.BlobHTTPHeaders{}, azblob.Metadata{}, azblob.BlobAccessConditions{})
if err != nil {
return fmt.Errorf("cannot create the blob URL: %v", err)
}
// Wrong MD5 does not seem to have any impact on the upload
_, err = blobURL.SetHTTPHeaders(ctx, azblob.BlobHTTPHeaders{ContentMD5: imageFileHash.Sum(nil)}, azblob.BlobAccessConditions{})
if err != nil {
return fmt.Errorf("cannot set the HTTP headers on the blob URL: %v", err)
}
// Create control variables
// This channel simulates behavior of a semaphore and bounds the number of parallel threads
var semaphore = make(chan int, threads)
// Forward error from goroutine to the caller
var errorInGoroutine = make(chan error, 1)
var counter int64 = 0
// Create buffered reader to speed up the upload
reader := bufio.NewReader(imageFile)
// Run the upload
run := true
var wg sync.WaitGroup
for run {
buffer := make([]byte, azblob.PageBlobMaxUploadPagesBytes)
n, err := reader.Read(buffer)
if err != nil {
if err == io.EOF {
run = false
} else {
return fmt.Errorf("reading the image failed: %v", err)
}
}
if n == 0 {
break
}
wg.Add(1)
semaphore <- 1
go func(counter int64, buffer []byte, n int) {
defer wg.Done()
_, err = blobURL.UploadPages(ctx, counter*azblob.PageBlobMaxUploadPagesBytes, bytes.NewReader(buffer[:n]), azblob.PageBlobAccessConditions{}, nil)
if err != nil {
err = fmt.Errorf("uploading a page failed: %v", err)
// Send the error to the error channel in a non-blocking way. If there is already an error, just discard this one
select {
case errorInGoroutine <- err:
default:
}
}
<-semaphore
}(counter, buffer, n)
counter++
}
// Wait for all goroutines to finish
wg.Wait()
// Check any errors during the transmission using a nonblocking read from the channel
select {
case err := <-errorInGoroutine:
return err
default:
}
// Check properties, specifically MD5 sum of the blob
props, err := blobURL.GetProperties(ctx, azblob.BlobAccessConditions{})
if err != nil {
return fmt.Errorf("getting the properties of the new blob failed: %v", err)
}
var blobChecksum []byte = props.ContentMD5()
var fileChecksum []byte = imageFileHash.Sum(nil)
if !bytes.Equal(blobChecksum, fileChecksum) {
return errors.New("error during image upload. the image seems to be corrupted")
}
return nil
}