feat: Support multi-arch images

FIxes #3

handler.go

@@ -48,7 +48,7 @@ func indexAndPush(ctx context.Context, repo string, digest string, registryUrl s
 		return logAndReturnError(ctx, "Remote registry initialization error", err)
 	}
 
-	err = registry.ValidateImageManifest(ctx, repo, digest)
+	digests, err := registry.GetImageDigests(ctx, repo, digest)
 	if err != nil {
 		log.Warn(ctx, fmt.Sprintf("Image manifest validation error: %v", err))
 		// Returning a non error to skip retries
@@ -67,32 +67,35 @@ func indexAndPush(ctx context.Context, repo string, digest string, registryUrl s
 		return logAndReturnError(ctx, "OCI storage initialization error", err)
 	}
 
-	desc, err := registry.Pull(ctx, repo, sociStore, digest)
-	if err != nil {
-		return logAndReturnError(ctx, "Image pull error", err)
-	}
+	for _, digest := range digests {
+		ctx := context.WithValue(ctx, "ImageDigest", digest)
+		desc, err := registry.Pull(ctx, repo, sociStore, digest)
+		if err != nil {
+			return logAndReturnError(ctx, "Image pull error", err)
+		}
 
-	image := images.Image{
-		Name:   repo + "@" + digest,
-		Target: *desc,
-	}
+		image := images.Image{
+			Name:   repo + "@" + digest,
+			Target: *desc,
+		}
 
-	indexDescriptor, err := buildIndex(ctx, dataDir, sociStore, image)
-	if err != nil {
-		if err.Error() == ErrEmptyIndex.Error() {
-			log.Warn(ctx, SkipPushOnEmptyIndexMessage)
-			return SkipPushOnEmptyIndexMessage, nil
+		indexDescriptor, err := buildIndex(ctx, dataDir, sociStore, image)
+		if err != nil {
+			if err.Error() == ErrEmptyIndex.Error() {
+				log.Warn(ctx, SkipPushOnEmptyIndexMessage)
+				return SkipPushOnEmptyIndexMessage, nil
+			}
+			return logAndReturnError(ctx, BuildFailedMessage, err)
 		}
-		return logAndReturnError(ctx, BuildFailedMessage, err)
-	}
-	ctx = context.WithValue(ctx, "SOCIIndexDigest", indexDescriptor.Digest.String())
+		ctx = context.WithValue(ctx, "SOCIIndexDigest", indexDescriptor.Digest.String())
 
-	err = registry.Push(ctx, sociStore, *indexDescriptor, repo)
-	if err != nil {
-		return logAndReturnError(ctx, PushFailedMessage, err)
-	}
+		err = registry.Push(ctx, sociStore, *indexDescriptor, repo)
+		if err != nil {
+			return logAndReturnError(ctx, PushFailedMessage, err)
+		}
 
-	log.Info(ctx, BuildAndPushSuccessMessage)
+		log.Info(ctx, BuildAndPushSuccessMessage)
+	}
 	return BuildAndPushSuccessMessage, nil
 }
 

utils/registry/registry.go

@@ -47,6 +47,12 @@ type Registry struct {
 
 var RegistryNotSupportingOciArtifacts = errors.New("Registry does not support OCI artifacts")
 
+type Manifest struct {
+	ocispec.Manifest
+
+	Manifests []ocispec.Descriptor `json:"manifests"`
+}
+
 // Initialize a remote registry
 func Init(ctx context.Context, registryUrl string, authToken string) (*Registry, error) {
 	log.Info(ctx, "Initializing registry client")
@@ -128,9 +134,9 @@ func (registry *Registry) HeadManifest(ctx context.Context, repositoryName strin
 
 // Call registry's getManifest and return the image's manifest
 // The image reference must be a digest because that's what oras-go FetchReference takes
-func (registry *Registry) GetManifest(ctx context.Context, repositoryName string, digest string) (ocispec.Manifest, error) {
+func (registry *Registry) GetManifest(ctx context.Context, repositoryName string, digest string) (Manifest, error) {
 	repo, err := registry.registry.Repository(ctx, repositoryName)
-	var manifest ocispec.Manifest
+	var manifest Manifest
 	if err != nil {
 		return manifest, err
 	}
@@ -173,6 +179,49 @@ func (registry *Registry) ValidateImageManifest(ctx context.Context, repositoryN
 	return fmt.Errorf("Unexpected config media type: %s, expected one of: %v.", manifest.Config.MediaType, ImageConfigMediaTypes)
 }
 
+// GetImageDigests inspects an image reference and returns all valid digets that need to be indexed.
+// For multi-arch images (docker manifest), that includes all digests mentioned by the manifest.
+// For normal images, it's just the image digest itself.
+func (registry *Registry) GetImageDigests(ctx context.Context, repositoryName string, digest string) (digests []string, err error) {
+	manifest, err := registry.GetManifest(ctx, repositoryName, digest)
+	if err != nil {
+		return
+	}
+
+	if manifest.MediaType == MediaTypeDockerManifestList {
+		// multi-arch iamge
+		for _, internalManifest := range manifest.Manifests {
+			if internalManifest.MediaType == MediaTypeDockerManifest {
+				internalDigest := fmt.Sprintf("%s:%s", internalManifest.Digest.Algorithm().String(), internalManifest.Digest.Encoded())
+				if registry.ValidateImageManifest(ctx, repositoryName, internalDigest) == nil {
+					digests = append(digests, internalDigest)
+				}
+			}
+		}
+
+		if len(digests) == 0 {
+			err = fmt.Errorf("Manifest contains no valid images.")
+		}
+
+		return
+	}
+
+	if manifest.Config.MediaType == "" {
+		err = fmt.Errorf("Empty config media type.")
+		return
+	}
+
+	for _, configMediaType := range ImageConfigMediaTypes {
+		if manifest.Config.MediaType == configMediaType {
+			digests = append(digests, digest)
+			return
+		}
+	}
+
+	err = fmt.Errorf("Unexpected config media type: %s, expected one of: %v.", manifest.Config.MediaType, ImageConfigMediaTypes)
+	return
+}
+
 // Check if a registry is an ECR registry
 func isEcrRegistry(registryUrl string) bool {
 	ecrRegistryUrlRegex := "\\d{12}\\.dkr\\.ecr\\.\\S+\\.amazonaws\\.com"