Add buf policy update command (#3959)

Author: emcfarlane

private/buf/bufcli/policy_key_provider.go

@@ -0,0 +1,36 @@
+// Copyright 2020-2025 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package bufcli
+
+import (
+	"buf.build/go/app/appext"
+	"github.com/bufbuild/buf/private/bufpkg/bufpolicy"
+	"github.com/bufbuild/buf/private/bufpkg/bufpolicy/bufpolicyapi"
+	"github.com/bufbuild/buf/private/bufpkg/bufregistryapi/bufregistryapipolicy"
+)
+
+// NewPolicyKeyProvider returns a new PolicyKeyProvider.
+func NewPolicyKeyProvider(container appext.Container) (bufpolicy.PolicyKeyProvider, error) {
+	clientConfig, err := NewConnectClientConfig(container)
+	if err != nil {
+		return nil, err
+	}
+	return bufpolicyapi.NewPolicyKeyProvider(
+		container.Logger(),
+		bufregistryapipolicy.NewClientProvider(
+			clientConfig,
+		),
+	), nil
+}

private/buf/cmd/buf/buf.go

@@ -69,6 +69,7 @@ import (
 	"github.com/bufbuild/buf/private/buf/cmd/buf/command/plugin/pluginupdate"
 	"github.com/bufbuild/buf/private/buf/cmd/buf/command/policy/policyprune"
 	"github.com/bufbuild/buf/private/buf/cmd/buf/command/policy/policypush"
+	"github.com/bufbuild/buf/private/buf/cmd/buf/command/policy/policyupdate"
 	"github.com/bufbuild/buf/private/buf/cmd/buf/command/push"
 	"github.com/bufbuild/buf/private/buf/cmd/buf/command/registry/module/modulecommit/modulecommitaddlabel"
 	"github.com/bufbuild/buf/private/buf/cmd/buf/command/registry/module/modulecommit/modulecommitinfo"
@@ -205,6 +206,7 @@ func NewRootCommand(name string) *appcmd.Command {
 				Hidden: true,
 				SubCommands: []*appcmd.Command{
 					policypush.NewCommand("push", builder),
+					policyupdate.NewCommand("update", builder),
 					policyprune.NewCommand("prune", builder),
 				},
 			},

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

@@ -0,0 +1,222 @@
+// Copyright 2020-2025 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package policyupdate
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"buf.build/go/app/appcmd"
+	"buf.build/go/app/appext"
+	"buf.build/go/standard/xslices"
+	"github.com/bufbuild/buf/private/buf/bufcli"
+	"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/pkg/syserror"
+	"github.com/spf13/pflag"
+)
+
+const (
+	onlyFlagName = "only"
+)
+
+// NewCommand returns a new Command.
+func NewCommand(
+	name string,
+	builder appext.SubCommandBuilder,
+) *appcmd.Command {
+	flags := newFlags()
+	return &appcmd.Command{
+		Use:   name + " <directory>",
+		Short: "Update pinned remote policies in a buf.lock",
+		Long: `Fetch the latest digests for the specified policy references in buf.yaml.
+
+The first argument is the directory of the local module to update.
+Defaults to "." if no argument is specified.`,
+		Args: appcmd.MaximumNArgs(1),
+		Run: builder.NewRunFunc(
+			func(ctx context.Context, container appext.Container) error {
+				return run(ctx, container, flags)
+			},
+		),
+		BindFlags: flags.Bind,
+	}
+}
+
+type flags struct {
+	Only []string
+}
+
+func newFlags() *flags {
+	return &flags{}
+}
+
+func (f *flags) Bind(flagSet *pflag.FlagSet) {
+	flagSet.StringSliceVar(
+		&f.Only,
+		onlyFlagName,
+		nil,
+		"The name of the policy to update. When set, only this policy is updated. May be provided multiple times",
+	)
+	// TODO FUTURE: implement
+	_ = flagSet.MarkHidden(onlyFlagName)
+}
+
+func run(
+	ctx context.Context,
+	container appext.Container,
+	flags *flags,
+) (retErr error) {
+	dirPath := "."
+	if container.NumArgs() > 0 {
+		dirPath = container.Arg(0)
+	}
+	if len(flags.Only) > 0 {
+		// TODO FUTURE: implement
+		return syserror.Newf("--%s is not implemented", onlyFlagName)
+	}
+
+	logger := container.Logger()
+	controller, err := bufcli.NewController(container)
+	if err != nil {
+		return err
+	}
+	workspaceDepManager, err := controller.GetWorkspaceDepManager(ctx, dirPath)
+	if err != nil {
+		return err
+	}
+	configuredRemotePolicyRefs, err := workspaceDepManager.ConfiguredRemotePolicyRefs(ctx)
+	if err != nil {
+		return err
+	}
+	policyKeyProvider, err := bufcli.NewPolicyKeyProvider(container)
+	if err != nil {
+		return err
+	}
+	configuredRemotePolicyKeys, err := policyKeyProvider.GetPolicyKeysForPolicyRefs(
+		ctx,
+		configuredRemotePolicyRefs,
+		bufpolicy.DigestTypeO1,
+	)
+	if err != nil {
+		return err
+	}
+	configuredPolicyNameToRemotePluginKeys, err := getPolicyKeyPluginKeysForPolicyKeys(
+		ctx,
+		container,
+		configuredRemotePolicyKeys,
+	)
+	if err != nil {
+		return err
+	}
+
+	// Store the existing buf.lock data.
+	existingRemotePolicyKeys, err := workspaceDepManager.ExistingBufLockFileRemotePolicyKeys(ctx)
+	if err != nil {
+		return err
+	}
+	if configuredRemotePolicyKeys == nil && existingRemotePolicyKeys == nil {
+		// 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
+		// place and we do not need to go through the entire operation of updating non-existent
+		// deps and building the image for tamper-proofing.
+		logger.Warn(fmt.Sprintf("No configured remote policies were found to update in %q.", dirPath))
+		return nil
+	}
+	existingDepModuleKeys, err := workspaceDepManager.ExistingBufLockFileDepModuleKeys(ctx)
+	if err != nil {
+		return err
+	}
+	existingRemotePluginKeys, err := workspaceDepManager.ExistingBufLockFileRemotePluginKeys(ctx)
+	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.
+	//
+	// TODO FUTURE: We should be able to update the buf.lock file in an in-memory bucket, then do the rebuild,
+	// and if the rebuild is successful, then actually write to disk. It shouldn't even be that much work - just
+	// overlay the new buf.lock file in a union bucket.
+	defer func() {
+		if retErr != nil {
+			retErr = errors.Join(retErr, workspaceDepManager.UpdateBufLockFile(
+				ctx, existingDepModuleKeys, existingRemotePluginKeys, existingRemotePolicyKeys, existingPolicyNameToRemotePluginKeys,
+			))
+		}
+	}()
+	// Edit the buf.lock file with the updated remote plugins.
+	if err := workspaceDepManager.UpdateBufLockFile(ctx, existingDepModuleKeys, existingRemotePluginKeys, configuredRemotePolicyKeys, configuredPolicyNameToRemotePluginKeys); err != nil {
+		return err
+	}
+	return nil
+}
+
+func getPolicyKeyPluginKeysForPolicyKeys(
+	ctx context.Context,
+	container appext.Container,
+	policyKeys []bufpolicy.PolicyKey,
+) (map[string][]bufplugin.PluginKey, error) {
+	policyDataProvider, err := bufcli.NewPolicyDataProvider(container)
+	if err != nil {
+		return nil, err
+	}
+	pluginKeyProvider, err := bufcli.NewPluginKeyProvider(container)
+	if err != nil {
+		return nil, err
+	}
+	policyDatas, err := policyDataProvider.GetPolicyDatasForPolicyKeys(ctx, policyKeys)
+	if err != nil {
+		return nil, err
+	}
+	policyNameToRemotePluginKeys := make(map[string][]bufplugin.PluginKey)
+	for _, policyData := range policyDatas {
+		policyConfig, err := policyData.Config()
+		if err != nil {
+			return nil, err
+		}
+		pluginConfigs := policyConfig.PluginConfigs()
+		pluginRefs, err := xslices.MapError(pluginConfigs, func(pluginConfig bufpolicy.PluginConfig) (bufparse.Ref, error) {
+			pluginRef := pluginConfig.Ref()
+			if pluginRef == nil {
+				return nil, fmt.Errorf("plugin config %q does not have a valid ref", pluginConfig.Name())
+			}
+			return pluginRef, nil
+		})
+		if err != nil {
+			return nil, fmt.Errorf("failed to get plugin refs for policy %q: %w", policyData.PolicyKey(), err)
+		}
+		remotePluginKeys, err := pluginKeyProvider.GetPluginKeysForPluginRefs(
+			ctx,
+			pluginRefs,
+			bufplugin.DigestTypeP1,
+		)
+		if err != nil {
+			return nil, err
+		}
+		policyName := policyData.PolicyKey().FullName().String()
+		if len(remotePluginKeys) > 0 {
+			policyNameToRemotePluginKeys[policyName] = remotePluginKeys
+		}
+	}
+	return policyNameToRemotePluginKeys, nil
+}