github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/grpc/xds/googledirectpath/googlec2p_test.go

/*
 *
 * Copyright 2021 gRPC authors.
 *
 * 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 googledirectpath

import (
	"strconv"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	grpc "github.com/hxx258456/ccgo/grpc"
	"github.com/hxx258456/ccgo/grpc/internal/envconfig"
	"github.com/hxx258456/ccgo/grpc/resolver"
	"github.com/hxx258456/ccgo/grpc/xds/internal/xdsclient"
	"github.com/hxx258456/ccgo/grpc/xds/internal/xdsclient/bootstrap"
	"github.com/hxx258456/ccgo/grpc/xds/internal/xdsclient/xdsresource/version"
	"google.golang.org/protobuf/testing/protocmp"
	"google.golang.org/protobuf/types/known/structpb"

	v3corepb "github.com/hxx258456/ccgo/go-control-plane/envoy/config/core/v3"
)

type emptyResolver struct {
	resolver.Resolver
	scheme string
}

func (er *emptyResolver) Build(_ resolver.Target, _ resolver.ClientConn, _ resolver.BuildOptions) (resolver.Resolver, error) {
	return er, nil
}

func (er *emptyResolver) Scheme() string {
	return er.scheme
}

func (er *emptyResolver) Close() {}

var (
	testDNSResolver = &emptyResolver{scheme: "dns"}
	testXDSResolver = &emptyResolver{scheme: "xds"}
)

func replaceResolvers() func() {
	var registerForTesting bool
	if resolver.Get(c2pScheme) == nil {
		// If env var to enable c2p is not set, the resolver isn't registered.
		// Need to register and unregister in defer.
		registerForTesting = true
		resolver.Register(&c2pResolverBuilder{})
	}
	oldDNS := resolver.Get("dns")
	resolver.Register(testDNSResolver)
	oldXDS := resolver.Get("xds")
	resolver.Register(testXDSResolver)
	return func() {
		if oldDNS != nil {
			resolver.Register(oldDNS)
		} else {
			resolver.UnregisterForTesting("dns")
		}
		if oldXDS != nil {
			resolver.Register(oldXDS)
		} else {
			resolver.UnregisterForTesting("xds")
		}
		if registerForTesting {
			resolver.UnregisterForTesting(c2pScheme)
		}
	}
}

// Test that when bootstrap env is set, fallback to DNS.
func TestBuildWithBootstrapEnvSet(t *testing.T) {
	defer replaceResolvers()()
	builder := resolver.Get(c2pScheme)

	for i, envP := range []*string{&envconfig.XDSBootstrapFileName, &envconfig.XDSBootstrapFileContent} {
		t.Run(strconv.Itoa(i), func(t *testing.T) {
			// Set bootstrap config env var.
			oldEnv := *envP
			*envP = "does not matter"
			defer func() { *envP = oldEnv }()

			// Build should return DNS, not xDS.
			r, err := builder.Build(resolver.Target{}, nil, resolver.BuildOptions{})
			if err != nil {
				t.Fatalf("failed to build resolver: %v", err)
			}
			if r != testDNSResolver {
				t.Fatalf("want dns resolver, got %#v", r)
			}
		})
	}
}

// Test that when not on GCE, fallback to DNS.
func TestBuildNotOnGCE(t *testing.T) {
	defer replaceResolvers()()
	builder := resolver.Get(c2pScheme)

	oldOnGCE := onGCE
	onGCE = func() bool { return false }
	defer func() { onGCE = oldOnGCE }()

	// Build should return DNS, not xDS.
	r, err := builder.Build(resolver.Target{}, nil, resolver.BuildOptions{})
	if err != nil {
		t.Fatalf("failed to build resolver: %v", err)
	}
	if r != testDNSResolver {
		t.Fatalf("want dns resolver, got %#v", r)
	}
}

type testXDSClient struct {
	xdsclient.XDSClient
	closed chan struct{}
}

func (c *testXDSClient) Close() {
	c.closed <- struct{}{}
}

// Test that when xDS is built, the client is built with the correct config.
func TestBuildXDS(t *testing.T) {
	defer replaceResolvers()()
	builder := resolver.Get(c2pScheme)

	oldOnGCE := onGCE
	onGCE = func() bool { return true }
	defer func() { onGCE = oldOnGCE }()

	const testZone = "test-zone"
	oldGetZone := getZone
	getZone = func(time.Duration) string { return testZone }
	defer func() { getZone = oldGetZone }()

	for _, tt := range []struct {
		name  string
		ipv6  bool
		tdURI string // traffic director URI will be overridden if this is set.
	}{
		{name: "ipv6 true", ipv6: true},
		{name: "ipv6 false", ipv6: false},
		{name: "override TD URI", ipv6: true, tdURI: "test-uri"},
	} {
		t.Run(tt.name, func(t *testing.T) {
			oldGetIPv6Capability := getIPv6Capable
			getIPv6Capable = func(time.Duration) bool { return tt.ipv6 }
			defer func() { getIPv6Capable = oldGetIPv6Capability }()

			if tt.tdURI != "" {
				oldURI := envconfig.C2PResolverTestOnlyTrafficDirectorURI
				envconfig.C2PResolverTestOnlyTrafficDirectorURI = tt.tdURI
				defer func() {
					envconfig.C2PResolverTestOnlyTrafficDirectorURI = oldURI
				}()
			}

			tXDSClient := &testXDSClient{closed: make(chan struct{}, 1)}

			configCh := make(chan *bootstrap.Config, 1)
			oldNewClient := newClientWithConfig
			newClientWithConfig = func(config *bootstrap.Config) (xdsclient.XDSClient, error) {
				configCh <- config
				return tXDSClient, nil
			}
			defer func() { newClientWithConfig = oldNewClient }()

			// Build should return DNS, not xDS.
			r, err := builder.Build(resolver.Target{}, nil, resolver.BuildOptions{})
			if err != nil {
				t.Fatalf("failed to build resolver: %v", err)
			}
			rr := r.(*c2pResolver)
			if rrr := rr.Resolver; rrr != testXDSResolver {
				t.Fatalf("want xds resolver, got %#v, ", rrr)
			}

			wantNode := &v3corepb.Node{
				Id:                   id,
				Metadata:             nil,
				Locality:             &v3corepb.Locality{Zone: testZone},
				UserAgentName:        gRPCUserAgentName,
				UserAgentVersionType: &v3corepb.Node_UserAgentVersion{UserAgentVersion: grpc.Version},
				ClientFeatures:       []string{clientFeatureNoOverprovisioning},
			}
			if tt.ipv6 {
				wantNode.Metadata = &structpb.Struct{
					Fields: map[string]*structpb.Value{
						ipv6CapableMetadataName: {
							Kind: &structpb.Value_BoolValue{BoolValue: true},
						},
					},
				}
			}
			wantConfig := &bootstrap.Config{
				XDSServer: &bootstrap.ServerConfig{
					ServerURI:    tdURL,
					TransportAPI: version.TransportV3,
					NodeProto:    wantNode,
				},
				ClientDefaultListenerResourceNameTemplate: "%s",
			}
			if tt.tdURI != "" {
				wantConfig.XDSServer.ServerURI = tt.tdURI
			}
			cmpOpts := cmp.Options{
				cmpopts.IgnoreFields(bootstrap.ServerConfig{}, "Creds"),
				cmp.AllowUnexported(bootstrap.ServerConfig{}),
				protocmp.Transform(),
			}
			select {
			case c := <-configCh:
				if diff := cmp.Diff(c, wantConfig, cmpOpts); diff != "" {
					t.Fatalf("%v", diff)
				}
			case <-time.After(time.Second):
				t.Fatalf("timeout waiting for client config")
			}

			r.Close()
			select {
			case <-tXDSClient.closed:
			case <-time.After(time.Second):
				t.Fatalf("timeout waiting for client close")
			}
		})
	}
}