From 917a6d782a3f31156c0b590a62833ce544b6ff71 Mon Sep 17 00:00:00 2001
From: Michail Resvanis <mresvani@redhat.com>
Date: Fri, 17 Jan 2025 13:39:48 +0100
Subject: [PATCH] Add support for mirror registries in IBI prepare

Signed-off-by: Michail Resvanis <mresvani@redhat.com>
---
 api/ibiconfig/ibiconfig.go                |  3 +
 internal/prep/prep.go                     |  8 ++-
 lca-cli/ibi-preparation/ibipreparation.go | 82 ++++++++++++++++++++---
 3 files changed, 82 insertions(+), 11 deletions(-)

diff --git a/api/ibiconfig/ibiconfig.go b/api/ibiconfig/ibiconfig.go
index 765ee4537..626035023 100644
--- a/api/ibiconfig/ibiconfig.go
+++ b/api/ibiconfig/ibiconfig.go
@@ -79,6 +79,9 @@ type IBIPrepareConfig struct {
 	// Provide the device id like /dev/by-id/ata-xxxxx
 	InstallationDisk string `json:"installationDisk"`
 
+	// ReleaseRegistry is the container image registry that hosts the OpenShift release image.
+	ReleaseRegistry string `json:"releaseRegistry,omitempty"`
+
 	// PrecacheBestEffort is a flag to enable best effort precaching.
 	// +optional
 	PrecacheBestEffort bool `json:"precacheBestEffort,omitempty"`
diff --git a/internal/prep/prep.go b/internal/prep/prep.go
index 7738a4c20..9decf20a5 100644
--- a/internal/prep/prep.go
+++ b/internal/prep/prep.go
@@ -123,7 +123,7 @@ func getDeploymentFromDeploymentID(deploymentID string) (string, error) {
 }
 
 func SetupStateroot(log logr.Logger, ops ops.Ops, ostreeClient ostreeclient.IClient,
-	rpmOstreeClient rpmostreeclient.IClient, seedImage, expectedVersion, imageListFile string, ibi bool) error {
+	rpmOstreeClient rpmostreeclient.IClient, seedImage, expectedVersion, tmpPath string, ibi bool) error {
 	log.Info("Start setupstateroot")
 
 	defer ops.UnmountAndRemoveImage(seedImage)
@@ -246,10 +246,14 @@ func SetupStateroot(log logr.Logger, ops ops.Ops, ostreeClient ostreeclient.ICli
 		return fmt.Errorf("failed to process etc.deletions: %w", err)
 	}
 
-	if err := common.CopyOutsideChroot(filepath.Join(mountpoint, "containers.list"), imageListFile); err != nil {
+	if err := common.CopyOutsideChroot(filepath.Join(mountpoint, "containers.list"), filepath.Join(tmpPath, "containers.list")); err != nil {
 		return fmt.Errorf("failed to copy image list file: %w", err)
 	}
 
+	if err := common.CopyOutsideChroot(filepath.Join(mountpoint, common.SeedClusterInfoFileName), filepath.Join(tmpPath, common.SeedClusterInfoFileName)); err != nil {
+		return fmt.Errorf("failed to copy %s file: %w", common.SeedClusterInfoFileName, err)
+	}
+
 	log.Info("Stateroot setup done successfully")
 	return nil
 }
diff --git a/lca-cli/ibi-preparation/ibipreparation.go b/lca-cli/ibi-preparation/ibipreparation.go
index 7d17aed0f..596d94264 100644
--- a/lca-cli/ibi-preparation/ibipreparation.go
+++ b/lca-cli/ibi-preparation/ibipreparation.go
@@ -5,8 +5,11 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/containers/image/v5/pkg/sysregistriesv2"
 	"github.com/go-logr/logr"
+	"github.com/pelletier/go-toml"
 	preinstallUtils "github.com/rh-ecosystem-edge/preinstall-utils/pkg"
+	"github.com/samber/lo"
 	"github.com/sirupsen/logrus"
 
 	"github.com/openshift-kni/lifecycle-agent/api/ibiconfig"
@@ -17,12 +20,20 @@ import (
 	"github.com/openshift-kni/lifecycle-agent/internal/prep"
 	"github.com/openshift-kni/lifecycle-agent/lca-cli/ops"
 	rpmostreeclient "github.com/openshift-kni/lifecycle-agent/lca-cli/ostreeclient"
+	"github.com/openshift-kni/lifecycle-agent/lca-cli/seedclusterinfo"
+	"github.com/openshift-kni/lifecycle-agent/utils"
 )
 
 const (
-	imageListFile    = "var/tmp/imageListFile"
-	rhcosOstreeIndex = 1
-	rhcosOstreePath  = "ostree/deploy/rhcos"
+	workspace          = "var/tmp/"
+	registriesConfFile = "etc/containers/registries.conf"
+	rhcosOstreeIndex   = 1
+	rhcosOstreePath    = "ostree/deploy/rhcos"
+)
+
+var (
+	imageListFile = filepath.Join(workspace, "containers.list")
+	seedInfoFile  = filepath.Join(workspace, common.SeedClusterInfoFileName)
 )
 
 type IBIPrepare struct {
@@ -62,11 +73,11 @@ func (i *IBIPrepare) Run() error {
 	common.OstreeDeployPathPrefix = "/mnt/"
 	// Setup state root
 	if err := prep.SetupStateroot(log, i.ops, i.ostreeClient, i.rpmostreeClient,
-		i.config.SeedImage, i.config.SeedVersion, imageListFile, true); err != nil {
+		i.config.SeedImage, i.config.SeedVersion, workspace, true); err != nil {
 		return fmt.Errorf("failed to setup stateroot: %w", err)
 	}
 
-	if err := i.precacheFlow(imageListFile); err != nil {
+	if err := i.precacheFlow(imageListFile, seedInfoFile); err != nil {
 		return fmt.Errorf("failed to precache: %w", err)
 	}
 
@@ -81,15 +92,28 @@ func (i *IBIPrepare) Run() error {
 	return i.shutdownNode()
 }
 
-func (i *IBIPrepare) precacheFlow(imageListFile string) error {
-	// TODO: add support for mirror registry
+func (i *IBIPrepare) precacheFlow(imageListFile, seedInfoFile string) error {
 	if i.config.PrecacheDisabled {
 		i.log.Info("Precache disabled, skipping it")
 		return nil
 	}
 
-	i.log.Info("Precaching imaging")
-	imageList, err := prep.ReadPrecachingList(imageListFile, "", "", false)
+	i.log.Info("Checking seed image info")
+	seedInfo, err := seedclusterinfo.ReadSeedClusterInfoFromFile(seedInfoFile)
+	if err != nil {
+		return fmt.Errorf("failed to read seed info: %s, %w", common.PathOutsideChroot(seedInfoFile), err)
+	}
+	i.log.Info("Collected seed info for precache: ", fmt.Sprintf("%+v", seedInfo))
+
+	i.log.Info("Checking whether to override seed registry")
+	shouldOverrideSeedRegistry, err := i.shouldOverrideSeedRegistry(seedInfo.MirrorRegistryConfigured, seedInfo.ReleaseRegistry)
+	if err != nil {
+		return fmt.Errorf("failed to check ShouldOverrideSeedRegistry %w", err)
+	}
+	i.log.Info("Should override seed registry: ", shouldOverrideSeedRegistry)
+
+	i.log.Info("Precaching images")
+	imageList, err := prep.ReadPrecachingList(imageListFile, i.config.ReleaseRegistry, seedInfo.ReleaseRegistry, shouldOverrideSeedRegistry)
 	if err != nil {
 		err = fmt.Errorf("failed to read pre-caching image file: %s, %w", common.PathOutsideChroot(imageListFile), err)
 		return err
@@ -211,3 +235,43 @@ func (i *IBIPrepare) cleanupRhcosSysroot() error {
 	}
 	return nil
 }
+
+func (i *IBIPrepare) shouldOverrideSeedRegistry(seedMirrorRegistryConfigured bool, seedReleaseRegistry string) (bool, error) {
+	targetMirroredRegistries, err := mirrorRegistrySourceRegistries()
+	if err != nil {
+		return false, err
+	}
+
+	targetMirrorRegistryConfigured := len(targetMirroredRegistries) > 0
+
+	// if the target SNO doesn't have mirror registry configured:
+	//   - and seed SNO has mirror registry configured then we should try to override the registry
+	//   - and seed SNO has no mirror registry configured then we shouldn't try to override the registry
+	if !targetMirrorRegistryConfigured {
+		return seedMirrorRegistryConfigured, nil
+	}
+
+	return !lo.Contains(targetMirroredRegistries, seedReleaseRegistry), nil
+}
+
+func mirrorRegistrySourceRegistries() ([]string, error) {
+	content, err := os.ReadFile(registriesConfFile)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read registry config file: %w", err)
+	}
+
+	config := &sysregistriesv2.V2RegistriesConf{}
+
+	if err := toml.Unmarshal(content, config); err != nil {
+		return nil, fmt.Errorf("failed to parse registry config: %w", err)
+	}
+
+	prefixes := make([]string, 0, len(config.Registries))
+	for _, registry := range config.Registries {
+		if registry.Prefix != "" {
+			prefixes = append(prefixes, utils.ExtractRegistryFromImage(registry.Prefix))
+		}
+	}
+
+	return prefixes, nil
+}