Resolve local policies remote plugin refs (#3981)

Author: emcfarlane

private/buf/bufworkspace/workspace_dep_manager.go

@@ -26,6 +26,7 @@ import (
 	"github.com/bufbuild/buf/private/bufpkg/bufparse"
 	"github.com/bufbuild/buf/private/bufpkg/bufplugin"
 	"github.com/bufbuild/buf/private/bufpkg/bufpolicy"
+	"github.com/bufbuild/buf/private/bufpkg/bufpolicy/bufpolicyconfig"
 	"github.com/bufbuild/buf/private/pkg/storage"
 	"github.com/bufbuild/buf/private/pkg/syserror"
 )
@@ -92,6 +93,16 @@ type WorkspaceDepManager interface {
 	//
 	// Sorted.
 	ConfiguredRemotePolicyRefs(ctx context.Context) ([]bufparse.Ref, error)
+	// ConfiguredLocalPolicyNameToRemotePluginRefs returns the configured remote plugins for each local policy of the Workspace.
+	//
+	// These come from buf.yaml files and the local buf.policy.yaml files.
+	//
+	// The PluginRefs for each Policy will be unique by FullName. If there are two PluginRefs
+	// in the buf.yaml for a given Policy with the same FullName but different Refs, an error will be given
+	// at workspace constructions.
+	//
+	// PluginRefs are sorted for each Policy.
+	ConfiguredLocalPolicyNameToRemotePluginRefs(ctx context.Context) (map[string][]bufparse.Ref, error)
 
 	isWorkspaceDepManager()
 }
@@ -242,6 +253,62 @@ func (w *workspaceDepManager) ConfiguredRemotePolicyRefs(ctx context.Context) ([
 	return policyRefs, nil
 }
 
+func (w *workspaceDepManager) ConfiguredLocalPolicyNameToRemotePluginRefs(ctx context.Context) (map[string][]bufparse.Ref, error) {
+	bufYAMLFile, err := bufconfig.GetBufYAMLFileForPrefix(ctx, w.bucket, w.targetSubDirPath)
+	if err != nil {
+		if !errors.Is(err, fs.ErrNotExist) {
+			return nil, err
+		}
+	}
+	if bufYAMLFile == nil {
+		return nil, nil
+	}
+	switch fileVersion := bufYAMLFile.FileVersion(); fileVersion {
+	case bufconfig.FileVersionV1Beta1, bufconfig.FileVersionV1:
+		if w.isV2 {
+			return nil, syserror.Newf("buf.yaml at %q did had version %v but expected v1beta1, v1", w.targetSubDirPath, fileVersion)
+		}
+		// Policys are not supported in versions less than v2.
+		return nil, nil
+	case bufconfig.FileVersionV2:
+		if !w.isV2 {
+			return nil, syserror.Newf("buf.yaml at %q did had version %v but expected v2", w.targetSubDirPath, fileVersion)
+		}
+	default:
+		return nil, syserror.Newf("unknown FileVersion: %v", fileVersion)
+	}
+	localPolicyNameToRemotePluginRefs := make(map[string][]bufparse.Ref)
+	for _, policyConfig := range bufYAMLFile.PolicyConfigs() {
+		if policyConfig.Ref() != nil {
+			continue // Only local policies refs are considered here.
+		}
+		localPolicyName := policyConfig.Name()
+		bufPolicyFile, err := bufpolicyconfig.GetBufPolicyYAMLFile(ctx, w.bucket, localPolicyName)
+		if err != nil {
+			return nil, err
+		}
+		pluginRefs := xslices.Filter(
+			xslices.Map(
+				bufPolicyFile.PluginConfigs(),
+				func(value bufpolicy.PluginConfig) bufparse.Ref {
+					return value.Ref()
+				},
+			),
+			func(value bufparse.Ref) bool {
+				return value != nil
+			},
+		)
+		sort.Slice(
+			pluginRefs,
+			func(i int, j int) bool {
+				return pluginRefs[i].FullName().String() < pluginRefs[j].FullName().String()
+			},
+		)
+		localPolicyNameToRemotePluginRefs[localPolicyName] = pluginRefs
+	}
+	return localPolicyNameToRemotePluginRefs, nil
+}
+
 func (w *workspaceDepManager) BufLockFileDigestType() bufmodule.DigestType {
 	if w.isV2 {
 		return bufmodule.DigestTypeB5

private/buf/cmd/buf/command/policy/policyupdate/policyupdate.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"maps"
 
 	"buf.build/go/app/appcmd"
 	"buf.build/go/app/appext"
@@ -103,6 +104,10 @@ func run(
 	if err != nil {
 		return err
 	}
+	configuredLocalPolicyNameToRemotePluginRefs, err := workspaceDepManager.ConfiguredLocalPolicyNameToRemotePluginRefs(ctx)
+	if err != nil {
+		return err
+	}
 	policyKeyProvider, err := bufcli.NewPolicyKeyProvider(container)
 	if err != nil {
 		return err
@@ -115,21 +120,40 @@ func run(
 	if err != nil {
 		return err
 	}
-	configuredPolicyNameToRemotePluginKeys, err := getPolicyKeyPluginKeysForPolicyKeys(
+	configuredRemotePolicyNameToRemotePluginKeys, err := getPolicyKeyPluginKeysForPolicyKeys(
 		ctx,
 		container,
 		configuredRemotePolicyKeys,
 	)
 	if err != nil {
 		return err
 	}
+	configuredLocalPolicyNameToRemotePluginKeys, err := getPolicyPluginKeysForPolicyNames(
+		ctx,
+		container,
+		configuredLocalPolicyNameToRemotePluginRefs,
+	)
+	if err != nil {
+		return err
+	}
+	var configuredPolicyNameToRemotePluginKeys map[string][]bufplugin.PluginKey
+	if policyCount := len(configuredRemotePolicyNameToRemotePluginKeys) + len(configuredLocalPolicyNameToRemotePluginKeys); policyCount > 0 {
+		configuredPolicyNameToRemotePluginKeys = make(map[string][]bufplugin.PluginKey, policyCount)
+		maps.Copy(configuredPolicyNameToRemotePluginKeys, configuredRemotePolicyNameToRemotePluginKeys)
+		maps.Copy(configuredPolicyNameToRemotePluginKeys, configuredLocalPolicyNameToRemotePluginKeys)
+	}
 
 	// Store the existing buf.lock data.
 	existingRemotePolicyKeys, err := workspaceDepManager.ExistingBufLockFileRemotePolicyKeys(ctx)
 	if err != nil {
 		return err
 	}
-	if configuredRemotePolicyKeys == nil && existingRemotePolicyKeys == nil {
+	existingPolicyNameToRemotePluginKeys, err := workspaceDepManager.ExistingBufLockFilePolicyNameToRemotePluginKeys(ctx)
+	if err != nil {
+		return err
+	}
+	if len(configuredRemotePolicyKeys) == 0 && len(configuredPolicyNameToRemotePluginKeys) == 0 &&
+		len(existingRemotePolicyKeys) == 0 && len(existingPolicyNameToRemotePluginKeys) == 0 {
 		// No new configured remote plugins were found, and no existing buf.lock deps were found, so there
 		// is nothing to update, we can return here.
 		// This ensures we do not create an empty buf.lock when one did not exist in the first
@@ -146,11 +170,6 @@ func run(
 	if err != nil {
 		return err
 	}
-	existingPolicyNameToRemotePluginKeys, err := workspaceDepManager.ExistingBufLockFilePolicyNameToRemotePluginKeys(ctx)
-	if err != nil {
-		return err
-	}
-
 	// We're about to edit the buf.lock file on disk. If we have a subsequent error,
 	// attempt to revert the buf.lock file.
 	//
@@ -176,6 +195,9 @@ func getPolicyKeyPluginKeysForPolicyKeys(
 	container appext.Container,
 	policyKeys []bufpolicy.PolicyKey,
 ) (map[string][]bufplugin.PluginKey, error) {
+	if len(policyKeys) == 0 {
+		return nil, nil
+	}
 	policyDataProvider, err := bufcli.NewPolicyDataProvider(container)
 	if err != nil {
 		return nil, err
@@ -220,3 +242,32 @@ func getPolicyKeyPluginKeysForPolicyKeys(
 	}
 	return policyNameToRemotePluginKeys, nil
 }
+
+func getPolicyPluginKeysForPolicyNames(
+	ctx context.Context,
+	container appext.Container,
+	localPolicyNameToRemotePluginRefs map[string][]bufparse.Ref,
+) (map[string][]bufplugin.PluginKey, error) {
+	if len(localPolicyNameToRemotePluginRefs) == 0 {
+		return nil, nil
+	}
+	pluginKeyProvider, err := bufcli.NewPluginKeyProvider(container)
+	if err != nil {
+		return nil, err
+	}
+	policyNameToRemotePluginKeys := make(map[string][]bufplugin.PluginKey)
+	for policyName, pluginRefs := range localPolicyNameToRemotePluginRefs {
+		remotePluginKeys, err := pluginKeyProvider.GetPluginKeysForPluginRefs(
+			ctx,
+			pluginRefs,
+			bufplugin.DigestTypeP1,
+		)
+		if err != nil {
+			return nil, err
+		}
+		if len(remotePluginKeys) > 0 {
+			policyNameToRemotePluginKeys[policyName] = remotePluginKeys
+		}
+	}
+	return policyNameToRemotePluginKeys, nil
+}