create a volume downloader and use it in cloudinit instead of custom code

Author: dmacvicar

libvirt/cloudinit_def.go

@@ -1,7 +1,6 @@
 package libvirt
 
 import (
-	"bufio"
 	"context"
 	"encoding/xml"
 	"errors"
@@ -109,7 +108,7 @@ func (ci *defCloudInit) UploadIso(ctx context.Context, client *Client, iso strin
 	}
 
 	// upload ISO file
-	err = img.Import(newCopier(virConn, &volume, uint64(size)), volumeDef)
+	err = img.Import(newVolumeUploader(virConn, &volume, uint64(size)), volumeDef)
 	if err != nil {
 		return "", fmt.Errorf("error while uploading cloudinit %s: %w", img.String(), err)
 	}
@@ -301,36 +300,16 @@ func readIso9660File(file os.FileInfo) ([]byte, error) {
 // Returns a pointer to the ISO file. Note well: you have to close this file
 // pointer when you are done.
 func downloadISO(virConn *libvirt.Libvirt, volume libvirt.StorageVol) (*os.File, error) {
-	// get Volume info (required to get size later)
-	_, size, _, err := virConn.StorageVolGetInfo(volume)
-	if err != nil {
-		return nil, fmt.Errorf("error retrieving info for volume: %w", err)
-	}
-
-	// create tmp file for the ISO
 	tmpFile, err := os.CreateTemp("", "cloudinit")
 	if err != nil {
 		return nil, fmt.Errorf("cannot create tmp file: %w", err)
 	}
 
-	w := bufio.NewWriterSize(tmpFile, int(size))
-
-	// download ISO file
-	if err := virConn.StorageVolDownload(volume, w, 0, size, 0); err != nil {
+	downloader := newVolumeDownloader(virConn, &volume)
+	if err := downloader(tmpFile); err != nil {
 		return tmpFile, fmt.Errorf("error while downloading volume: %w", err)
 	}
 
-	bytesCopied := w.Buffered()
-	err = w.Flush()
-	if err != nil {
-		return tmpFile, fmt.Errorf("error while copying remote volume to local disk: %w", err)
-	}
-
-	log.Printf("%d bytes downloaded", bytesCopied)
-	if uint64(bytesCopied) != size {
-		return tmpFile, fmt.Errorf("error while copying remote volume to local disk, bytesCopied %d !=  %d volume.size", bytesCopied, size)
-	}
-
 	if _, err := tmpFile.Seek(0, 0); err != nil {
 		return nil, err
 	}

libvirt/coreos_ignition_def.go

@@ -98,7 +98,7 @@ func (ign *defIgnition) CreateAndUpload(ctx context.Context, client *Client) (st
 	}
 
 	// upload ignition file
-	err = img.Import(newCopier(virConn, &volume, volumeDef.Capacity.Value), volumeDef)
+	err = img.Import(newVolumeUploader(virConn, &volume, volumeDef.Capacity.Value), volumeDef)
 	if err != nil {
 		return "", fmt.Errorf("error while uploading ignition file %s: %w", img.String(), err)
 	}

libvirt/resource_libvirt_volume.go

@@ -256,7 +256,7 @@ be smaller than the backing store specified with
 
 	// upload source if present
 	if _, ok := d.GetOk("source"); ok {
-		err = img.Import(newCopier(virConn, &volume, volumeDef.Capacity.Value), volumeDef)
+		err = img.Import(newVolumeUploader(virConn, &volume, volumeDef.Capacity.Value), volumeDef)
 		if err != nil {
 			//  don't save volume ID  in case of error. This will taint the volume after.
 			// If we don't throw away the id, we will keep instead a broken volume.

libvirt/utils_volume.go

@@ -13,20 +13,36 @@ import (
 )
 
 const (
-	copierBufferSize = 4 * 1024 * 1024
+	uploaderBufferSize = 4 * 1024 * 1024
 )
 
-func newCopier(virConn *libvirt.Libvirt, volume *libvirt.StorageVol, size uint64) func(src io.Reader) error {
-	copier := func(src io.Reader) error {
+func newVolumeUploader(virConn *libvirt.Libvirt, volume *libvirt.StorageVol, size uint64) func(src io.Reader) error {
+	return func(src io.Reader) error {
 		start := time.Now()
-		if err := virConn.StorageVolUpload(*volume, bufio.NewReaderSize(src, copierBufferSize), 0, size, 0); err != nil {
+		if err := virConn.StorageVolUpload(*volume, bufio.NewReaderSize(src, uploaderBufferSize), 0, size, 0); err != nil {
 			return fmt.Errorf("error while uploading volume %w", err)
 		}
 		log.Printf("[DEBUG] upload took %d ms", time.Since(start).Milliseconds())
 
 		return nil
 	}
-	return copier
+}
+
+// returns a function you can give a writer to download the volume content
+// the function will return the downloaded size
+func newVolumeDownloader(virConn *libvirt.Libvirt, volume *libvirt.StorageVol) func(src io.Writer) error {
+	return func(dst io.Writer) error {
+		start := time.Now()
+
+		bufdst := bufio.NewWriter(dst)
+		if err := virConn.StorageVolDownload(*volume, bufdst, 0, 0, 0); err != nil {
+			return fmt.Errorf("error while downloading volume: %w", err)
+		}
+
+		log.Printf("[DEBUG] download took %d ms", time.Since(start).Milliseconds())
+
+		return bufdst.Flush()
+	}
 }
 
 //nolint:mnd

libvirt/utils_volume_test.go

@@ -1,14 +1,20 @@
 package libvirt
 
 import (
+	"crypto/sha256"
 	"fmt"
+	"io"
 	"os"
+	"path/filepath"
 	"testing"
 	"time"
 
 	libvirt "github.com/digitalocean/go-libvirt"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/acctest"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestTimeFromEpoch(t *testing.T) {
@@ -29,27 +35,17 @@ func TestTimeFromEpoch(t *testing.T) {
 	}
 }
 
-func TestAccUtilsVolume_UploadVolumeCopier(t *testing.T) {
+func TestAccUtilsVolume_UploaderDownloader(t *testing.T) {
 	var volume libvirt.StorageVol
 	randomVolumeResource := acctest.RandStringFromCharSet(10, acctest.CharSetAlpha)
 	randomVolumeName := acctest.RandStringFromCharSet(10, acctest.CharSetAlpha)
 	randomPoolName := acctest.RandStringFromCharSet(10, acctest.CharSetAlpha)
 	randomPoolPath := t.TempDir()
 
-	tmpfile, err := os.CreateTemp(t.TempDir(), "test-image-")
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	// simulate uploading a 1G file usig a sparse file
-	if err := tmpfile.Truncate(1024 * 1024 * 1024); err != nil {
-		t.Fatal(err)
-	}
+	imagePath, err := filepath.Abs("testdata/test.qcow2")
+	require.NoError(t, err)
 
-	defer os.Remove(tmpfile.Name())
-	defer tmpfile.Close()
-
-	url := fmt.Sprintf("file://%s", tmpfile.Name())
+	url := fmt.Sprintf("file://%s", imagePath)
 
 	config := fmt.Sprintf(`
     resource "libvirt_pool" "%s" {
@@ -74,8 +70,32 @@ func TestAccUtilsVolume_UploadVolumeCopier(t *testing.T) {
 					testAccCheckLibvirtVolumeExists("libvirt_volume."+randomVolumeResource, &volume),
 					resource.TestCheckResourceAttr(
 						"libvirt_volume."+randomVolumeResource, "name", randomVolumeName),
+					func(state *terraform.State) error {
+						virConn := testAccProvider.Meta().(*Client).libvirt
+
+						file, err := os.CreateTemp("", "downloader-")
+						require.NoError(t, err)
+
+						defer os.Remove(file.Name())
+						defer file.Close()
+
+						downloader := newVolumeDownloader(virConn, &volume)
+						err = downloader(file)
+						require.NoError(t, err)
+
+						_, err = file.Seek(0, 0)
+						require.NoError(t, err)
+
+						h := sha256.New()
+						_, err = io.Copy(h, file)
+						require.NoError(t, err)
+
+						assert.Equal(t, "0f71acdc66da59b04121b939573bec2e5be78a6cdf829b64142cf0a93a7076f5", fmt.Sprintf("%x", h.Sum(nil)))
+						return nil
+					},
 				),
 			},
 		},
 	})
+
 }

libvirt/volume_image.go

@@ -61,7 +61,7 @@ func (i *localImage) IsQCOW2() (bool, error) {
 	return isQCOW2Header(buf)
 }
 
-func (i *localImage) Import(copier func(io.Reader) error, vol libvirtxml.StorageVolume) error {
+func (i *localImage) Import(uploader func(io.Reader) error, vol libvirtxml.StorageVolume) error {
 	file, err := os.Open(i.path)
 	defer file.Close()
 	if err != nil {
@@ -80,7 +80,7 @@ func (i *localImage) Import(copier func(io.Reader) error, vol libvirtxml.Storage
 		}
 	}
 
-	return copier(file)
+	return uploader(file)
 }
 
 type httpImage struct {
@@ -158,7 +158,7 @@ func (i *httpImage) IsQCOW2() (bool, error) {
 	return isQCOW2Header(header)
 }
 
-func (i *httpImage) Import(copier func(io.Reader) error, vol libvirtxml.StorageVolume) error {
+func (i *httpImage) Import(uploader func(io.Reader) error, vol libvirtxml.StorageVolume) error {
 	// number of download retries on non client errors (eg. 5xx)
 	const maxHTTPRetries int = 3
 	// wait time between retries
@@ -187,7 +187,7 @@ func (i *httpImage) Import(copier func(io.Reader) error, vol libvirtxml.StorageV
 		if response.StatusCode == http.StatusNotModified {
 			return nil
 		} else if response.StatusCode == http.StatusOK {
-			return copier(response.Body)
+			return uploader(response.Body)
 		} else if response.StatusCode < http.StatusInternalServerError {
 			break
 		} else if retryCount < maxHTTPRetries {

libvirt/volume_image_test.go

@@ -123,12 +123,12 @@ func TestLocalImageDownload(t *testing.T) {
 		Mtime: fmt.Sprintf("%d.%d", tmpfileStat.ModTime().Unix(), tmpfileStat.ModTime().Nanosecond()),
 	}
 
-	copier := func(r io.Reader) error {
+	uploader := func(r io.Reader) error {
 		require.FailNow(t, fmt.Sprintf("This should not be run, as image has not changed. url: %s", url))
 		return nil
 	}
 
-	if err = image.Import(copier, vol); err != nil {
+	if err = image.Import(uploader, vol); err != nil {
 		require.NoError(t, err, "As the image was not modified and not copied, no error was expected. url: %s", tmpfile.Name())
 	}
 
@@ -156,7 +156,7 @@ func TestRemoteImageDownloadRetry(t *testing.T) {
 				}))
 	}
 
-	copier := func(r io.Reader) error {
+	uploader := func(r io.Reader) error {
 		_, err := io.ReadAll(r)
 		return err
 	}
@@ -169,7 +169,7 @@ func TestRemoteImageDownloadRetry(t *testing.T) {
 		t.Errorf("Could not create image object: %v", err)
 	}
 	start := time.Now()
-	if err = image.Import(copier, vol); err != nil {
+	if err = image.Import(uploader, vol); err != nil {
 		t.Fatalf("Expected to retry: %v", err)
 	}
 	if time.Since(start).Seconds() < 4 {
@@ -184,7 +184,7 @@ func TestRemoteImageDownloadRetry(t *testing.T) {
 	if err != nil {
 		t.Errorf("Could not create image object: %v", err)
 	}
-	if err = image.Import(copier, vol); err == nil {
+	if err = image.Import(uploader, vol); err == nil {
 		t.Fatalf("Expected %s to fail with status 4xx", server.URL)
 	}
 	if time.Since(start).Seconds() < 2 {
@@ -219,11 +219,11 @@ func TestRemoteImageDownload(t *testing.T) {
 	vol.Target.Timestamps = &libvirtxml.StorageVolumeTargetTimestamps{
 		Mtime: fmt.Sprintf("%d.%d", tmpfileStat.ModTime().Unix(), tmpfileStat.ModTime().Nanosecond()),
 	}
-	copier := func(r io.Reader) error {
+	uploader := func(r io.Reader) error {
 		t.Fatalf("ERROR: starting copy of %s... but the file is the same!", url)
 		return nil
 	}
-	if err = image.Import(copier, vol); err != nil {
+	if err = image.Import(uploader, vol); err != nil {
 		t.Fatalf("Could not copy image from %s: %v", url, err)
 	}
 	t.Log("File not copied because modification time was the same")