Repository: grpc/grpc-go
File: credentials/credentials_ext_test.go (Go)

Code before:
func (FakeAuthInfo) ValidateAuthority(authority string) error {
	if authority == "auth.test.example.com" {
		return nil
	} else {
		return fmt.Errorf("invalid authority")
	}
}

// Clone creates a copy of FakeCredsWithAuthValidator.
func (c *FakeCredsWithAuthValidator) Clone() credentials.TransportCredentials {
	return c
}

// Info provides protocol information.
func (c *FakeCredsWithAuthValidator) Info() credentials.ProtocolInfo {
	return credentials.ProtocolInfo{}
}

// OverrideServerName overrides the server name used for verification.
func (c *FakeCredsWithAuthValidator) OverrideServerName(serverName string) error {
	return nil
}

// ServerHandshake performs the server-side handshake.
// Returns a test AuthInfo object to satisfy the interface requirements.
func (c *FakeCredsWithAuthValidator) ServerHandshake(rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
	return rawConn, FakeAuthInfo{}, nil
}

// TestCorrectAuthorityWithCustomCreds tests the CallAuthority call option
// with custom credentials that implement AuthorityValidator and verifies
// it with both correct and incorrect authority override.
func (s) TestCorrectAuthorityWithCustomCreds(t *testing.T) {
	tests := []struct {
		name         string
		creds        credentials.TransportCredentials
		expectedAuth string
		wantRPCError bool
	}{
		{
			name:         "CorrectAuthorityWithFakeCreds",
			expectedAuth: "auth.test.example.com",
			creds:        &FakeCredsWithAuthValidator{},
			wantRPCError: false,
		},
		{
			name:         "IncorrectAuthorityWithFakeCreds",
			expectedAuth: "auth.example.com",
			creds:        &FakeCredsWithAuthValidator{},
			wantRPCError: true,
		},

Review comment (security):
Basically, we can have a single `testCreds` type that contains a boolean field that indicates whether or not it should return an `AuthInfo` that implements `AuthorityValidator` or not. We don't need two test credentials types here.

Code after:
// override is specified.
// In both cases, the RPC is expected to fail with an `UNAVAILABLE` status code.
func (s) TestAuthorityValidationFailureWithCustomCreds(t *testing.T) {
	tests := []struct {
		name      string
		creds     credentials.TransportCredentials
		authority string
	}{
		{
			name:      "IncorrectAuthorityWithFakeCreds",
			authority: "auth.example.com",
			creds:     &testCreds{authority: "auth.test.example.com"},
		},
		{
			name:      "FakeCredsWithNoAuthValidator",
			creds:     &testCreds{},
			authority: "auth.test.example.com",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			serverCalled := make(chan struct{})
			ss := stubserver.StubServer{
				EmptyCallF: func(ctx context.Context, _ *testpb.Empty) (*testpb.Empty, error) {
					close(serverCalled)
					return nil, nil
				},
			}
			if err := ss.StartServer(); err != nil {
				t.Fatalf("Failed to start stub server: %v", err)
			}
			defer ss.Stop()

			cc, err := grpc.NewClient(ss.Address, grpc.WithTransportCredentials(tt.creds))
			if err != nil {
				t.Fatalf("grpc.NewClient(%q) = %v", ss.Address, err)
			}
			defer cc.Close()

			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
			defer cancel()
			if _, err = testgrpc.NewTestServiceClient(cc).EmptyCall(ctx, &testpb.Empty{}, grpc.CallAuthority(tt.authority)); status.Code(err) != codes.Unavailable {
				t.Fatalf("EmptyCall() returned status %v, want %v", status.Code(err), codes.Unavailable)
			}
			select {
			case <-serverCalled:
				t.Fatalf("Server should not have been called")
			case <-time.After(defaultTestShortTimeout):
			}
		})
	}