From 58d8ef6103fc792567c6d257a1a25de12c623731 Mon Sep 17 00:00:00 2001
From: Yury Bushmelev <jay4mail@gmail.com>
Date: Thu, 9 Mar 2023 15:34:34 +0400
Subject: [PATCH] Do not prune template-referenced downloads

---
 cmd/limactl/prune.go | 49 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 47 insertions(+), 2 deletions(-)

diff --git a/cmd/limactl/prune.go b/cmd/limactl/prune.go
index 1af8a248565..4cd6b3fbb26 100644
--- a/cmd/limactl/prune.go
+++ b/cmd/limactl/prune.go
@@ -8,7 +8,9 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/lima-vm/lima/pkg/limayaml"
 	"github.com/lima-vm/lima/pkg/store"
+	"github.com/lima-vm/lima/pkg/templatestore"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 )
@@ -75,11 +77,21 @@ func pruneDownloadsAction(cmd *cobra.Command, args []string) error {
 
 	logrus.Infof("Pruning unreferenced downloads in %q", downloadsDir)
 
-	refrd, err := getReferencedDigests()
+	refrd, err := getVMReferencedDigests()
 	if err != nil {
 		return err
 	}
 
+	tmplRefrd, err := getTemplateReferencedDigests()
+	if err != nil {
+		return err
+	}
+
+	// merge template-referenced digests into the VM-referenced
+	for k, v := range tmplRefrd {
+		refrd[k] = v
+	}
+
 	cached, err := getCachedDigests(downloadsDir)
 	if err != nil {
 		return err
@@ -115,7 +127,7 @@ func makeDigest(s string) string {
 }
 
 // Collect all downloads referenced in VM definitions
-func getReferencedDigests() (map[string]string, error) {
+func getVMReferencedDigests() (map[string]string, error) {
 	digests := make(map[string]string)
 	instanceNames, err := store.Instances()
 	if err != nil {
@@ -142,6 +154,39 @@ func getReferencedDigests() (map[string]string, error) {
 	return digests, nil
 }
 
+// Collect all downloads referenced in bundled templates
+func getTemplateReferencedDigests() (map[string]string, error) {
+	digests := make(map[string]string)
+
+	if templates, err := templatestore.Templates(); err == nil {
+		for _, t := range templates {
+			yBytes, err := templatestore.Read(t.Name)
+			if err != nil {
+				logrus.WithError(err).Errorf("Cannot find template %q. Skipping...", t.Name)
+				continue
+			}
+			config, err := limayaml.Load(yBytes, t.Location)
+			if err != nil {
+				logrus.WithError(err).Errorf("Cannot parse template %q. Skipping...", t.Name)
+				continue
+			}
+			for _, img := range config.Images {
+				digests[makeDigest(img.File.Location)] = img.File.Location
+				if img.Kernel != nil {
+					digests[makeDigest(img.Kernel.File.Location)] = img.Kernel.File.Location
+				}
+				if img.Initrd != nil {
+					digests[makeDigest((*img.Initrd).Location)] = (*img.Initrd).Location
+				}
+			}
+			for _, archive := range config.Containerd.Archives {
+				digests[makeDigest(archive.Location)] = archive.Location
+			}
+		}
+	}
+	return digests, nil
+}
+
 // Collect all cached downloads
 func getCachedDigests(downloadsDir string) (map[string]string, error) {
 	digests := make(map[string]string)