From f7b3a0b9082af6ad13526ffbfd51273e7e0a52a6 Mon Sep 17 00:00:00 2001
From: Edward McFarlane <emcfarlane@buf.build>
Date: Fri, 20 Sep 2024 10:44:12 -0400
Subject: [PATCH] Add WKT resolver for bufcurl

---
 private/buf/bufcurl/resolver.go          | 33 ++++++++++++++++++++++++
 private/buf/cmd/buf/command/curl/curl.go | 11 ++++++--
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/private/buf/bufcurl/resolver.go b/private/buf/bufcurl/resolver.go
index b83e10025e..ea38b900d2 100644
--- a/private/buf/bufcurl/resolver.go
+++ b/private/buf/bufcurl/resolver.go
@@ -15,8 +15,13 @@
 package bufcurl
 
 import (
+	"context"
+
 	"github.com/bufbuild/buf/private/bufpkg/bufimage"
+	"github.com/bufbuild/buf/private/bufpkg/bufmodule"
+	"github.com/bufbuild/buf/private/gen/data/datawkt"
 	"github.com/bufbuild/buf/private/pkg/protoencoding"
+	"github.com/bufbuild/buf/private/pkg/tracing"
 	"google.golang.org/protobuf/reflect/protoreflect"
 )
 
@@ -96,3 +101,31 @@ func (c *combinedResolver) ListServices() ([]protoreflect.FullName, error) {
 	}
 	return serviceNames, nil
 }
+
+// NewWKTResolver returns a Resolver that can resolve all well-known types.
+func NewWKTResolver(ctx context.Context, tracer tracing.Tracer) (Resolver, error) {
+	moduleSet, err := bufmodule.NewModuleSetBuilder(
+		ctx,
+		tracer,
+		bufmodule.NopModuleDataProvider,
+		bufmodule.NopCommitProvider,
+	).AddLocalModule(
+		datawkt.ReadBucket,
+		".",
+		true,
+	).Build()
+	if err != nil {
+		return nil, err
+	}
+	module := bufmodule.ModuleSetToModuleReadBucketWithOnlyProtoFiles(moduleSet)
+	image, err := bufimage.BuildImage(
+		ctx,
+		tracer,
+		module,
+		bufimage.WithExcludeSourceCodeInfo(),
+	)
+	if err != nil {
+		return nil, err
+	}
+	return ResolverForImage(image), nil
+}
diff --git a/private/buf/cmd/buf/command/curl/curl.go b/private/buf/cmd/buf/command/curl/curl.go
index 6a275b49d8..bdb1a31770 100644
--- a/private/buf/cmd/buf/command/curl/curl.go
+++ b/private/buf/cmd/buf/command/curl/curl.go
@@ -40,6 +40,7 @@ import (
 	"github.com/bufbuild/buf/private/pkg/app/appext"
 	"github.com/bufbuild/buf/private/pkg/netrc"
 	"github.com/bufbuild/buf/private/pkg/stringutil"
+	"github.com/bufbuild/buf/private/pkg/tracing"
 	"github.com/bufbuild/buf/private/pkg/verbose"
 	"github.com/quic-go/quic-go"
 	"github.com/quic-go/quic-go/http3"
@@ -965,7 +966,7 @@ func run(ctx context.Context, container appext.Container, f *flags) (err error)
 		}
 	}
 
-	resolvers := make([]bufcurl.Resolver, 0, len(f.Schemas)+1)
+	resolvers := make([]bufcurl.Resolver, 0, len(f.Schemas)+2)
 	if f.Reflect {
 		reflectHeaders, _, err := bufcurl.LoadHeaders(f.ReflectHeaders, "", requestHeaders)
 		if err != nil {
@@ -1010,7 +1011,13 @@ func run(ctx context.Context, container appext.Container, f *flags) (err error)
 		}
 		resolvers = append(resolvers, bufcurl.ResolverForImage(image))
 	}
-	res := bufcurl.CombineResolvers(resolvers...)
+	// Add a WKT resolver to the end of the end of the list. This is used
+	// for printing a WKT encoded in a "google.protobuf.Any" type as JSON.
+	wktResolver, err := bufcurl.NewWKTResolver(ctx, tracing.NewTracer(container.Tracer()))
+	if err != nil {
+		return err
+	}
+	res := bufcurl.CombineResolvers(append(resolvers, wktResolver)...)
 
 	switch {
 	case f.ListServices || f.ListMethods: