diff --git a/.gitignore b/.gitignore index 39c25ac02..1c8887b0a 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ __pycache__ /osbuild-composer /osbuild-worker /osbuild-pipeline +/osbuild-upload-azure diff --git a/Makefile b/Makefile index eaf77ae8d..63d786b4a 100644 --- a/Makefile +++ b/Makefile @@ -3,6 +3,7 @@ build: go build -o osbuild-composer ./cmd/osbuild-composer/ go build -o osbuild-worker ./cmd/osbuild-worker/ go build -o osbuild-pipeline ./cmd/osbuild-pipeline/ + go build -o osbuild-upload-azure ./cmd/osbuild-upload-azure/ .PHONY: install install: diff --git a/cmd/osbuild-upload-azure/main.go b/cmd/osbuild-upload-azure/main.go new file mode 100644 index 000000000..2e6943e53 --- /dev/null +++ b/cmd/osbuild-upload-azure/main.go @@ -0,0 +1,148 @@ +package main + +import ( + "bufio" + "bytes" + "context" + "flag" + "fmt" + "io" + "log" + "net/url" + "os" + "path" + "sync" + + "github.com/Azure/azure-storage-blob-go/azblob" +) + +func handleErrors(err error) { + if err != nil { + if serr, ok := err.(azblob.StorageError); ok { // This error is a Service-specific + switch serr.ServiceCode() { // Compare serviceCode to ServiceCodeXxx constants + case azblob.ServiceCodeContainerAlreadyExists: + // This error is not fatal + fmt.Println("Received 409. Container already exists") + return + } + } + // All other error causes the program to exit + fmt.Println(err) + os.Exit(1) + } +} + +type azureCredentials struct { + storageAccount string + storageAccessKey string +} + +type azureImageMetadata struct { + containerName string + imageName string +} + +func azureUploadImage(credentials azureCredentials, metadata azureImageMetadata, fileName string, threads int) { + // Create a default request pipeline using your storage account name and account key. + credential, err := azblob.NewSharedKeyCredential(credentials.storageAccount, credentials.storageAccessKey) + handleErrors(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) + handleErrors(err) + + // Stat image to get the file size + stat, err := imageFile.Stat() + handleErrors(err) + + // Create page blob URL. Page blob is required for VM images + blobURL := containerURL.NewPageBlobURL(metadata.imageName) + _, err = blobURL.Create(ctx, stat.Size(), 0, azblob.BlobHTTPHeaders{}, azblob.Metadata{}, azblob.BlobAccessConditions{}) + handleErrors(err) + + // Create control variables + // This channel simulates behavior of a semaphore and bounds the number of parallel threads + var semaphore = make(chan int, threads) + var counter int64 = 0 + + // Create buffered reader to speed up the upload + reader := bufio.NewReader(imageFile) + imageSize := stat.Size() + // 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 { + panic(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 { + log.Fatal(err) + } + <-semaphore + }(counter, buffer, n) + fmt.Printf("\rProgress: uploading bytest %d-%d from %d bytes", counter*azblob.PageBlobMaxUploadPagesBytes, counter*azblob.PageBlobMaxUploadPagesBytes+int64(n), imageSize) + counter++ + } + wg.Wait() +} + +func checkStringNotEmpty(variable string, errorMessage string) { + if variable == "" { + log.Fatal(errorMessage) + } +} + +func main() { + var storageAccount string + var storageAccessKey string + var fileName string + var containerName string + var threads int + flag.StringVar(&storageAccount, "storage-account", "", "Azure storage account (mandatory)") + flag.StringVar(&storageAccessKey, "storage-access-key", "", "Azure storage access key (mandatory)") + flag.StringVar(&fileName, "image", "", "image to upload (mandatory)") + flag.StringVar(&containerName, "container", "", "name of storage container (see Azure docs for explanation, mandatory)") + flag.IntVar(&threads, "threads", 16, "number of threads for parallel upload") + flag.Parse() + + checkStringNotEmpty(storageAccount, "You need to specify storage account") + checkStringNotEmpty(storageAccessKey, "You need to specify storage access key") + checkStringNotEmpty(fileName, "You need to specify image file") + checkStringNotEmpty(containerName, "You need to specify container name") + + fmt.Println("Image to upload is:", fileName) + + azureUploadImage(azureCredentials{ + storageAccount: storageAccount, + storageAccessKey: storageAccessKey, + }, azureImageMetadata{ + imageName: path.Base(fileName), + containerName: containerName, + }, fileName, threads) + +} diff --git a/go.mod b/go.mod index 590be8e4b..4e5aecd82 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/osbuild/osbuild-composer go 1.12 require ( + github.com/Azure/azure-storage-blob-go v0.8.0 github.com/coreos/go-systemd v0.0.0-20190719114852-fd7a80b32e1f github.com/gobwas/glob v0.2.3 github.com/google/uuid v1.1.1 diff --git a/go.sum b/go.sum index 3c65591e1..7f8b2a91b 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,7 @@ +github.com/Azure/azure-pipeline-go v0.2.1 h1:OLBdZJ3yvOn2MezlWvbrBMTEUQC72zAftRZOMdj5HYo= +github.com/Azure/azure-pipeline-go v0.2.1/go.mod h1:UGSo8XybXnIGZ3epmeBw7Jdz+HiUVpqIlpz/HKHylF4= +github.com/Azure/azure-storage-blob-go v0.8.0 h1:53qhf0Oxa0nOjgbDeeYPUeyiNmafAFEY95rZLK0Tj6o= +github.com/Azure/azure-storage-blob-go v0.8.0/go.mod h1:lPI3aLPpuLTeUwh1sViKXFxwl2B6teiRqI0deQUvsw0= github.com/coreos/go-systemd v0.0.0-20190719114852-fd7a80b32e1f h1:JOrtw2xFKzlg+cbHpyrpLDmnN1HqhBfnX7WDiW7eG2c= github.com/coreos/go-systemd v0.0.0-20190719114852-fd7a80b32e1f/go.mod h1:F5haX7vjVVG0kc13fIWeqUViNPyEJxv/OmvnBo0Yme4= github.com/gobwas/glob v0.2.3 h1:A4xDbljILXROh+kObIiy5kIaPYD8e96x1tgBhUI5J+Y= @@ -6,3 +10,9 @@ github.com/google/uuid v1.1.1 h1:Gkbcsh/GbpXz7lPftLA3P6TYMwjCLYm83jiFQZF/3gY= github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/julienschmidt/httprouter v1.2.0 h1:TDTW5Yz1mjftljbcKqRcrYhd4XeOoI98t+9HbQbYf7g= github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w= +github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/mattn/go-ieproxy v0.0.0-20190610004146-91bb50d98149 h1:HfxbT6/JcvIljmERptWhwa8XzP7H3T+Z2N26gTsaDaA= +github.com/mattn/go-ieproxy v0.0.0-20190610004146-91bb50d98149/go.mod h1:31jz6HNzdxOmlERGGEc4v/dMssOfmp2p5bT/okiKFFc= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/golang-github-osbuild-composer.spec b/golang-github-osbuild-composer.spec index 6dbed4f6c..c659b895b 100644 --- a/golang-github-osbuild-composer.spec +++ b/golang-github-osbuild-composer.spec @@ -21,6 +21,7 @@ Source0: %{gosource} BuildRequires: systemd-rpm-macros BuildRequires: systemd +BuildRequires: golang-github-azure-storage-blob-devel BuildRequires: golang(github.com/coreos/go-systemd/activation) BuildRequires: golang(github.com/google/uuid) BuildRequires: golang(github.com/julienschmidt/httprouter)