From 4d51afdb048f53faa312678b9b370eb8749365a4 Mon Sep 17 00:00:00 2001
From: Edward McFarlane <emcfarlane@buf.build>
Date: Tue, 10 Dec 2024 19:47:15 +0000
Subject: [PATCH] Validate remote PluginConfigs are present in buf.lock

---
 private/buf/bufctl/controller.go              | 39 ++++++++++++++++---
 .../bufpkg/bufplugin/plugin_key_provider.go   | 16 ++++----
 2 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/private/buf/bufctl/controller.go b/private/buf/bufctl/controller.go
index d4543c502e..d0e101f094 100644
--- a/private/buf/bufctl/controller.go
+++ b/private/buf/bufctl/controller.go
@@ -1203,30 +1203,38 @@ Declare %q in the deps key in your buf.yaml.`,
 //
 // Remote plugins Refs are resolved to PluginKeys from the workspace buf.lock file.
 // If the Ref is a MessageRef, we use the current directory buf.lock file.
+// PluginConfigs are validated to ensure that all remote Refs are pinned in the buf.lock file.
 func (c *controller) getPluginKeyProviderForRef(
 	ctx context.Context,
 	ref buffetch.Ref,
 	functionOptions *functionOptions,
 ) (_ bufplugin.PluginKeyProvider, retErr error) {
+	var (
+		pluginConfigs []bufconfig.PluginConfig
+		pluginKeys    []bufplugin.PluginKey
+	)
 	switch t := ref.(type) {
 	case buffetch.ProtoFileRef:
 		workspace, err := c.getWorkspaceForProtoFileRef(ctx, t, functionOptions)
 		if err != nil {
 			return nil, err
 		}
-		return bufplugin.NewStaticPluginKeyProvider(workspace.RemotePluginKeys())
+		pluginConfigs = workspace.PluginConfigs()
+		pluginKeys = workspace.RemotePluginKeys()
 	case buffetch.SourceRef:
 		workspace, err := c.getWorkspaceForSourceRef(ctx, t, functionOptions)
 		if err != nil {
 			return nil, err
 		}
-		return bufplugin.NewStaticPluginKeyProvider(workspace.RemotePluginKeys())
+		pluginConfigs = workspace.PluginConfigs()
+		pluginKeys = workspace.RemotePluginKeys()
 	case buffetch.ModuleRef:
 		workspace, err := c.getWorkspaceForModuleRef(ctx, t, functionOptions)
 		if err != nil {
 			return nil, err
 		}
-		return bufplugin.NewStaticPluginKeyProvider(workspace.RemotePluginKeys())
+		pluginConfigs = workspace.PluginConfigs()
+		pluginKeys = workspace.RemotePluginKeys()
 	case buffetch.MessageRef:
 		bucket, err := c.storageosProvider.NewReadWriteBucket(
 			".",
@@ -1247,10 +1255,11 @@ func (c *controller) getPluginKeyProviderForRef(
 			}
 			// We did not find a buf.yaml in our current directory,
 			// and there was no config override.
+			// Remote plugins are not available.
 			return bufplugin.NopPluginKeyProvider, nil
 		}
-		var pluginKeys []bufplugin.PluginKey
 		if bufYAMLFile.FileVersion() == bufconfig.FileVersionV2 {
+			pluginConfigs = bufYAMLFile.PluginConfigs()
 			bufLockFile, err := bufconfig.GetBufLockFileForPrefix(
 				ctx,
 				bucket,
@@ -1267,11 +1276,31 @@ func (c *controller) getPluginKeyProviderForRef(
 			}
 			pluginKeys = bufLockFile.RemotePluginKeys()
 		}
-		return bufplugin.NewStaticPluginKeyProvider(pluginKeys)
 	default:
 		// This is a system error.
 		return nil, syserror.Newf("invalid Ref: %T", ref)
 	}
+	// Validate the all remote PluginConfigs are in the buf.lock file.
+	pluginKeysByFullName, err := slicesext.ToUniqueValuesMap(pluginKeys, func(pluginKey bufplugin.PluginKey) string {
+		return pluginKey.FullName().String()
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to validate remote PluginKeys: %w", err)
+	}
+	remotePluginRefs := slicesext.Filter(
+		slicesext.Map(pluginConfigs, func(pluginConfig bufconfig.PluginConfig) bufparse.Ref {
+			return pluginConfig.Ref()
+		}),
+		func(pluginRef bufparse.Ref) bool {
+			return pluginRef != nil
+		},
+	)
+	for _, remotePluginRef := range remotePluginRefs {
+		if _, ok := pluginKeysByFullName[remotePluginRef.FullName().String()]; !ok {
+			return nil, fmt.Errorf(`remote plugin %q is not in the buf.lock file, use "buf plugin update" to pin remote refs`, remotePluginRef)
+		}
+	}
+	return bufplugin.NewStaticPluginKeyProvider(pluginKeys)
 }
 
 // handleFileAnnotationSetError will attempt to handle the error as a FileAnnotationSet, and if so, print
diff --git a/private/bufpkg/bufplugin/plugin_key_provider.go b/private/bufpkg/bufplugin/plugin_key_provider.go
index cef430f5db..0fba69f8b6 100644
--- a/private/bufpkg/bufplugin/plugin_key_provider.go
+++ b/private/bufpkg/bufplugin/plugin_key_provider.go
@@ -59,13 +59,8 @@ func NewStaticPluginKeyProvider(pluginKeys []PluginKey) (PluginKeyProvider, erro
 	if err != nil {
 		return nil, err
 	}
-	digetType, err := UniqueDigestTypeForPluginKeys(pluginKeys)
-	if err != nil {
-		return nil, err
-	}
 	return staticPluginKeyProvider{
 		pluginKeysByFullName: pluginKeysByFullName,
-		digestType:           digetType,
 	}, nil
 }
 
@@ -83,7 +78,6 @@ func (nopPluginKeyProvider) GetPluginKeysForPluginRefs(
 
 type staticPluginKeyProvider struct {
 	pluginKeysByFullName map[string]PluginKey
-	digestType           DigestType
 }
 
 func (s staticPluginKeyProvider) GetPluginKeysForPluginRefs(
@@ -91,9 +85,6 @@ func (s staticPluginKeyProvider) GetPluginKeysForPluginRefs(
 	refs []bufparse.Ref,
 	digestType DigestType,
 ) ([]PluginKey, error) {
-	if digestType != s.digestType {
-		return nil, fmt.Errorf("expected DigestType %v, got %v", s.digestType, digestType)
-	}
 	pluginKeys := make([]PluginKey, len(refs))
 	for i, ref := range refs {
 		// Only the FullName is used to match the PluginKey. The Ref is not
@@ -103,6 +94,13 @@ func (s staticPluginKeyProvider) GetPluginKeysForPluginRefs(
 		if !ok {
 			return nil, fs.ErrNotExist
 		}
+		digest, err := pluginKey.Digest()
+		if err != nil {
+			return nil, err
+		}
+		if digest.Type() != digestType {
+			return nil, fmt.Errorf("expected DigestType %v, got %v", digestType, digest.Type())
+		}
 		pluginKeys[i] = pluginKey
 	}
 	return pluginKeys, nil