Handle nil keyring gracefully and fix gRPC v3 broker startup race

pkg/plugins/core_cli_helper.go

@@ -2,6 +2,7 @@ package plugins
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"runtime"
 	"sync"
@@ -163,15 +164,29 @@ type pendingKeychainValue struct {
 }
 
 var (
+	errKeychainUnavailable          = errors.New("system keychain is unavailable")
 	keychainVisibilityRetryTimeout  = 1500 * time.Millisecond
 	keychainVisibilityRetryEnabled  = runtime.GOOS == "darwin"
 	keychainVisibilityNow           = time.Now
 	keychainVisibilityPendingMu     sync.Mutex
 	keychainVisibilityPendingValues = map[string]pendingKeychainValue{}
 )
 
+func keyRing() (keyring.Keyring, error) {
+	if config.KeyRing == nil {
+		return nil, errKeychainUnavailable
+	}
+
+	return config.KeyRing, nil
+}
+
 func readKeychainPassword(key string) (string, bool, error) {
-	item, err := config.KeyRing.Get(key)
+	ring, err := keyRing()
+	if err != nil {
+		return "", false, err
+	}
+
+	item, err := ring.Get(key)
 	if err == nil {
 		return string(item.Data), true, nil
 	}
@@ -267,7 +282,12 @@ func (h *coreCLIHelper) KeychainGetPassword(key string) (string, bool, error) {
 
 // KeychainSetPassword stores a password in the system keychain.
 func (h *coreCLIHelper) KeychainSetPassword(key string, value string) error {
-	if err := config.KeyRing.Set(keyring.Item{
+	ring, err := keyRing()
+	if err != nil {
+		return err
+	}
+
+	if err := ring.Set(keyring.Item{
 		Key:   key,
 		Data:  []byte(value),
 		Label: key,
@@ -283,13 +303,18 @@ func (h *coreCLIHelper) KeychainSetPassword(key string, value string) error {
 func (h *coreCLIHelper) KeychainDeletePassword(key string) (bool, error) {
 	clearPendingKeychainValue(key)
 
-	existingKeys, err := config.KeyRing.Keys()
+	ring, err := keyRing()
+	if err != nil {
+		return false, err
+	}
+
+	existingKeys, err := ring.Keys()
 	if err != nil {
 		return false, err
 	}
 	for _, k := range existingKeys {
 		if k == key {
-			if err := config.KeyRing.Remove(key); err != nil {
+			if err := ring.Remove(key); err != nil {
 				return false, err
 			}
 			return true, nil
@@ -300,7 +325,12 @@ func (h *coreCLIHelper) KeychainDeletePassword(key string) (bool, error) {
 
 // KeychainFindCredentials lists all keys stored in the keychain for this service.
 func (h *coreCLIHelper) KeychainFindCredentials() ([]string, error) {
-	return config.KeyRing.Keys()
+	ring, err := keyRing()
+	if err != nil {
+		return nil, err
+	}
+
+	return ring.Keys()
 }
 
 // RunPeerPlugin looks up and runs the named plugin with the given arguments.

pkg/plugins/core_cli_helper_test.go

@@ -146,6 +146,20 @@ func TestKeychainGetPasswordReturnsNotFoundWithoutRetry(t *testing.T) {
 	require.Equal(t, 1, ring.getCalls)
 }
 
+func TestKeychainGetPasswordReturnsErrorWhenKeyringUnavailable(t *testing.T) {
+	originalKeyRing := config.KeyRing
+	config.KeyRing = nil
+	t.Cleanup(func() {
+		config.KeyRing = originalKeyRing
+	})
+
+	coreCLIHelper := NewCoreCLIHelper(context.Background(), nil, afero.NewMemMapFs())
+	value, found, err := coreCLIHelper.KeychainGetPassword("missing.key")
+	require.ErrorIs(t, err, errKeychainUnavailable)
+	require.False(t, found)
+	require.Empty(t, value)
+}
+
 func TestKeychainGetPasswordReturnsUnexpectedError(t *testing.T) {
 	originalKeyRing := config.KeyRing
 	originalEnabled := keychainVisibilityRetryEnabled
@@ -210,6 +224,18 @@ func TestKeychainSetPasswordMakesRecentWriteVisibleWhenKeychainHasNotCaughtUp(t
 	require.Equal(t, 1, ring.getCalls)
 }
 
+func TestKeychainSetPasswordReturnsErrorWhenKeyringUnavailable(t *testing.T) {
+	originalKeyRing := config.KeyRing
+	config.KeyRing = nil
+	t.Cleanup(func() {
+		config.KeyRing = originalKeyRing
+	})
+
+	coreCLIHelper := NewCoreCLIHelper(context.Background(), nil, afero.NewMemMapFs())
+	err := coreCLIHelper.KeychainSetPassword("test.key", "sk_test_123")
+	require.ErrorIs(t, err, errKeychainUnavailable)
+}
+
 func TestKeychainGetPasswordPrefersRecentWriteOverStaleVisibleValue(t *testing.T) {
 	originalKeyRing := config.KeyRing
 	originalEnabled := keychainVisibilityRetryEnabled

pkg/plugins/interface_grpc_3.go

@@ -2,6 +2,9 @@ package plugins
 
 import (
 	"context"
+	"errors"
+	"sync"
+	"time"
 
 	hcplugin "github.com/hashicorp/go-plugin"
 	"google.golang.org/grpc"
@@ -33,40 +36,105 @@ func (p *CLIPluginV3) GRPCServer(broker *hcplugin.GRPCBroker, s *grpc.Server) er
 func (p *CLIPluginV3) GRPCClient(ctx context.Context, broker *hcplugin.GRPCBroker, c *grpc.ClientConn) (interface{}, error) {
 	return &GRPCClientV3{
 		client: proto.NewMainClient(c),
-		broker: broker,
+		broker: grpcBrokerAdapter{broker: broker},
 	}, nil
 }
 
 // GRPCClientV3 is an implementation of the gRPC client that talks over gRPC.
 type GRPCClientV3 struct {
 	client proto.MainClient
+	broker grpcBrokerClient
+}
+
+type grpcBrokerServer interface {
+	AcceptAndServe(id uint32, newGRPCServer func([]grpc.ServerOption) *grpc.Server)
+}
+
+type grpcBrokerClient interface {
+	grpcBrokerServer
+	nextID() uint32
+}
+
+type grpcBrokerAdapter struct {
 	broker *hcplugin.GRPCBroker
 }
 
-// RunCommand calls the RPC.
-func (m *GRPCClientV3) RunCommand(additionalInfo *proto.AdditionalInfo, args []string, coreCLIHelper CoreCLIHelper) error {
-	coreCLIHelperServer := &CoreCLIHelperServer{Impl: coreCLIHelper}
+func (b grpcBrokerAdapter) nextID() uint32 {
+	return b.broker.NextId()
+}
+
+func (b grpcBrokerAdapter) AcceptAndServe(id uint32, newGRPCServer func([]grpc.ServerOption) *grpc.Server) {
+	b.broker.AcceptAndServe(id, newGRPCServer)
+}
 
-	var s *grpc.Server
-	serverFunc := func(opts []grpc.ServerOption) *grpc.Server {
-		s = grpc.NewServer(opts...)
-		proto.RegisterCoreCLIHelperServer(s, coreCLIHelperServer)
-		return s
+var errCoreCLIHelperBrokerServerStart = errors.New("failed to start CoreCLIHelper broker server")
+var coreCLIHelperBrokerPublishDelay = 25 * time.Millisecond
+var coreCLIHelperBrokerServerStartTimeout = 5 * time.Second
+
+func startCoreCLIHelperBrokerServer(broker grpcBrokerServer, brokerID uint32, coreCLIHelper CoreCLIHelper) (func(), error) {
+	startedCh := make(chan struct{})
+	doneCh := make(chan struct{})
+
+	var server *grpc.Server
+	go func() {
+		defer close(doneCh)
+
+		broker.AcceptAndServe(brokerID, func(opts []grpc.ServerOption) *grpc.Server {
+			server = grpc.NewServer(opts...)
+			proto.RegisterCoreCLIHelperServer(server, &CoreCLIHelperServer{Impl: coreCLIHelper})
+			close(startedCh)
+			return server
+		})
+	}()
+
+	select {
+	case <-startedCh:
+	case <-doneCh:
+		return nil, errCoreCLIHelperBrokerServerStart
+	case <-time.After(coreCLIHelperBrokerServerStartTimeout):
+		return nil, errCoreCLIHelperBrokerServerStart
 	}
 
-	brokerID := m.broker.NextId()
-	go m.broker.AcceptAndServe(brokerID, serverFunc)
+	var cleanupOnce sync.Once
+	cleanup := func() {
+		cleanupOnce.Do(func() {
+			if server != nil {
+				server.Stop()
+			}
+		})
+	}
 
-	_, err := m.client.RunCommand(context.Background(), &proto.RunCommandRequest{
-		AdditionalInfo:  additionalInfo,
-		Args:            args,
-		CoreCliHelperId: brokerID,
-	})
+	return cleanup, nil
+}
+
+// RunCommand calls the RPC.
+func (m *GRPCClientV3) RunCommand(additionalInfo *proto.AdditionalInfo, args []string, coreCLIHelper CoreCLIHelper) error {
+	brokerID := m.broker.nextID()
+	errCh := make(chan error, 1)
+	go func() {
+		_, err := m.client.RunCommand(context.Background(), &proto.RunCommandRequest{
+			AdditionalInfo:  additionalInfo,
+			Args:            args,
+			CoreCliHelperId: brokerID,
+		})
+		errCh <- err
+	}()
+
+	// Non-Go plugins can drop unsolicited broker connection info if it arrives
+	// before they register a pending dial for this service ID.
+	time.Sleep(coreCLIHelperBrokerPublishDelay)
+
+	cleanup, err := startCoreCLIHelperBrokerServer(m.broker, brokerID, coreCLIHelper)
+	if err != nil {
+		return err
+	}
+	defer cleanup()
+
+	err = <-errCh
 	if err != nil {
 		return err
 	}
 
-	s.Stop()
 	return nil
 }