Improve key manager template

Signed-off-by: Guilherme Carvalho <guilhermocc@proton.me>

templates/agent/keymanager/keymanager.go

@@ -50,22 +50,8 @@ type Plugin struct {
 	logger hclog.Logger
 }
 
-// SetLogger is called by the framework when the plugin is loaded and provides
-// the plugin with a logger wired up to SPIRE's logging facilities.
-// TODO: Remove if the plugin does not need the logger.
-func (p *Plugin) SetLogger(logger hclog.Logger) {
-	p.logger = logger
-}
-
-// BrokerHostServices is called by the framework when the plugin is loaded to
-// give the plugin a chance to obtain clients to SPIRE host services.
-// TODO: Remove if the plugin does not need host services.
-func (p *Plugin) BrokerHostServices(broker pluginsdk.ServiceBroker) error {
-	// TODO: Use the broker to obtain host service clients
-	return nil
-}
-
-// GenerateKey implements the KeyManager GenerateKey RPC
+// GenerateKey implements the KeyManager GenerateKey RPC. Generates a new private key with the given ID.
+// If a key already exists under that ID, it is overwritten and given a different fingerprint.
 func (p *Plugin) GenerateKey(ctx context.Context, req *keymanagerv1.GenerateKeyRequest) (*keymanagerv1.GenerateKeyResponse, error) {
 	config, err := p.getConfig()
 	if err != nil {
@@ -80,7 +66,8 @@ func (p *Plugin) GenerateKey(ctx context.Context, req *keymanagerv1.GenerateKeyR
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-// GetPublicKey implements the KeyManager GetPublicKey RPC
+// GetPublicKey implements the KeyManager GetPublicKey RPC. Gets the public key information for the private key managed
+// by the plugin with the given ID. If a key with the given ID does not exist, NOT_FOUND is returned.
 func (p *Plugin) GetPublicKey(ctx context.Context, req *keymanagerv1.GetPublicKeyRequest) (*keymanagerv1.GetPublicKeyResponse, error) {
 	config, err := p.getConfig()
 	if err != nil {
@@ -95,7 +82,8 @@ func (p *Plugin) GetPublicKey(ctx context.Context, req *keymanagerv1.GetPublicKe
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-// GetPublicKeys implements the KeyManager GetPublicKeys RPC
+// GetPublicKeys implements the KeyManager GetPublicKeys RPC. Gets all public key information for the private keys
+// managed by the plugin.
 func (p *Plugin) GetPublicKeys(ctx context.Context, req *keymanagerv1.GetPublicKeysRequest) (*keymanagerv1.GetPublicKeysResponse, error) {
 	config, err := p.getConfig()
 	if err != nil {
@@ -110,7 +98,9 @@ func (p *Plugin) GetPublicKeys(ctx context.Context, req *keymanagerv1.GetPublicK
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-// SignData implements the KeyManager SignData RPC
+// SignData implements the KeyManager SignData RPC. Signs data with the private key identified by the given ID. If a key
+// with the given ID does not exist, NOT_FOUND is returned. The response contains the signed data and the fingerprint of
+// the key used to sign the data. See the PublicKey message for more details on the role of the fingerprint.
 func (p *Plugin) SignData(ctx context.Context, req *keymanagerv1.SignDataRequest) (*keymanagerv1.SignDataResponse, error) {
 	config, err := p.getConfig()
 	if err != nil {
@@ -125,6 +115,21 @@ func (p *Plugin) SignData(ctx context.Context, req *keymanagerv1.SignDataRequest
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
+// SetLogger is called by the framework when the plugin is loaded and provides
+// the plugin with a logger wired up to SPIRE's logging facilities.
+// TODO: Remove if the plugin does not need the logger.
+func (p *Plugin) SetLogger(logger hclog.Logger) {
+	p.logger = logger
+}
+
+// BrokerHostServices is called by the framework when the plugin is loaded to
+// give the plugin a chance to obtain clients to SPIRE host services.
+// TODO: Remove if the plugin does not need host services.
+func (p *Plugin) BrokerHostServices(broker pluginsdk.ServiceBroker) error {
+	// TODO: Use the broker to obtain host service clients
+	return nil
+}
+
 // Configure configures the plugin. This is invoked by SPIRE when the plugin is
 // first loaded. In the future, it may be invoked to reconfigure the plugin.
 // As such, it should replace the previous configuration atomically.

templates/agent/keymanager/keymanager_test.go

@@ -1,6 +1,7 @@
 package keymanager_test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/spiffe/spire-plugin-sdk/pluginsdk"
@@ -8,6 +9,8 @@ import (
 	keymanagerv1 "github.com/spiffe/spire-plugin-sdk/proto/spire/plugin/agent/keymanager/v1"
 	configv1 "github.com/spiffe/spire-plugin-sdk/proto/spire/service/common/config/v1"
 	"github.com/spiffe/spire-plugin-sdk/templates/agent/keymanager"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func Test(t *testing.T) {
@@ -32,5 +35,24 @@ func Test(t *testing.T) {
 		},
 	})
 
-	// TODO: Invoke methods on the clients and assert the results
+	ctx := context.Background()
+
+	// TODO: Remove if no configuration is required.
+	_, err := configClient.Configure(ctx, &configv1.ConfigureRequest{
+		CoreConfiguration: &configv1.CoreConfiguration{TrustDomain: "example.org"},
+		HclConfiguration:  `{}`,
+	})
+	assert.NoError(t, err)
+
+	require.True(t, kmClient.IsInitialized())
+
+	// TODO: Make assertions using the desired plugin behavior.
+	_, err = kmClient.GenerateKey(ctx, &keymanagerv1.GenerateKeyRequest{})
+	assert.EqualError(t, err, "rpc error: code = Unimplemented desc = not implemented")
+	_, err = kmClient.GetPublicKeys(ctx, &keymanagerv1.GetPublicKeysRequest{})
+	assert.EqualError(t, err, "rpc error: code = Unimplemented desc = not implemented")
+	_, err = kmClient.GetPublicKey(ctx, &keymanagerv1.GetPublicKeyRequest{})
+	assert.EqualError(t, err, "rpc error: code = Unimplemented desc = not implemented")
+	_, err = kmClient.SignData(ctx, &keymanagerv1.SignDataRequest{})
+	assert.EqualError(t, err, "rpc error: code = Unimplemented desc = not implemented")
 }

templates/agent/nodeattestor/nodeattestor.go

@@ -51,21 +51,6 @@ type Plugin struct {
 	logger hclog.Logger
 }
 
-// SetLogger is called by the framework when the plugin is loaded and provides
-// the plugin with a logger wired up to SPIRE's logging facilities.
-// TODO: Remove if the plugin does not need the logger.
-func (p *Plugin) SetLogger(logger hclog.Logger) {
-	p.logger = logger
-}
-
-// BrokerHostServices is called by the framework when the plugin is loaded to
-// give the plugin a chance to obtain clients to SPIRE host services.
-// TODO: Remove if the plugin does not need host services.
-func (p *Plugin) BrokerHostServices(broker pluginsdk.ServiceBroker) error {
-	// TODO: Use the broker to obtain host service clients
-	return nil
-}
-
 // AidAttestation implements the NodeAttestor AidAttestation RPC. AidAttestation facilitates attestation by returning
 // the attestation payload and participating in attestation challenge/response. This RPC uses a bidirectional stream for
 // communication.
@@ -83,6 +68,21 @@ func (p *Plugin) AidAttestation(stream nodeattestorv1.NodeAttestor_AidAttestatio
 	return status.Error(codes.Unimplemented, "not implemented")
 }
 
+// SetLogger is called by the framework when the plugin is loaded and provides
+// the plugin with a logger wired up to SPIRE's logging facilities.
+// TODO: Remove if the plugin does not need the logger.
+func (p *Plugin) SetLogger(logger hclog.Logger) {
+	p.logger = logger
+}
+
+// BrokerHostServices is called by the framework when the plugin is loaded to
+// give the plugin a chance to obtain clients to SPIRE host services.
+// TODO: Remove if the plugin does not need host services.
+func (p *Plugin) BrokerHostServices(broker pluginsdk.ServiceBroker) error {
+	// TODO: Use the broker to obtain host service clients
+	return nil
+}
+
 // Configure configures the plugin. This is invoked by SPIRE when the plugin is
 // first loaded. In the future, it may be invoked to reconfigure the plugin.
 // As such, it should replace the previous configuration atomically.

templates/agent/nodeattestor/nodeattestor_test.go

@@ -47,7 +47,7 @@ func Test(t *testing.T) {
 	require.True(t, naClient.IsInitialized())
 
 	// TODO: Make assertions using the desired plugin behavior.
-	resp, err := naClient.AidAttestation(context.Background())
+	resp, err := naClient.AidAttestation(ctx)
 	require.NoError(t, err)
 	_, err = resp.Recv()
 	assert.EqualError(t, err, "rpc error: code = Unimplemented desc = not implemented")

templates/server/keymanager/keymanager.go

@@ -50,22 +50,8 @@ type Plugin struct {
 	logger hclog.Logger
 }
 
-// SetLogger is called by the framework when the plugin is loaded and provides
-// the plugin with a logger wired up to SPIRE's logging facilities.
-// TODO: Remove if the plugin does not need the logger.
-func (p *Plugin) SetLogger(logger hclog.Logger) {
-	p.logger = logger
-}
-
-// BrokerHostServices is called by the framework when the plugin is loaded to
-// give the plugin a chance to obtain clients to SPIRE host services.
-// TODO: Remove if the plugin does not need host services.
-func (p *Plugin) BrokerHostServices(broker pluginsdk.ServiceBroker) error {
-	// TODO: Use the broker to obtain host service clients
-	return nil
-}
-
-// GenerateKey implements the KeyManager GenerateKey RPC
+// GenerateKey implements the KeyManager GenerateKey RPC. Generates a new private key with the given ID.
+// If a key already exists under that ID, it is overwritten and given a different fingerprint.
 func (p *Plugin) GenerateKey(ctx context.Context, req *keymanagerv1.GenerateKeyRequest) (*keymanagerv1.GenerateKeyResponse, error) {
 	config, err := p.getConfig()
 	if err != nil {
@@ -80,7 +66,8 @@ func (p *Plugin) GenerateKey(ctx context.Context, req *keymanagerv1.GenerateKeyR
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-// GetPublicKey implements the KeyManager GetPublicKey RPC
+// GetPublicKey implements the KeyManager GetPublicKey RPC. Gets the public key information for the private key managed
+// by the plugin with the given ID. If a key with the given ID does not exist, NOT_FOUND is returned.
 func (p *Plugin) GetPublicKey(ctx context.Context, req *keymanagerv1.GetPublicKeyRequest) (*keymanagerv1.GetPublicKeyResponse, error) {
 	config, err := p.getConfig()
 	if err != nil {
@@ -95,7 +82,8 @@ func (p *Plugin) GetPublicKey(ctx context.Context, req *keymanagerv1.GetPublicKe
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-// GetPublicKeys implements the KeyManager GetPublicKeys RPC
+// GetPublicKeys implements the KeyManager GetPublicKeys RPC. Gets all public key information for the private keys
+// managed by the plugin.
 func (p *Plugin) GetPublicKeys(ctx context.Context, req *keymanagerv1.GetPublicKeysRequest) (*keymanagerv1.GetPublicKeysResponse, error) {
 	config, err := p.getConfig()
 	if err != nil {
@@ -110,7 +98,9 @@ func (p *Plugin) GetPublicKeys(ctx context.Context, req *keymanagerv1.GetPublicK
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
-// SignData implements the KeyManager SignData RPC
+// SignData implements the KeyManager SignData RPC. Signs data with the private key identified by the given ID. If a key
+// with the given ID does not exist, NOT_FOUND is returned. The response contains the signed data and the fingerprint of
+// the key used to sign the data. See the PublicKey message for more details on the role of the fingerprint.
 func (p *Plugin) SignData(ctx context.Context, req *keymanagerv1.SignDataRequest) (*keymanagerv1.SignDataResponse, error) {
 	config, err := p.getConfig()
 	if err != nil {
@@ -125,6 +115,21 @@ func (p *Plugin) SignData(ctx context.Context, req *keymanagerv1.SignDataRequest
 	return nil, status.Error(codes.Unimplemented, "not implemented")
 }
 
+// SetLogger is called by the framework when the plugin is loaded and provides
+// the plugin with a logger wired up to SPIRE's logging facilities.
+// TODO: Remove if the plugin does not need the logger.
+func (p *Plugin) SetLogger(logger hclog.Logger) {
+	p.logger = logger
+}
+
+// BrokerHostServices is called by the framework when the plugin is loaded to
+// give the plugin a chance to obtain clients to SPIRE host services.
+// TODO: Remove if the plugin does not need host services.
+func (p *Plugin) BrokerHostServices(broker pluginsdk.ServiceBroker) error {
+	// TODO: Use the broker to obtain host service clients
+	return nil
+}
+
 // Configure configures the plugin. This is invoked by SPIRE when the plugin is
 // first loaded. In the future, it may be invoked to reconfigure the plugin.
 // As such, it should replace the previous configuration atomically.

templates/server/keymanager/keymanager_test.go

@@ -1,6 +1,7 @@
 package keymanager_test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/spiffe/spire-plugin-sdk/pluginsdk"
@@ -8,6 +9,8 @@ import (
 	keymanagerv1 "github.com/spiffe/spire-plugin-sdk/proto/spire/plugin/server/keymanager/v1"
 	configv1 "github.com/spiffe/spire-plugin-sdk/proto/spire/service/common/config/v1"
 	"github.com/spiffe/spire-plugin-sdk/templates/server/keymanager"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func Test(t *testing.T) {
@@ -32,5 +35,24 @@ func Test(t *testing.T) {
 		},
 	})
 
-	// TODO: Invoke methods on the clients and assert the results
+	ctx := context.Background()
+
+	// TODO: Remove if no configuration is required.
+	_, err := configClient.Configure(ctx, &configv1.ConfigureRequest{
+		CoreConfiguration: &configv1.CoreConfiguration{TrustDomain: "example.org"},
+		HclConfiguration:  `{}`,
+	})
+	assert.NoError(t, err)
+
+	require.True(t, kmClient.IsInitialized())
+
+	// TODO: Make assertions using the desired plugin behavior.
+	_, err = kmClient.GenerateKey(ctx, &keymanagerv1.GenerateKeyRequest{})
+	assert.EqualError(t, err, "rpc error: code = Unimplemented desc = not implemented")
+	_, err = kmClient.GetPublicKeys(ctx, &keymanagerv1.GetPublicKeysRequest{})
+	assert.EqualError(t, err, "rpc error: code = Unimplemented desc = not implemented")
+	_, err = kmClient.GetPublicKey(ctx, &keymanagerv1.GetPublicKeyRequest{})
+	assert.EqualError(t, err, "rpc error: code = Unimplemented desc = not implemented")
+	_, err = kmClient.SignData(ctx, &keymanagerv1.SignDataRequest{})
+	assert.EqualError(t, err, "rpc error: code = Unimplemented desc = not implemented")
 }

templates/server/nodeattestor/nodeattestor.go

@@ -50,21 +50,6 @@ type Plugin struct {
 	logger hclog.Logger
 }
 
-// SetLogger is called by the framework when the plugin is loaded and provides
-// the plugin with a logger wired up to SPIRE's logging facilities.
-// TODO: Remove if the plugin does not need the logger.
-func (p *Plugin) SetLogger(logger hclog.Logger) {
-	p.logger = logger
-}
-
-// BrokerHostServices is called by the framework when the plugin is loaded to
-// give the plugin a chance to obtain clients to SPIRE host services.
-// TODO: Remove if the plugin does not need host services.
-func (p *Plugin) BrokerHostServices(broker pluginsdk.ServiceBroker) error {
-	// TODO: Use the broker to obtain host service clients
-	return nil
-}
-
 // Attest implements the NodeAttestor Attest RPC. Attest attests attestation payload received from the agent and
 // optionally participates in challenge/response attestation mechanics. This RPC uses a bidirectional stream for
 // communication.
@@ -82,6 +67,21 @@ func (p *Plugin) Attest(stream nodeattestorv1.NodeAttestor_AttestServer) error {
 	return status.Error(codes.Unimplemented, "not implemented")
 }
 
+// SetLogger is called by the framework when the plugin is loaded and provides
+// the plugin with a logger wired up to SPIRE's logging facilities.
+// TODO: Remove if the plugin does not need the logger.
+func (p *Plugin) SetLogger(logger hclog.Logger) {
+	p.logger = logger
+}
+
+// BrokerHostServices is called by the framework when the plugin is loaded to
+// give the plugin a chance to obtain clients to SPIRE host services.
+// TODO: Remove if the plugin does not need host services.
+func (p *Plugin) BrokerHostServices(broker pluginsdk.ServiceBroker) error {
+	// TODO: Use the broker to obtain host service clients
+	return nil
+}
+
 // Configure configures the plugin. This is invoked by SPIRE when the plugin is
 // first loaded. In the future, it may be invoked to reconfigure the plugin.
 // As such, it should replace the previous configuration atomically.

templates/server/nodeattestor/nodeattestor_test.go

@@ -47,7 +47,7 @@ func Test(t *testing.T) {
 	require.True(t, naClient.IsInitialized())
 
 	// TODO: Make assertions using the desired plugin behavior.
-	resp, err := naClient.Attest(context.Background())
+	resp, err := naClient.Attest(ctx)
 	require.NoError(t, err)
 	_, err = resp.Recv()
 	assert.EqualError(t, err, "rpc error: code = Unimplemented desc = not implemented")