Allow using WKT module as a dependency (#2906)

This fixes an issue where we were ignoring the use of
`buf.build/protocolbuffers/wellknowntypes`
as a configured dependency by always skipping WKT import paths and then
resolving the WKT
import at compile time.
This instead only skips the WKT path if it is not found in our module
set, and otherwise adds
the module containing the WKT path as a dependency.

Co-authored-by: bufdev <[email protected]>

private/bufpkg/bufmodule/bufmodule_test.go

@@ -147,10 +147,6 @@ func TestBasic(t *testing.T) {
 				"module2.proto": []byte(
 					`syntax = proto3; package module2; import "module1.proto"; import "extdep1.proto";`,
 				),
-				// Add a WKT, which is imported by a dependency. Make sure we have no cycle error.
-				"google/protobuf/timestamp.proto": []byte(
-					`syntax = proto3; package google.protobuf;`,
-				),
 				// module2 is excluded by path, but imports a Module that is not imported anywhere
 				// else. We want to make sure this path is not targeted, but extdep3 is still
 				// a dependency of the Module.
@@ -212,14 +208,12 @@ func TestBasic(t *testing.T) {
 		t,
 		module2,
 		"foo/module2_excluded.proto",
-		"google/protobuf/timestamp.proto",
 		"module2.proto",
 	)
 	// These are the target files. We excluded foo, so we only have module2.proto.
 	testTargetFilePaths(
 		t,
 		module2,
-		"google/protobuf/timestamp.proto",
 		"module2.proto",
 	)
 	//module2ProtoFileInfo, err := module2.StatFileInfo(ctx, "module2.proto")

private/bufpkg/bufmodule/module_dep.go

@@ -167,15 +167,22 @@ func getModuleDepsRec(
 				return fmt.Errorf("%s: %w", fileInfo.Path(), err)
 			}
 			for _, imp := range fastscanResult.Imports {
-				// Do not include as a dependency. If someone checks in the WKT, and one of their dependencies also
-				// imports a WKT (as is common), this algorithm will treat the original module as a dependency of
-				// the dependency, resulting in a module cycle.
-				if datawkt.Exists(imp.Path) {
-					continue
-				}
 				potentialModuleDep, err := moduleSet.getModuleForFilePath(ctx, imp.Path)
 				if err != nil {
 					if errors.Is(err, fs.ErrNotExist) {
+						// It is OK to not have a module with the WKT specified as a dependency. In this case,
+						// we do not include WKTs in our digest calculations. We've determined this is OK since
+						// WKTs are not downloaded and not subject to supply-side attacks.
+						//
+						// It is also OK for a module to have the WKT, such as buf.build/protocolbuffers/wellknowntypes.
+						// In this case, we want this to be recognized.
+						//
+						// Note that if someone checks in a WKT as part of their module, this will result in this module
+						// being a dependency for all other modules that import that WKT.
+						// This could result in unintended module cycles.
+						if datawkt.Exists(imp.Path) {
+							continue
+						}
 						// We specifically handle ImportNotExistErrors with exit code 100 in buf.go.
 						//
 						// We don't want to return a FileAnnotationSet here as we never have line

private/bufpkg/bufmodule/module_set.go

@@ -20,7 +20,6 @@ import (
 	"io/fs"
 	"sort"
 
-	"github.com/bufbuild/buf/private/gen/data/datawkt"
 	"github.com/bufbuild/buf/private/pkg/cache"
 	"github.com/bufbuild/buf/private/pkg/dag"
 	"github.com/bufbuild/buf/private/pkg/slicesext"
@@ -30,10 +29,6 @@ import (
 	"github.com/gofrs/uuid/v5"
 )
 
-// errIsWKT is the error returned by getFastscanResultForPath or getModuleForFilePath if the
-// input filePath is a well-known type.
-var errIsWKT = errors.New("wkt")
-
 // ModuleSet is a set of Modules constructed by a ModuleBuilder.
 //
 // A ModuleSet is expected to be self-contained, that is Modules only import
@@ -80,7 +75,6 @@ type ModuleSet interface {
 	// of getModuleDeps will result in nasty bugs down the line - we effectively ignore WKTs in this function,
 	// but WKTs may be included in one of the Modules as files.
 	//
-	// returns errIsWKT if the filePath is a WKT.
 	// returns an error with fs.ErrNotExist if the file is not found.
 	getModuleForFilePath(ctx context.Context, filePath string) (Module, error)
 
@@ -372,13 +366,11 @@ func (m *moduleSet) getModuleForFilePathUncached(ctx context.Context, filePath s
 	}
 	switch len(matchingOpaqueIDs) {
 	case 0:
-		// We do not include WKTs in our digest calculations. We've determined this is OK since
-		// WKTs are not downloaded and not subject to supply-side attacks.
-		if datawkt.Exists(filePath) {
-			return nil, errIsWKT
-		}
 		// This will happen if there is a file path we cannot find in our modules, which will result
 		// in an error on ModuleDeps() or Digest().
+		//
+		// Note that we will commonly call getModuleForFilePathUncached for a WKT, which will result
+		// in this error (this is desired).
 		return nil, &fs.PathError{Op: "stat", Path: filePath, Err: fs.ErrNotExist}
 	case 1:
 		var matchingOpaqueID string