From f6ab86e51f3a8b64f20bfe7bf2a9b5f2ee1880f9 Mon Sep 17 00:00:00 2001
From: Kevin Cui <bh@bugs.cc>
Date: Thu, 18 Apr 2024 13:49:57 +0800
Subject: [PATCH] feat(cli): support extend share dir (#43)

https://oomol.atlassian.net/browse/OOM-1349

Signed-off-by: Kevin Cui <bh@bugs.cc>
---
 pkg/cli/cli.go     |  2 ++
 pkg/cli/setup.go   | 19 +++++++++++++++++++
 pkg/vfkit/mount.go | 10 ++++++++++
 pkg/vfkit/vfkit.go |  4 ++++
 4 files changed, 35 insertions(+)

diff --git a/pkg/cli/cli.go b/pkg/cli/cli.go
index 7ed7061..42124c8 100644
--- a/pkg/cli/cli.go
+++ b/pkg/cli/cli.go
@@ -25,6 +25,7 @@ var (
 	bindPID         int
 	powerSaveMode   bool
 	kernelDebug     bool
+	extendShareDir  string
 )
 
 func Parse() {
@@ -44,6 +45,7 @@ func Parse() {
 	flag.IntVar(&bindPID, "bind-pid", 0, "OVM will exit when the bound pid exited")
 	flag.BoolVar(&powerSaveMode, "power-save-mode", false, "Enable power save mode")
 	flag.BoolVar(&kernelDebug, "kernel-debug", false, "Enable kernel debug")
+	flag.StringVar(&extendShareDir, "extend-share-dir", "", "Extends share directory with the guest. e.g. --extend-share-dir=host-tmp:/tmp,host-var:/var")
 
 	flag.Parse()
 
diff --git a/pkg/cli/setup.go b/pkg/cli/setup.go
index f093012..ca0f3c1 100644
--- a/pkg/cli/setup.go
+++ b/pkg/cli/setup.go
@@ -29,6 +29,7 @@ type Context struct {
 	EventSocketPath string
 	PowerSaveMode   bool
 	KernelDebug     bool
+	ExtendShareDir  map[string]string
 
 	Endpoint          string
 	SSHPort           int
@@ -112,6 +113,24 @@ func (c *Context) basic() error {
 		c.LockFile = lockPrefixPath + "/" + hash + "-" + name + ".pid"
 	}
 
+	c.ExtendShareDir = make(map[string]string)
+	if extendShareDir != "" {
+		for _, item := range strings.Split(extendShareDir, ",") {
+			parts := strings.Split(item, ":")
+			if len(parts) != 2 {
+				return fmt.Errorf("invalid extend share dir: %s", item)
+			}
+
+			if info, err := os.Stat(parts[1]); err != nil {
+				return fmt.Errorf("extend share dir %s not exists: %w", parts[1], err)
+			} else if !info.IsDir() {
+				return fmt.Errorf("extend share dir %s is not a directory", parts[1])
+			}
+
+			c.ExtendShareDir[parts[0]] = parts[1]
+		}
+	}
+
 	return nil
 }
 
diff --git a/pkg/vfkit/mount.go b/pkg/vfkit/mount.go
index b3b887a..6e70e54 100644
--- a/pkg/vfkit/mount.go
+++ b/pkg/vfkit/mount.go
@@ -35,6 +35,16 @@ var mounts = &_mounts{
 	},
 }
 
+func (m *_mounts) extend(tag, shareDir string) {
+	for _, fs := range mounts.list {
+		if fs.tag == tag || fs.shareDir == shareDir {
+			return
+		}
+	}
+
+	m.list = append(m.list, fs{tag: tag, shareDir: shareDir})
+}
+
 func (m *_mounts) toVFKit() (devices []config.VirtioDevice) {
 	for _, fs := range m.list {
 		d, _ := config.VirtioFsNew(fs.shareDir, fs.tag)
diff --git a/pkg/vfkit/vfkit.go b/pkg/vfkit/vfkit.go
index a95a238..c793a48 100644
--- a/pkg/vfkit/vfkit.go
+++ b/pkg/vfkit/vfkit.go
@@ -26,6 +26,10 @@ func Run(ctx context.Context, g *errgroup.Group, opt *cli.Context) error {
 	runtime.LockOSThread()
 	defer runtime.UnlockOSThread()
 
+	for tag, dir := range opt.ExtendShareDir {
+		mounts.extend(tag, dir)
+	}
+
 	log, err := logger.New(opt.LogPath, opt.Name+"-vfkit")
 	if err != nil {
 		return fmt.Errorf("create vfkit logger error: %v", err)