Repository: grpc/grpc-go
File: test/creds_test.go (Go)

Code before:
const testAuthority = "test.auth.ori.ty"

type authorityCheckCreds struct {
	credentials.TransportCredentials
	got string
}

func (c *authorityCheckCreds) ClientHandshake(_ context.Context, authority string, rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
	c.got = authority
	return rawConn, nil, nil
}
func (c *authorityCheckCreds) Info() credentials.ProtocolInfo {
	return credentials.ProtocolInfo{}
}
func (c *authorityCheckCreds) Clone() credentials.TransportCredentials {
	return c
}

// This test makes sure that the authority client handshake gets is the endpoint
// in dial target, not the resolved ip address.
func (s) TestCredsHandshakeAuthority(t *testing.T) {
	lis, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatal(err)
	}
	defer lis.Close()

	cred := &authorityCheckCreds{}
	s := grpc.NewServer()
	go s.Serve(lis)
	defer s.Stop()

	r := manual.NewBuilderWithScheme("whatever")

	r.InitialState(resolver.State{
		Addresses: []resolver.Address{{Addr: lis.Addr().String()}},
	})

	cc, err := grpc.NewClient(r.Scheme()+":///"+testAuthority, grpc.WithTransportCredentials(cred), grpc.WithResolvers(r))
	if err != nil {
		t.Fatalf("grpc.NewClient(%q) = %v", lis.Addr().String(), err)
	}
	cc.Connect()
	defer cc.Close()

	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	testutils.AwaitState(ctx, t, cc, connectivity.Ready)

	if cred.got != testAuthority {
		t.Fatalf("client creds got authority: %q, want: %q", cred.got, testAuthority)

Review comment (suggestion):
I think calling `s.Stop()` below would automatically close the listener. This may be redundant. Same comment for other places listener is being closed explicitly.

Code after:
const testAuthority = "test.auth.ori.ty"

type authorityCheckCreds struct {
	credentials.TransportCredentials
	got string
}

func (c *authorityCheckCreds) ClientHandshake(_ context.Context, authority string, rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
	c.got = authority
	return rawConn, nil, nil
}
func (c *authorityCheckCreds) Info() credentials.ProtocolInfo {
	return credentials.ProtocolInfo{}
}
func (c *authorityCheckCreds) Clone() credentials.TransportCredentials {
	return c
}

// This test makes sure that the authority client handshake gets is the endpoint
// in dial target, not the resolved ip address.
func (s) TestCredsHandshakeAuthority(t *testing.T) {
	lis, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatal(err)
	}
	cred := &authorityCheckCreds{}
	s := grpc.NewServer()
	go s.Serve(lis)
	defer s.Stop()

	r := manual.NewBuilderWithScheme("whatever")

	cc, err := grpc.NewClient(r.Scheme()+":///"+testAuthority, grpc.WithTransportCredentials(cred), grpc.WithResolvers(r))
	if err != nil {
		t.Fatalf("grpc.NewClient(%q) = %v", lis.Addr().String(), err)
	}
	defer cc.Close()
	cc.Connect()
	r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: lis.Addr().String()}}})

	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	testutils.AwaitState(ctx, t, cc, connectivity.Ready)

	if cred.got != testAuthority {
		t.Fatalf("client creds got authority: %q, want: %q", cred.got, testAuthority)
	}
}

// This test makes sure that the authority client handshake gets is the endpoint
// of the ServerName of the address when it is set.