From 670a268bc09dd53f016cb285337d1e7819554c1c Mon Sep 17 00:00:00 2001
From: Justin Alvarez <alvajus@amazon.com>
Date: Thu, 23 Feb 2023 12:20:40 -0500
Subject: [PATCH] feat: mount additional disks when using vz

Signed-off-by: Justin Alvarez <alvajus@amazon.com>
---
 pkg/driver/driver.go       | 10 ++++++++
 pkg/qemu/qemu.go           |  2 +-
 pkg/vz/vm_darwin.go        | 51 ++++++++++++++++++++++++++++++++++++++
 pkg/vz/vz_driver_darwin.go |  6 +++++
 4 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/pkg/driver/driver.go b/pkg/driver/driver.go
index 6415567a11f0..b420c88a3afe 100644
--- a/pkg/driver/driver.go
+++ b/pkg/driver/driver.go
@@ -12,6 +12,8 @@ type Driver interface {
 
 	CreateDisk() error
 
+	InspectDisk(diskName string) (*store.Disk, error)
+
 	Start(_ context.Context) (chan error, error)
 
 	Stop(_ context.Context) error
@@ -36,6 +38,14 @@ func (d *BaseDriver) CreateDisk() error {
 	return nil
 }
 
+func (d *BaseDriver) InspectDisk(diskName string) (*store.Disk, error) {
+	disk, err := store.InspectDisk(diskName)
+	if err != nil {
+		return nil, err
+	}
+	return disk, nil
+}
+
 func (d *BaseDriver) Start(_ context.Context) (chan error, error) {
 	return nil, nil
 }
diff --git a/pkg/qemu/qemu.go b/pkg/qemu/qemu.go
index 4f803e15ef76..aa28e9b9ee34 100644
--- a/pkg/qemu/qemu.go
+++ b/pkg/qemu/qemu.go
@@ -111,7 +111,7 @@ func CreateDataDisk(dir string, size int) error {
 		return err
 	}
 
-	args := []string{"create", "-f", "qcow2", dataDisk, strconv.Itoa(size)}
+	args := []string{"create", "-f", "raw", dataDisk, strconv.Itoa(size)}
 	cmd := exec.Command("qemu-img", args...)
 	if out, err := cmd.CombinedOutput(); err != nil {
 		return fmt.Errorf("failed to run %v: %q: %w", cmd.Args, string(out), err)
diff --git a/pkg/vz/vm_darwin.go b/pkg/vz/vm_darwin.go
index 2c2c6182cc37..501bdd37ee34 100644
--- a/pkg/vz/vm_darwin.go
+++ b/pkg/vz/vm_darwin.go
@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"net"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"strconv"
 	"syscall"
@@ -334,6 +335,56 @@ func attachDisks(driver *driver.BaseDriver, vmConfig *vz.VirtualMachineConfigura
 	}
 	configurations = append(configurations, diffDisk)
 
+	for _, diskName := range driver.Yaml.AdditionalDisks {
+		d, err := driver.InspectDisk(diskName)
+		if err != nil {
+			return fmt.Errorf("failed to run load disk %q: %q", diskName, err)
+		}
+
+		if d.Instance != "" {
+			return fmt.Errorf("failed to run attach disk %q, in use by instance %q", diskName, d.Instance)
+		}
+		logrus.Infof("Mounting disk %q on %q", diskName, d.MountPoint)
+		err = d.Lock(driver.Instance.Dir)
+		if err != nil {
+			return fmt.Errorf("failed to run lock disk %q: %q", diskName, err)
+		}
+		extraDiskPath := filepath.Join(d.Dir, filenames.DataDisk)
+
+		extraDiskFormat, err := imgutil.DetectFormat(extraDiskPath)
+		if err != nil {
+			return fmt.Errorf("failed to run detect disk format %q: %q", diskName, err)
+		}
+		if extraDiskFormat != "raw" {
+			rawPath := fmt.Sprintf("%s.raw", extraDiskPath)
+			if err = imgutil.QCOWToRaw(extraDiskPath, rawPath); err != nil {
+				return fmt.Errorf("failed to convert qcow2 disk %q to raw for vz driver: %w", diskName, err)
+			}
+			cmd := exec.Command("mv", extraDiskPath, fmt.Sprintf("%s.qcow2", extraDiskPath))
+			if out, err := cmd.CombinedOutput(); err != nil {
+				return fmt.Errorf("failed to run %v: %q: %w", cmd.Args, string(out), err)
+			}
+			cmd = exec.Command("mv", rawPath, extraDiskPath)
+			logrus.Info("Running command: %s", cmd.String())
+			if out, err := cmd.CombinedOutput(); err != nil {
+				return fmt.Errorf("failed to run %v: %q: %w", cmd.Args, string(out), err)
+			}
+		}
+
+		if err = validateDiskFormat(extraDiskPath); err != nil {
+			return fmt.Errorf("failed to validate extra disk %q: %w", extraDiskPath, err)
+		}
+		extraDiskPathAttachment, err := vz.NewDiskImageStorageDeviceAttachmentWithCacheAndSync(extraDiskPath, false, vz.DiskImageCachingModeAutomatic, vz.DiskImageSynchronizationModeFsync)
+		if err != nil {
+			return fmt.Errorf("failed to create disk attachment for extra disk %q: %w", extraDiskPath, err)
+		}
+		extraDisk, err := vz.NewVirtioBlockDeviceConfiguration(extraDiskPathAttachment)
+		if err != nil {
+			return fmt.Errorf("failed to create new virtio block device config for extra disk %q: %w", extraDiskPath, err)
+		}
+		configurations = append(configurations, extraDisk)
+	}
+
 	if err = validateDiskFormat(ciDataPath); err != nil {
 		return err
 	}
diff --git a/pkg/vz/vz_driver_darwin.go b/pkg/vz/vz_driver_darwin.go
index 05c4856befe1..84125ad37e69 100644
--- a/pkg/vz/vz_driver_darwin.go
+++ b/pkg/vz/vz_driver_darwin.go
@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/lima-vm/lima/pkg/reflectutil"
+	"github.com/lima-vm/lima/pkg/store"
 
 	"github.com/Code-Hex/vz/v3"
 
@@ -64,6 +65,7 @@ func (l *LimaVzDriver) Validate() error {
 		"PropagateProxyEnv",
 		"CACertificates",
 		"Rosetta",
+		"AdditionalDisks",
 	); len(unknown) > 0 {
 		logrus.Warnf("Ignoring: vmType %s: %+v", *l.Yaml.VMType, unknown)
 	}
@@ -106,6 +108,10 @@ func (l *LimaVzDriver) CreateDisk() error {
 	return nil
 }
 
+func (l *LimaVzDriver) InspectDisk(diskName string) (*store.Disk, error) {
+	return l.BaseDriver.InspectDisk(diskName)
+}
+
 func (l *LimaVzDriver) Start(ctx context.Context) (chan error, error) {
 	logrus.Infof("Starting VZ (hint: to watch the boot progress, see %q)", filepath.Join(l.Instance.Dir, filenames.SerialLog))
 	vm, errCh, err := startVM(ctx, l.BaseDriver)