incus-mirror/cmd/incusd/daemon_images.go

package main

import (
	"context"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"slices"
	"time"

	incus "github.com/lxc/incus/v7/client"
	internalIO "github.com/lxc/incus/v7/internal/io"
	"github.com/lxc/incus/v7/internal/server/db"
	"github.com/lxc/incus/v7/internal/server/db/cluster"
	"github.com/lxc/incus/v7/internal/server/locking"
	"github.com/lxc/incus/v7/internal/server/operations"
	"github.com/lxc/incus/v7/internal/server/project"
	"github.com/lxc/incus/v7/internal/server/response"
	"github.com/lxc/incus/v7/internal/server/state"
	localUtil "github.com/lxc/incus/v7/internal/server/util"
	internalUtil "github.com/lxc/incus/v7/internal/util"
	"github.com/lxc/incus/v7/internal/version"
	"github.com/lxc/incus/v7/shared/api"
	"github.com/lxc/incus/v7/shared/cancel"
	"github.com/lxc/incus/v7/shared/ioprogress"
	"github.com/lxc/incus/v7/shared/logger"
	"github.com/lxc/incus/v7/shared/units"
	"github.com/lxc/incus/v7/shared/util"
	"github.com/lxc/incus/v7/shared/validate"
)

// imageDownloadArgs used with imageDownload.
type imageDownloadArgs struct {
	ProjectName       string
	Server            string
	Protocol          string
	Certificate       string
	Secret            string
	Alias             string
	Type              string
	SetCached         bool
	PreferCached      bool
	AutoUpdate        bool
	Public            bool
	StoragePool       string
	Budget            int64
	SourceProjectName string
}

// imageOperationLock acquires a lock for operating on an image and returns the unlock function.
func imageOperationLock(ctx context.Context, fingerprint string) (locking.UnlockFunc, error) {
	l := logger.AddContext(logger.Ctx{"fingerprint": fingerprint})
	l.Debug("Acquiring lock for image")
	defer l.Debug("Lock acquired for image")

	return locking.Lock(ctx, fmt.Sprintf("ImageOperation_%s", fingerprint))
}

// imageDownload resolves the image fingerprint and if not in the database, downloads it.
func imageDownload(ctx context.Context, r *http.Request, s *state.State, op *operations.Operation, args *imageDownloadArgs) (*api.Image, bool, error) {
	var err error
	var ctxMap logger.Ctx

	var remote incus.ImageServer
	var info *api.Image

	// Check if the project allows retrieving the image.
	err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
		err := project.AllowImageDownload(tx, args.ProjectName, args.Server)
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return nil, false, err
	}

	// Default protocol is Incus. Copy so that local modifications aren't propagated to args.
	protocol := args.Protocol
	if protocol == "" {
		protocol = "incus"
	}

	// Copy so that local modifications aren't propagated to args.
	alias := args.Alias

	// Default the fingerprint to the alias string we received
	fp := alias

	// Attempt to resolve the alias
	if args.Server != "" && slices.Contains([]string{"incus", "lxd", "oci", "simplestreams"}, protocol) {
		clientArgs := &incus.ConnectionArgs{
			TLSServerCert: args.Certificate,
			UserAgent:     version.UserAgent,
			Proxy:         s.Proxy,
			CachePath:     s.OS.CacheDir,
			CacheExpiry:   time.Hour,
			SkipGetEvents: true,
			SkipGetServer: true,
			TempPath:      internalUtil.VarPath("images"),
		}

		if slices.Contains([]string{"incus", "lxd"}, protocol) {
			// Setup client
			remote, err = incus.ConnectPublicIncus(args.Server, clientArgs)
			if err != nil {
				return nil, false, fmt.Errorf("Failed to connect to the server %q: %w", args.Server, err)
			}

			server, ok := remote.(incus.InstanceServer)
			if ok {
				remote = server.UseProject(args.SourceProjectName)
			}
		} else if protocol == "oci" {
			// Setup OCI client
			remote, err = incus.ConnectOCI(args.Server, clientArgs)
			if err != nil {
				return nil, false, fmt.Errorf("Failed to connect to oci server %q: %w", args.Server, err)
			}
		} else if protocol == "simplestreams" {
			// Setup simplestreams client
			remote, err = incus.ConnectSimpleStreams(args.Server, clientArgs)
			if err != nil {
				return nil, false, fmt.Errorf("Failed to connect to simple streams server %q: %w", args.Server, err)
			}
		}

		// For public images, handle aliases and initial metadata
		if args.Secret == "" {
			// Look for a matching alias.  Note, this err message is lost!
			entry, _, err := remote.GetImageAliasType(args.Type, fp)
			if err == nil {
				fp = entry.Target
			} else if args.Type != "" {
				// If no match was found for the requested type, check whether
				// the other instance type has one to give a more helpful hint.
				otherType := "virtual-machine"
				if args.Type == "virtual-machine" {
					otherType = "container"
				}

				_, _, otherErr := remote.GetImageAliasType(otherType, fp)
				if otherErr == nil {
					return nil, false, fmt.Errorf("The requested image couldn't be found for instance type %q, but one was found for instance type %q", args.Type, otherType)
				}
			}

			// Expand partial fingerprints
			info, _, err = remote.GetImage(fp)
			if err != nil {
				return nil, false, fmt.Errorf("Failed getting remote image info: %w", err)
			}

			fp = info.Fingerprint
		}
	}

	// Ensure we are the only ones operating on this image.
	unlock, err := imageOperationLock(ctx, fp)
	if err != nil {
		return nil, false, err
	}

	defer unlock()

	// If auto-update is on and we're being given the image by
	// alias, try to use a locally cached image matching the given
	// server/protocol/alias, regardless of whether it's stale or
	// not (we can assume that it will be not *too* stale since
	// auto-update is on).
	interval := s.GlobalConfig.ImagesAutoUpdateIntervalHours()

	if args.PreferCached && interval > 0 && alias != fp {
		err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
			for _, architecture := range s.OS.Architectures {
				cachedFingerprint, err := tx.GetCachedImageSourceFingerprint(ctx, args.Server, args.Protocol, alias, args.Type, architecture)
				if err == nil && cachedFingerprint != fp {
					fp = cachedFingerprint
					break
				}
			}

			return nil
		})
		if err != nil {
			return nil, false, err
		}
	}

	var imgInfo *api.Image

	err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
		// Check if the image already exists in this project (partial hash match).
		_, imgInfo, err = tx.GetImage(ctx, fp, cluster.ImageFilter{Project: &args.ProjectName})

		return err
	})
	if err == nil {
		var nodeAddress string

		err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
			// Check if the image is available locally or it's on another node.
			nodeAddress, err = tx.LocateImage(ctx, imgInfo.Fingerprint)

			return err
		})
		if err != nil {
			return nil, false, fmt.Errorf("Failed locating image %q in the cluster: %w", imgInfo.Fingerprint, err)
		}

		if nodeAddress != "" {
			// The image is available from another node, let's try to import it.
			err = instanceImageTransfer(s, r, args.ProjectName, imgInfo.Fingerprint, nodeAddress)
			if err != nil {
				return nil, false, fmt.Errorf("Failed transferring image %q from %q: %w", imgInfo.Fingerprint, nodeAddress, err)
			}

			err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
				// As the image record already exists in the project, just add the node ID to the image.
				return tx.AddImageToLocalNode(ctx, args.ProjectName, imgInfo.Fingerprint)
			})
			if err != nil {
				return nil, false, fmt.Errorf("Failed adding transferred image %q to local cluster member: %w", imgInfo.Fingerprint, err)
			}
		}
	} else if response.IsNotFoundError(err) {
		err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
			// Check if the image already exists in some other project.
			_, imgInfo, err = tx.GetImageFromAnyProject(ctx, fp)

			return err
		})
		if err == nil {
			var nodeAddress string

			err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
				// Check if the image is available locally or it's on another node. Do this before creating
				// the missing DB record so we don't include ourself in the search results.
				nodeAddress, err = tx.LocateImage(ctx, imgInfo.Fingerprint)
				if err != nil {
					return fmt.Errorf("Locate image %q in the cluster: %w", imgInfo.Fingerprint, err)
				}

				// We need to insert the database entry for this project, including the node ID entry.
				err = tx.CreateImage(ctx, args.ProjectName, imgInfo.Fingerprint, imgInfo.Filename, imgInfo.Size, args.Public, imgInfo.AutoUpdate, imgInfo.Architecture, imgInfo.CreatedAt, imgInfo.ExpiresAt, imgInfo.Properties, imgInfo.Type, nil)
				if err != nil {
					return fmt.Errorf("Failed creating image record for project: %w", err)
				}

				// Mark the image as "cached" if downloading for an instance.
				if args.SetCached {
					err = tx.SetImageCachedAndLastUseDate(ctx, args.ProjectName, imgInfo.Fingerprint, time.Now().UTC())
					if err != nil {
						return fmt.Errorf("Failed setting cached flag and last use date: %w", err)
					}
				}

				var id int

				id, imgInfo, err = tx.GetImage(ctx, fp, cluster.ImageFilter{Project: &args.ProjectName})
				if err != nil {
					return err
				}

				return tx.CreateImageSource(ctx, id, args.Server, args.Protocol, args.Certificate, alias)
			})
			if err != nil {
				return nil, false, err
			}

			// Transfer image if needed (after database record has been created above).
			if nodeAddress != "" {
				// The image is available from another node, let's try to import it.
				err = instanceImageTransfer(s, r, args.ProjectName, info.Fingerprint, nodeAddress)
				if err != nil {
					return nil, false, fmt.Errorf("Failed transferring image: %w", err)
				}
			}
		}
	}

	if imgInfo != nil {
		info = imgInfo
		ctxMap = logger.Ctx{"fingerprint": info.Fingerprint}
		logger.Debug("Image already exists in the DB", ctxMap)

		// If not requested in a particular pool, we're done.
		if args.StoragePool == "" {
			return info, false, nil
		}

		ctxMap["pool"] = args.StoragePool

		var poolID int64
		var poolIDs []int64

		err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
			// Get the ID of the target storage pool.
			poolID, err = tx.GetStoragePoolID(ctx, args.StoragePool)
			if err != nil {
				return err
			}

			// Check if the image is already in the pool.
			poolIDs, err = tx.GetPoolsWithImage(ctx, info.Fingerprint)

			return err
		})
		if err != nil {
			return nil, false, err
		}

		if slices.Contains(poolIDs, poolID) {
			logger.Debug("Image already exists on storage pool", ctxMap)
			return info, false, nil
		}

		// Import the image in the pool.
		logger.Debug("Image does not exist on storage pool", ctxMap)

		err = imageCreateInPool(s, info, args.StoragePool)
		if err != nil {
			ctxMap["err"] = err
			logger.Debug("Failed to create image on storage pool", ctxMap)
			return nil, false, fmt.Errorf("Failed to create image %q on storage pool %q: %w", info.Fingerprint, args.StoragePool, err)
		}

		logger.Debug("Created image on storage pool", ctxMap)
		return info, false, nil
	}

	// Begin downloading
	if op == nil {
		ctxMap = logger.Ctx{"alias": alias, "server": args.Server}
	} else {
		ctxMap = logger.Ctx{"trigger": op.URL(), "fingerprint": fp, "operation": op.ID(), "alias": alias, "server": args.Server}
	}

	logger.Info("Downloading image", ctxMap)

	// The fingerprint is used as a file name, so reject anything that isn't a
	// partial or full hex fingerprint to avoid path traversal. It's re-validated
	// as a full SHA-256 once expanded from the remote below.
	match, _ := regexp.MatchString("^[0-9a-f]{1,64}$", fp)
	if !match {
		return nil, false, errors.New("Invalid image fingerprint")
	}

	// Cleanup any leftover from a past attempt
	destDir := internalUtil.VarPath("images")
	destName := filepath.Join(destDir, fp)

	failure := true
	cleanup := func() {
		if failure {
			_ = os.Remove(destName)
			_ = os.Remove(destName + ".rootfs")
		}
	}
	defer cleanup()

	// Setup a progress handler
	progress := func(progress ioprogress.ProgressData) {
		if op == nil {
			return
		}

		meta := op.Metadata()
		if meta == nil {
			meta = make(map[string]any)
		}

		if meta["download_progress"] != progress.Text {
			_ = op.ExtendMetadata(map[string]any{"download_progress": progress.Text})
		}
	}

	var canceler *cancel.HTTPRequestCanceller
	if op != nil {
		canceler = cancel.NewHTTPRequestCanceller()
		op.SetCanceler(canceler)
	}

	if slices.Contains([]string{"incus", "lxd", "oci", "simplestreams"}, protocol) {
		// Create the target files
		dest, err := os.Create(destName)
		if err != nil {
			return nil, false, err
		}

		defer logger.WarnOnErrorExcept(dest.Close, []error{os.ErrClosed}, "Failed to close image file")

		destRootfs, err := os.Create(destName + ".rootfs")
		if err != nil {
			return nil, false, err
		}

		defer logger.WarnOnErrorExcept(destRootfs.Close, []error{os.ErrClosed}, "Failed to close rootfs file")

		// Get the image information
		if info == nil {
			if args.Secret != "" {
				info, _, err = remote.GetPrivateImage(fp, args.Secret)
				if err != nil {
					return nil, false, err
				}

				// Expand the fingerprint now and mark alias string to match
				fp = info.Fingerprint
				alias = info.Fingerprint
			} else {
				info, _, err = remote.GetImage(fp)
				if err != nil {
					return nil, false, err
				}
			}
		}

		// Compatibility with older servers
		if info.Type == "" {
			info.Type = "container"
		}

		if args.Budget > 0 && info.Size > args.Budget {
			return nil, false, fmt.Errorf("Remote image with size %d exceeds allowed bugdget of %d", info.Size, args.Budget)
		}

		// Download the image
		var resp *incus.ImageFileResponse
		request := incus.ImageFileRequest{
			MetaFile:        io.ReadWriteSeeker(dest),
			RootfsFile:      io.ReadWriteSeeker(destRootfs),
			ProgressHandler: progress,
			Canceler:        canceler,
			DeltaSourceRetriever: func(fingerprint string, file string) string {
				path := internalUtil.VarPath("images", fmt.Sprintf("%s.%s", fingerprint, file))
				if util.PathExists(path) {
					return path
				}

				return ""
			},
		}

		if args.Secret != "" {
			resp, err = remote.GetPrivateImageFile(fp, args.Secret, request)
		} else {
			resp, err = remote.GetImageFile(fp, request)
		}

		if err != nil {
			return nil, false, err
		}

		// Truncate down to size
		if resp.RootfsSize > 0 {
			err = destRootfs.Truncate(resp.RootfsSize)
			if err != nil {
				return nil, false, err
			}
		}

		err = dest.Truncate(resp.MetaSize)
		if err != nil {
			return nil, false, err
		}

		// Deal with unified images
		if resp.RootfsSize == 0 {
			err := os.Remove(destName + ".rootfs")
			if err != nil {
				return nil, false, err
			}
		}

		err = dest.Close()
		if err != nil {
			return nil, false, err
		}

		err = destRootfs.Close()
		if err != nil {
			return nil, false, err
		}
	} else if protocol == "direct" {
		// Setup HTTP client
		httpClient, err := localUtil.HTTPClient(args.Certificate, s.Proxy)
		if err != nil {
			return nil, false, err
		}

		// Use relatively short response header timeout so as not to hold the image lock open too long.
		httpTransport, ok := httpClient.Transport.(*http.Transport)
		if !ok {
			return nil, false, errors.New("Unexpected HTTP transport type")
		}

		httpTransport.ResponseHeaderTimeout = 30 * time.Second

		req, err := http.NewRequest("GET", args.Server, nil)
		if err != nil {
			return nil, false, err
		}

		req.Header.Set("User-Agent", version.UserAgent)

		// Make the request
		raw, doneCh, err := cancel.CancelableDownload(canceler, httpClient.Do, req)
		if err != nil {
			return nil, false, err
		}

		defer close(doneCh)

		if raw.StatusCode != http.StatusOK {
			return nil, false, fmt.Errorf("Unable to fetch %q: %s", args.Server, raw.Status)
		}

		// Progress handler
		body := &ioprogress.ProgressReader{
			ReadCloser: raw.Body,
			Tracker: &ioprogress.ProgressTracker{
				Length: raw.ContentLength,
				Handler: func(percent int64, speed int64) {
					progress(ioprogress.ProgressData{Text: fmt.Sprintf("%d%% (%s/s)", percent, units.GetByteSizeString(speed, 2))})
				},
			},
		}

		// Create the target files
		f, err := os.Create(destName)
		if err != nil {
			return nil, false, err
		}

		defer logger.WarnOnErrorExcept(f.Close, []error{os.ErrClosed}, "Failed to close image file")

		// Hashing
		hash256 := sha256.New()

		// Download the image
		writer := internalIO.NewQuotaWriter(io.MultiWriter(f, hash256), args.Budget)
		size, err := util.SafeCopy(writer, body)
		if err != nil {
			return nil, false, err
		}

		// Validate hash
		result := fmt.Sprintf("%x", hash256.Sum(nil))
		if result != fp {
			return nil, false, fmt.Errorf("Hash mismatch for %q: %s != %s", args.Server, result, fp)
		}

		// Parse the image
		imageMeta, imageType, err := getImageMetadata(destName)
		if err != nil {
			return nil, false, err
		}

		info = &api.Image{}
		info.Fingerprint = fp
		info.Size = size
		info.Architecture = imageMeta.Architecture
		info.Properties = imageMeta.Properties
		info.Type = imageType

		if imageMeta.CreationDate > 0 {
			info.CreatedAt = time.Unix(imageMeta.CreationDate, 0)
		}

		if imageMeta.ExpiryDate > 0 {
			info.ExpiresAt = time.Unix(imageMeta.ExpiryDate, 0)
		}

		err = f.Close()
		if err != nil {
			return nil, false, err
		}
	} else {
		return nil, false, fmt.Errorf("Unsupported protocol: %v", protocol)
	}

	// Override visibility
	info.Public = args.Public

	// We want to enable auto-update only if we were passed an
	// alias name, so we can figure when the associated
	// fingerprint changes in the remote.
	if alias != fp {
		info.AutoUpdate = args.AutoUpdate
	}

	err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
		// Create the database entry
		return tx.CreateImage(ctx, args.ProjectName, info.Fingerprint, info.Filename, info.Size, info.Public, info.AutoUpdate, info.Architecture, info.CreatedAt, info.ExpiresAt, info.Properties, info.Type, nil)
	})
	if err != nil && api.StatusErrorCheck(err, http.StatusConflict) {
		// Another cluster member created the record concurrently, reuse it and just register this member.
		err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
			_, info, err = tx.GetImage(ctx, info.Fingerprint, cluster.ImageFilter{Project: &args.ProjectName})
			if err != nil {
				return err
			}

			return tx.AddImageToLocalNode(ctx, args.ProjectName, info.Fingerprint)
		})
	}

	if err != nil {
		return nil, false, fmt.Errorf("Failed creating image record: %w", err)
	}

	// Image is in the DB now, don't wipe on-disk files on failure
	failure = false

	// Re-validate the fingerprint as it may have been updated from the remote
	// server's response, and it is used as a file name below.
	err = validate.IsSHA256(fp)
	if err != nil {
		return nil, false, errors.New("Invalid image fingerprint")
	}

	// Check if the image path changed (private images)
	newDestName := filepath.Join(destDir, fp)
	if newDestName != destName {
		err = internalUtil.FileMove(destName, newDestName)
		if err != nil {
			return nil, false, err
		}

		if util.PathExists(destName + ".rootfs") {
			err = internalUtil.FileMove(destName+".rootfs", newDestName+".rootfs")
			if err != nil {
				return nil, false, err
			}
		}
	}

	// Record the image source
	if alias != fp {
		err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
			id, _, err := tx.GetImage(ctx, fp, cluster.ImageFilter{Project: &args.ProjectName})
			if err != nil {
				return err
			}

			return tx.CreateImageSource(ctx, id, args.Server, protocol, args.Certificate, alias)
		})
		if err != nil {
			return nil, false, err
		}
	}

	// Import into the requested storage pool
	if args.StoragePool != "" {
		err = imageCreateInPool(s, info, args.StoragePool)
		if err != nil {
			return nil, false, err
		}
	}

	// Mark the image as "cached" if downloading for an instance
	if args.SetCached {
		err = s.DB.Cluster.Transaction(ctx, func(ctx context.Context, tx *db.ClusterTx) error {
			return tx.SetImageCachedAndLastUseDate(ctx, args.ProjectName, fp, time.Now().UTC())
		})
		if err != nil {
			return nil, false, fmt.Errorf("Failed setting cached flag and last use date: %w", err)
		}
	}

	logger.Info("Image downloaded", ctxMap)

	return info, true, nil
}