From e28d21f8c9d82982dca0e2399ba38e9f840aa995 Mon Sep 17 00:00:00 2001
From: "Philip K. Warren" <pkwarren@users.noreply.github.com>
Date: Mon, 12 Dec 2022 10:59:57 -0600
Subject: [PATCH] optimize checking for current digest (#1665)

If the current digest has the same image ID as the image to push, we can
greatly speed up looking through every tag (which may contain multiple
duplicates - potentially one per version/revision of a plugin). This
should greatly speed up the lookup in the case where the plugin already
exists on the OCI registry.
---
 .../alpha/plugin/pluginpush/pluginpush.go     | 55 ++++++++++++++-----
 1 file changed, 41 insertions(+), 14 deletions(-)

diff --git a/private/buf/cmd/buf/command/alpha/plugin/pluginpush/pluginpush.go b/private/buf/cmd/buf/command/alpha/plugin/pluginpush/pluginpush.go
index 59539ab158..a4d2cfd654 100644
--- a/private/buf/cmd/buf/command/alpha/plugin/pluginpush/pluginpush.go
+++ b/private/buf/cmd/buf/command/alpha/plugin/pluginpush/pluginpush.go
@@ -236,6 +236,7 @@ func run(
 			Revision: 0, // get latest revision for the plugin version.
 		}),
 	)
+	var currentImageDigest string
 	var nextRevision uint32
 	if err != nil {
 		if connect.CodeOf(err) != connect.CodeNotFound {
@@ -244,6 +245,7 @@ func run(
 		nextRevision = 1
 	} else {
 		nextRevision = latestPluginResp.Msg.Plugin.Revision + 1
+		currentImageDigest = latestPluginResp.Msg.Plugin.ContainerImageDigest
 	}
 	machine, err := netrc.GetMachineForName(container, pluginConfig.Name.Remote())
 	if err != nil {
@@ -255,7 +257,7 @@ func run(
 		authConfig.Username = machine.Login()
 		authConfig.Password = machine.Password()
 	}
-	imageDigest, err := findExistingDigestForImageID(ctx, pluginConfig, authConfig, imageID)
+	imageDigest, err := findExistingDigestForImageID(ctx, pluginConfig, authConfig, imageID, currentImageDigest)
 	if err != nil {
 		return err
 	}
@@ -366,7 +368,13 @@ func pushImage(
 // - For each tag:
 //   - Fetch image: GET /v2/{owner}/{plugin}/manifests/{tag}
 //   - If image manifest matches imageID, we can use the image digest for the image.
-func findExistingDigestForImageID(ctx context.Context, plugin *bufpluginconfig.Config, authConfig *bufplugindocker.RegistryAuthConfig, imageID string) (string, error) {
+func findExistingDigestForImageID(
+	ctx context.Context,
+	plugin *bufpluginconfig.Config,
+	authConfig *bufplugindocker.RegistryAuthConfig,
+	imageID string,
+	currentImageDigest string,
+) (string, error) {
 	pluginsRemote := plugin.Name.Remote()
 	if !strings.HasPrefix(pluginsRemote, bufplugindocker.PluginsImagePrefix) {
 		pluginsRemote = bufplugindocker.PluginsImagePrefix + pluginsRemote
@@ -376,7 +384,19 @@ func findExistingDigestForImageID(ctx context.Context, plugin *bufpluginconfig.C
 		return "", err
 	}
 	auth := &authn.Basic{Username: authConfig.Username, Password: authConfig.Password}
-	tags, err := remote.List(repo, remote.WithContext(ctx), remote.WithAuth(auth))
+	remoteOpts := []remote.Option{remote.WithContext(ctx), remote.WithAuth(auth)}
+	// First attempt to see if the current image digest matches the image ID
+	if currentImageDigest != "" {
+		remoteImageID, _, err := getImageIDAndDigestFromReference(repo.Digest(currentImageDigest), remoteOpts...)
+		if err != nil {
+			return "", err
+		}
+		if remoteImageID == imageID {
+			return currentImageDigest, nil
+		}
+	}
+	// List all tags and check for a match
+	tags, err := remote.List(repo, remoteOpts...)
 	if err != nil {
 		structuredErr := new(transport.Error)
 		if errors.As(err, &structuredErr) {
@@ -388,27 +408,34 @@ func findExistingDigestForImageID(ctx context.Context, plugin *bufpluginconfig.C
 	}
 	existingImageDigest := ""
 	for _, tag := range tags {
-		image, err := remote.Image(repo.Tag(tag), remote.WithContext(ctx), remote.WithAuth(auth))
+		remoteImageID, imageDigest, err := getImageIDAndDigestFromReference(repo.Tag(tag), remoteOpts...)
 		if err != nil {
 			return "", err
 		}
-		manifest, err := image.Manifest()
-		if err != nil {
-			return "", err
-		}
-		remoteImageID := manifest.Config.Digest.String()
 		if remoteImageID == imageID {
-			imageDigest, err := image.Digest()
-			if err != nil {
-				return "", err
-			}
-			existingImageDigest = imageDigest.String()
+			existingImageDigest = imageDigest
 			break
 		}
 	}
 	return existingImageDigest, nil
 }
 
+func getImageIDAndDigestFromReference(ref name.Reference, options ...remote.Option) (string, string, error) {
+	image, err := remote.Image(ref, options...)
+	if err != nil {
+		return "", "", err
+	}
+	imageDigest, err := image.Digest()
+	if err != nil {
+		return "", "", err
+	}
+	manifest, err := image.Manifest()
+	if err != nil {
+		return "", "", err
+	}
+	return manifest.Config.Digest.String(), imageDigest.String(), nil
+}
+
 func unzipPluginToSourceBucket(ctx context.Context, pluginZip string, size int64, bucket storage.ReadWriteBucket) (retErr error) {
 	f, err := os.Open(pluginZip)
 	if err != nil {