Repository: google/exposure-notifications-server
Pull request #847 (2020-08-07 22:32:32)

Add interfaces for various key handlers
* Add new interfaces for in-memory key management (used for testing)

**Release Note**

```release-note
NONE
```

google-oss-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/847#" title="Author self-approved">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/main/OWNERS)~~ [sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/internal/integration/helpers.go b/internal/integration/helpers.go
index e5b3e43..e9f6554 100644
--- a/internal/integration/helpers.go
+++ b/internal/integration/helpers.go
@@ -137,10 +137,10 @@ func testServer(tb testing.TB) (*serverenv.ServerEnv, *http.Client, testutil.JWT
 	if err != nil {
 		tb.Fatal(err)
 	}
-	if err := km.AddEncryptionKey("tokenkey"); err != nil {
+	if _, err := km.CreateEncryptionKey("tokenkey"); err != nil {
 		tb.Fatal(err)
 	}
-	if err := km.AddSigningKey("signingkey"); err != nil {
+	if _, err := km.CreateSigningKey("signingkey"); err != nil {
 		tb.Fatal(err)
 	}
 	// create an initial revision key.
diff --git a/internal/serverenv/env.go b/internal/serverenv/env.go
index 0f35045..caf86f8 100644
--- a/internal/serverenv/env.go
+++ b/internal/serverenv/env.go
@@ -178,7 +178,9 @@ func (s *ServerEnv) Close(ctx context.Context) error {
 	}
 
 	if s.observabilityExporter != nil {
-		s.observabilityExporter.Flush()
+		if err := s.observabilityExporter.Close(); err != nil {
+			return nil
+		}
 	}
 
 	return nil
diff --git a/internal/setup/setup.go b/internal/setup/setup.go
index 9d32deb..28574f8 100644
--- a/internal/setup/setup.go
+++ b/internal/setup/setup.go
@@ -170,18 +170,21 @@ func SetupWith(ctx context.Context, config interface{}, l envconfig.Lookuper) (*
 	// Configure and initialize the observability exporter.
 	if provider, ok := config.(ObservabilityExporterConfigProvider); ok {
 		logger.Info("configuring observability exporter")
+
 		oeConfig := provider.ObservabilityExporterConfig()
 		oe, err := observability.NewFromEnv(ctx, oeConfig)
 		if err != nil {
-			return nil, fmt.Errorf("unable to create ObservabilityExporter provider: %v", err)
+			return nil, fmt.Errorf("unable to create observability provider: %v", err)
 		}
-		if err := oe.InitExportOnce(); err != nil {
-			return nil, fmt.Errorf("error initializing observability exporter: %v", err)
+		if err := oe.StartExporter(); err != nil {
+			return nil, fmt.Errorf("failed to start observability: %w", err)
 		}
+		exporter := serverenv.WithObservabilityExporter(oe)
 
-		serverEnvOpts = append(serverEnvOpts, serverenv.WithObservabilityExporter(oe))
+		// Update serverEnv setup.
+		serverEnvOpts = append(serverEnvOpts, exporter)
 
-		logger.Infow("observability exporter", "config", oeConfig)
+		logger.Infow("observability", "config", oeConfig)
 	}
 
 	// Configure blob storage.
diff --git a/pkg/keys/in_memory.go b/pkg/keys/in_memory.go
index 47498ea..5b43ae3 100644
--- a/pkg/keys/in_memory.go
+++ b/pkg/keys/in_memory.go
@@ -27,6 +27,12 @@ import (
 	"sync"
 )
 
+var _ KeyManager = (*InMemory)(nil)
+var _ SigningKeyCreator = (*InMemory)(nil)
+var _ SigningKeyAdder = (*InMemory)(nil)
+var _ EncryptionKeyCreator = (*InMemory)(nil)
+var _ EncryptionKeyAdder = (*InMemory)(nil)
+
 // InMemory is useful for testing. Do NOT use in a running system as all
 // keys are only kept in memory and will be lost across server reboots.
 type InMemory struct {
@@ -43,42 +49,68 @@ func NewInMemory(ctx context.Context) (*InMemory, error) {
 	}, nil
 }
 
-// AddSigningKey generates a new ECDSA P256 Signing Key identified by
+// CreateSigningKey generates a new ECDSA P256 Signing Key identified by
 // the provided keyID
-func (k *InMemory) AddSigningKey(keyID string) error {
+func (k *InMemory) CreateSigningKey(keyID string) (*ecdsa.PrivateKey, error) {
 	k.mu.Lock()
 	defer k.mu.Unlock()
 
 	if _, ok := k.signingKeys[keyID]; ok {
-		return fmt.Errorf("key already exists: %v", keyID)
+		return nil, fmt.Errorf("key already exists: %v", keyID)
 	}
 
 	pk, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
 	if err != nil {
-		return fmt.Errorf("unable to generate private key: %w", err)
+		return nil, fmt.Errorf("unable to generate private key: %w", err)
+	}
+
+	k.signingKeys[keyID] = pk
+	return pk, nil
+}
+
+// AddSigningKey adds a new ECDSA P256 Signing Key identified by the provided
+// keyID.
+func (k *InMemory) AddSigningKey(keyID string, pk *ecdsa.PrivateKey) error {
+	k.mu.Lock()
+	defer k.mu.Unlock()
+
+	if _, ok := k.signingKeys[keyID]; ok {
+		return fmt.Errorf("key already exists: %v", keyID)
 	}
 
 	k.signingKeys[keyID] = pk
 	return nil
 }
 
-// AddEncryptionKey generates a new encryption key identified by
-// the provided keyID.
-func (k *InMemory) AddEncryptionKey(keyID string) error {
+// CreateEncryptionKey generates and stores new encryption key identified by the
+// provided keyID.
+func (k *InMemory) CreateEncryptionKey(keyID string) ([]byte, error) {
 	k.mu.Lock()
 	defer k.mu.Unlock()
 
 	if _, ok := k.cryptoKeys[keyID]; ok {
-		return fmt.Errorf("key already exists: %v", keyID)
+		return nil, fmt.Errorf("key already exists: %v", keyID)
 	}
 
 	key := make([]byte, 32)
 	if _, err := io.ReadFull(rand.Reader, key); err != nil {
-		return fmt.Errorf("failed to read random bytes: %w", err)
+		return nil, fmt.Errorf("failed to read random bytes: %w", err)
 	}
 
 	k.cryptoKeys[keyID] = key
+	return key, nil
+}
+
+// AddEncryptionKey stores the key on the system.
+func (k *InMemory) AddEncryptionKey(keyID string, key []byte) error {
+	k.mu.Lock()
+	defer k.mu.Unlock()
 
+	if _, ok := k.cryptoKeys[keyID]; ok {
+		return fmt.Errorf("key already exists: %v", keyID)
+	}
+
+	k.cryptoKeys[keyID] = key
 	return nil
 }
 
diff --git a/pkg/keys/keys.go b/pkg/keys/keys.go
index 1ae1e93..d8bdf15 100644
--- a/pkg/keys/keys.go
+++ b/pkg/keys/keys.go
@@ -23,6 +23,7 @@ package keys
 import (
 	"context"
 	"crypto"
+	"crypto/ecdsa"
 	"fmt"
 )
 
@@ -48,6 +49,26 @@ type KeyManager interface {
 	Decrypt(ctx context.Context, keyID string, ciphertext []byte, aad []byte) ([]byte, error)
 }
 
+// EncryptionKeyCreator supports creating encryption keys.
+type EncryptionKeyCreator interface {
+	CreateEncryptionKey(string) ([]byte, error)
+}
+
+// EncryptionKeyAdder supports creating encryption keys.
+type EncryptionKeyAdder interface {
+	AddEncryptionKey(string, []byte) error
+}
+
+// SigningKeyCreator supports creating signing keys.
+type SigningKeyCreator interface {
+	CreateSigningKey(string) (*ecdsa.PrivateKey, error)
+}
+
+// SigningKeyAdder supports creating signing keys.
+type SigningKeyAdder interface {
+	AddSigningKey(string, *ecdsa.PrivateKey) error
+}
+
 // KeyManagerFor returns the appropriate key manager for the given type.
 func KeyManagerFor(ctx context.Context, typ KeyManagerType) (KeyManager, error) {
 	switch typ {
diff --git a/pkg/observability/config.go b/pkg/observability/config.go
index 316d477..cd39a3e 100644
--- a/pkg/observability/config.go
+++ b/pkg/observability/config.go
@@ -29,30 +29,31 @@ const (
 type Config struct {
 	ExporterType ExporterType `env:"OBSERVABILITY_EXPORTER, default=STACKDRIVER"`
 
-	OpenCensusConfig
-	StackdriverConfig
-	OCAgentConfig
+	OpenCensus  *OpenCensusConfig
+	Stackdriver *StackdriverConfig
 }
 
 // OpenCensusConfig holds the configuration options for the open census exporter
 type OpenCensusConfig struct {
-	TraceProbabilitySampleRate float64 `env:"TRACE_PROBABILITY, default=0.40"`
+	SampleRate float64 `env:"TRACE_PROBABILITY, default=0.40"`
+
+	Insecure bool   `env:"OCAGENT_INSECURE"`
+	Endpoint string `env:"OCAGENT_TRACE_EXPORTER_ENDPOINT"`
 }
 
 // StackdriverConfig holds the configuration options for the stackdriver exporter
 type StackdriverConfig struct {
-	ProjectID string `env:"PROJECT_ID, default=$GOOGLE_CLOUD_PROJECT"`
-	// Knative+Cloud Run container contract env vars
-	// https://cloud.google.com/run/docs/reference/container-contract#env-vars
-	// If present, can be used to configured the Stackdriver MonitoredResource correctly.
-	Service  string `env:"K_SERVICE"`
-	Revision string `env:"K_REVISION"`
+	SampleRate float64 `env:"TRACE_PROBABILITY, default=0.40"`
 
-	Namespace string `env:"K_CONFIGURATION, default=ens"`
-}
+	ProjectID string `env:"PROJECT_ID, default=$GOOGLE_CLOUD_PROJECT"`
 
-// OCAgentConfig holds the configuration options for the default opencensus exporter
-type OCAgentConfig struct {
-	Insecure bool   `env:"OCAGENT_INSECURE, default=true"`
-	Endpoint string `env:"OCAGENT_TRACE_EXPORTER_ENDPOINT"`
+	// Knative+Cloud Run container contract envvars:
+	//
+	// https://cloud.google.com/run/docs/reference/container-contract#env-vars
+	//
+	// If present, can be used to configured the Stackdriver MonitoredResource
+	// correctly.
+	Service   string `env:"K_SERVICE"`
+	Revision  string `env:"K_REVISION"`
+	Namespace string `env:"K_CONFIGURATION, default=en"`
 }
diff --git a/pkg/observability/generic.go b/pkg/observability/generic.go
deleted file mode 100644
index 95102bd..0000000
--- a/pkg/observability/generic.go
+++ /dev/null
@@ -1,81 +0,0 @@
-// Copyright 2020 Google LLC
-//
-// 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 observability sets up and configures observability tools.
-package observability
-
-import (
-	"fmt"
-	"sync"
-
-	"go.opencensus.io/plugin/ocgrpc"
-	"go.opencensus.io/plugin/ochttp"
-	"go.opencensus.io/stats/view"
-	"go.opencensus.io/trace"
-)
-
-// Compile-time check to verify implements interface.
-var _ Exporter = (*GenericExporter)(nil)
-
-var initExporterOnce sync.Once
-
-type FlushFn func()
-
-type traceAndViewExporter interface {
-	trace.Exporter
-	view.Exporter
-}
-
-// GenericExporter is a standard implementation of an exporter that wraps the opencensus interfaces
-// with custom configuration
-type GenericExporter struct {
-	exporter   traceAndViewExporter
-	sampleRate float64
-	flushFn    FlushFn
-}
-
-func (g *GenericExporter) InitExportOnce() error {
-	var err error
-	initExporterOnce.Do(func() {
-		err = g.initExporter()
-	})
-	return err
-}
-
-func (g *GenericExporter) Flush() {
-	g.flushFn()
-}
-
-func (g *GenericExporter) initExporter() error {
-	if g.exporter == nil {
-		return nil
-	}
-	trace.ApplyConfig(trace.Config{
-		DefaultSampler: trace.ProbabilitySampler(g.sampleRate),
-	})
-	trace.RegisterExporter(g.exporter)
-	view.RegisterExporter(g.exporter)
-
-	// Record the various HTTP view to collect metrics.
-	httpViews := append(ochttp.DefaultServerViews, ochttp.DefaultClientViews...)
-	if err := view.Register(httpViews...); err != nil {
-		return fmt.Errorf("failed to register http views for observability exporter: %v", err)
-	}
-	// Register the various gRPC views to collect metrics.
-	gRPCViews := append(ocgrpc.DefaultServerViews, ocgrpc.DefaultClientViews...)
-	if err := view.Register(gRPCViews...); err != nil {
-		return fmt.Errorf("failed to register grpc views for observability exporter: %v", err)
-	}
-	return nil
-}
diff --git a/pkg/observability/noop.go b/pkg/observability/noop.go
index c66633d..93a4e63 100644
--- a/pkg/observability/noop.go
+++ b/pkg/observability/noop.go
@@ -15,14 +15,22 @@
 // Package observability sets up and configures observability tools.
 package observability
 
+import "context"
+
 // Compile-time check to verify implements interface.
-var _ Exporter = (*NoopExporter)(nil)
+var _ Exporter = (*noopExporter)(nil)
+
+// noopExporter is an observability exporter that does nothing.
+type noopExporter struct{}
 
-// NoopExporter is an observability exporter that does nothing.
-type NoopExporter struct{}
+func NewNoop(_ context.Context) (Exporter, error) {
+	return &noopExporter{}, nil
+}
 
-func (g *NoopExporter) InitExportOnce() error {
+func (g *noopExporter) StartExporter() error {
 	return nil
 }
 
-func (g *NoopExporter) Flush() {}
+func (g *noopExporter) Close() error {
+	return nil
+}
diff --git a/pkg/observability/observability.go b/pkg/observability/observability.go
index ae1d746..425d907 100644
--- a/pkg/observability/observability.go
+++ b/pkg/observability/observability.go
@@ -18,64 +18,48 @@ package observability
 import (
 	"context"
 	"fmt"
-	"time"
+	"io"
 
-	"contrib.go.opencensus.io/exporter/ocagent"
-	"contrib.go.opencensus.io/exporter/stackdriver"
-	"github.com/google/exposure-notifications-server/internal/logging"
+	"go.opencensus.io/plugin/ocgrpc"
+	"go.opencensus.io/plugin/ochttp"
+	"go.opencensus.io/stats/view"
 )
 
 // Exporter defines the minimum shared functionality for an observability exporter
 // used by this application.
 type Exporter interface {
-	InitExportOnce() error
-	Flush()
+	io.Closer
+	StartExporter() error
 }
 
 // NewFromEnv returns the observability exporter given the provided configuration, or an error
 // if it failed to be created.
 func NewFromEnv(ctx context.Context, config *Config) (Exporter, error) {
 	switch config.ExporterType {
-	default:
-		return &NoopExporter{}, nil
-
 	case ExporterNoop:
-		return &NoopExporter{}, nil
-
+		return NewNoop(ctx)
 	case ExporterStackdriver:
-		if config.StackdriverConfig.ProjectID == "" {
-			return nil, fmt.Errorf("configuration PROJECT_ID is required to use the Stackdriver observability exporter")
-		}
-		logger := logging.FromContext(ctx).Named("stackdriver")
-
-		monitoredResource := NewStackdriverMonitoredResoruce(&config.StackdriverConfig)
-
-		sde, err := stackdriver.NewExporter(stackdriver.Options{
-			ProjectID:         config.StackdriverConfig.ProjectID,
-			ReportingInterval: time.Minute, // stackdriver export interval minimum
-			MonitoredResource: monitoredResource,
-			OnError: func(err error) {
-				logger.Errorf("stackdriver export error: %v", err)
-			},
-		})
-		if err != nil {
-			return nil, fmt.Errorf("failed to create Stackdriver observability exporter: %v", err)
-		}
-		return &GenericExporter{sde, config.TraceProbabilitySampleRate, func() { sde.Flush() }}, nil
-
+		return NewStackdriver(ctx, config.Stackdriver)
 	case ExporterOCAgent, ExporterPrometheus:
-		var opts []ocagent.ExporterOption
-		if config.OCAgentConfig.Insecure {
-			opts = append(opts, ocagent.WithInsecure())
-		}
-		if config.OCAgentConfig.Endpoint != "" {
-			opts = append(opts, ocagent.WithAddress(config.OCAgentConfig.Endpoint))
-		}
+		return NewOpenCensus(ctx, config.OpenCensus)
+	default:
+		return nil, fmt.Errorf("unknown observability exporter type %v", config.ExporterType)
+	}
+}
 
-		oce, err := ocagent.NewExporter(opts...)
-		if err != nil {
-			return nil, fmt.Errorf("failed to create OpenCensus observability exporter: %v", err)
-		}
-		return &GenericExporter{oce, config.TraceProbabilitySampleRate, func() { oce.Flush() }}, nil
+// registerViews registers the necessary tracing views.
+func registerViews() error {
+	// Record the various HTTP view to collect metrics.
+	httpViews := append(ochttp.DefaultServerViews, ochttp.DefaultClientViews...)
+	if err := view.Register(httpViews...); err != nil {
+		return fmt.Errorf("failed to register http views: %w", err)
 	}
+
+	// Register the various gRPC views to collect metrics.
+	gRPCViews := append(ocgrpc.DefaultServerViews, ocgrpc.DefaultClientViews...)
+	if err := view.Register(gRPCViews...); err != nil {
+		return fmt.Errorf("failed to register grpc views: %w", err)
+	}
+
+	return nil
 }
diff --git a/pkg/observability/opencensus.go b/pkg/observability/opencensus.go
new file mode 100644
index 0000000..24ce19c
--- /dev/null
+++ b/pkg/observability/opencensus.go
@@ -0,0 +1,76 @@
+// Copyright 2020 Google LLC
+//
+// 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 observability
+
+import (
+	"context"
+	"fmt"
+
+	"contrib.go.opencensus.io/exporter/ocagent"
+	"go.opencensus.io/stats/view"
+	"go.opencensus.io/trace"
+)
+
+var _ Exporter = (*opencensusExporter)(nil)
+
+type opencensusExporter struct {
+	exporter *ocagent.Exporter
+	config   *OpenCensusConfig
+}
+
+// NewOpenCensus creates a new metrics and trace exporter for OpenCensus.
+func NewOpenCensus(ctx context.Context, config *OpenCensusConfig) (Exporter, error) {
+	var opts []ocagent.ExporterOption
+	if config.Insecure {
+		opts = append(opts, ocagent.WithInsecure())
+	}
+	if config.Endpoint != "" {
+		opts = append(opts, ocagent.WithAddress(config.Endpoint))
+	}
+
+	oc, err := ocagent.NewExporter(opts...)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create opencensus exporter: %w", err)
+	}
+	return &opencensusExporter{oc, config}, nil
+}
+
+// StartExporter starts the exporter.
+func (e *opencensusExporter) StartExporter() error {
+	trace.ApplyConfig(trace.Config{
+		DefaultSampler: trace.ProbabilitySampler(e.config.SampleRate),
+	})
+	trace.RegisterExporter(e.exporter)
+	view.RegisterExporter(e.exporter)
+
+	if err := registerViews(); err != nil {
+		return fmt.Errorf("failed to start opencensus exporter: %w", err)
+	}
+
+	return nil
+}
+
+// Close halts the exporter.
+func (e *opencensusExporter) Close() error {
+	if err := e.exporter.Stop(); err != nil {
+		return fmt.Errorf("failed to stop exporter: %w", err)
+	}
+
+	// Unregister the exporter
+	trace.UnregisterExporter(e.exporter)
+	view.UnregisterExporter(e.exporter)
+
+	return nil
+}
diff --git a/pkg/observability/resource.go b/pkg/observability/resource.go
index 46e297a..c195c78 100644
--- a/pkg/observability/resource.go
+++ b/pkg/observability/resource.go
@@ -22,9 +22,9 @@ import (
 	"github.com/google/uuid"
 )
 
-var _ monitoredresource.Interface = (*StackdriverMonitoredResoruce)(nil)
+var _ monitoredresource.Interface = (*stackdriverMonitoredResource)(nil)
 
-type StackdriverMonitoredResoruce struct {
+type stackdriverMonitoredResource struct {
 	resource string
 	labels   map[string]string
 }
@@ -65,12 +65,12 @@ func NewStackdriverMonitoredResoruce(c *StackdriverConfig) monitoredresource.Int
 	}
 	labels["namespace"] = c.Namespace
 
-	return &StackdriverMonitoredResoruce{
+	return &stackdriverMonitoredResource{
 		resource: resource,
 		labels:   labels,
 	}
 }
 
-func (s *StackdriverMonitoredResoruce) MonitoredResource() (string, map[string]string) {
+func (s *stackdriverMonitoredResource) MonitoredResource() (string, map[string]string) {
 	return s.resource, s.labels
 }
diff --git a/pkg/observability/stackdriver.go b/pkg/observability/stackdriver.go
new file mode 100644
index 0000000..fa9e066
--- /dev/null
+++ b/pkg/observability/stackdriver.go
@@ -0,0 +1,90 @@
+// Copyright 2020 Google LLC
+//
+// 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 observability sets up and configures observability tools.
+package observability
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"contrib.go.opencensus.io/exporter/stackdriver"
+	"github.com/google/exposure-notifications-server/internal/logging"
+	"go.opencensus.io/stats/view"
+	"go.opencensus.io/trace"
+)
+
+var _ Exporter = (*stackdriverExporter)(nil)
+
+type stackdriverExporter struct {
+	exporter *stackdriver.Exporter
+	config   *StackdriverConfig
+}
+
+// NewStackdriver creates a new metrics and trace exporter for Stackdriver.
+func NewStackdriver(ctx context.Context, config *StackdriverConfig) (Exporter, error) {
+	logger := logging.FromContext(ctx).Named("stackdriver")
+
+	projectID := config.ProjectID
+	if projectID == "" {
+		return nil, fmt.Errorf("missing PROJECT_ID in Stackdriver exporter")
+	}
+
+	monitoredResource := NewStackdriverMonitoredResoruce(config)
+
+	exporter, err := stackdriver.NewExporter(stackdriver.Options{
+		ProjectID:         projectID,
+		ReportingInterval: time.Minute, // stackdriver export interval minimum
+		MonitoredResource: monitoredResource,
+		OnError: func(err error) {
+			logger.Errorw("failed to export metric", "error", err)
+		},
+	})
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to create Stackdriver exporter: %w", err)
+	}
+	return &stackdriverExporter{exporter, config}, nil
+}
+
+// StartExporter starts the exporter.
+func (e *stackdriverExporter) StartExporter() error {
+	if err := e.exporter.StartMetricsExporter(); err != nil {
+		return fmt.Errorf("failed to start stackdriver exporter: %w", err)
+	}
+
+	trace.ApplyConfig(trace.Config{
+		DefaultSampler: trace.ProbabilitySampler(e.config.SampleRate),
+	})
+	trace.RegisterExporter(e.exporter)
+	view.RegisterExporter(e.exporter)
+
+	return nil
+}
+
+// Close halts the exporter.
+func (e *stackdriverExporter) Close() error {
+	// Flush any existing metrics
+	e.exporter.Flush()
+
+	// Stop the exporter
+	e.exporter.StopMetricsExporter()
+
+	// Unregister the exporter
+	trace.UnregisterExporter(e.exporter)
+	view.UnregisterExporter(e.exporter)
+
+	return nil
+}


Test patch:
diff --git a/internal/keyrotation/rotate_test.go b/internal/keyrotation/rotate_test.go
index f1587cc..41b4147 100644
--- a/internal/keyrotation/rotate_test.go
+++ b/internal/keyrotation/rotate_test.go
@@ -37,7 +37,9 @@ func TestRotateKeys(t *testing.T) {
 
 	kms, _ := keys.NewInMemory(context.Background())
 	keyID := "testKeyID"
-	kms.AddEncryptionKey(keyID)
+	if _, err := kms.CreateEncryptionKey(keyID); err != nil {
+		t.Fatal(err)
+	}
 	config := &Config{
 		RevisionToken:      revision.Config{KeyID: keyID},
 		DeleteOldKeyPeriod: 14 * 24 * time.Hour, // two weeks
diff --git a/internal/keyrotation/server_test.go b/internal/keyrotation/server_test.go
index 7511d18..3a36d5f 100644
--- a/internal/keyrotation/server_test.go
+++ b/internal/keyrotation/server_test.go
@@ -62,7 +62,9 @@ func TestNewRotationHandler(t *testing.T) {
 			t.Parallel()
 
 			keyID := "test" + t.Name()
-			kms.AddEncryptionKey(keyID)
+			if _, err := kms.CreateEncryptionKey(keyID); err != nil {
+				t.Fatal(err)
+			}
 			config := &Config{
 				RevisionToken: revision.Config{KeyID: keyID},
 			}
diff --git a/internal/publish/publish_test.go b/internal/publish/publish_test.go
index 6445efa..644b764 100644
--- a/internal/publish/publish_test.go
+++ b/internal/publish/publish_test.go
@@ -371,7 +371,9 @@ func TestPublishWithBypass(t *testing.T) {
 			t.Fatalf("can't make kms: %v", err)
 		}
 		keyID := "rev"
-		kms.AddEncryptionKey(keyID)
+		if _, err := kms.CreateEncryptionKey(keyID); err != nil {
+			t.Fatal(err)
+		}
 		tokenAAD := make([]byte, 16)
 		if _, err := rand.Read(tokenAAD); err != nil {
 			t.Fatalf("not enough entropy: %v", err)
@@ -682,7 +684,9 @@ func TestKeyRevision(t *testing.T) {
 		t.Fatalf("can't make kms: %v", err)
 	}
 	keyID := "rev"
-	kms.AddEncryptionKey(keyID)
+	if _, err := kms.CreateEncryptionKey(keyID); err != nil {
+		t.Fatal(err)
+	}
 	tokenAAD := make([]byte, 16)
 	if _, err := rand.Read(tokenAAD); err != nil {
 		t.Fatalf("not enough entropy: %v", err)
diff --git a/internal/revision/database/revision_test.go b/internal/revision/database/revision_test.go
index a020e94..2e1d8a2 100644
--- a/internal/revision/database/revision_test.go
+++ b/internal/revision/database/revision_test.go
@@ -40,7 +40,7 @@ func TestRevisionKey(t *testing.T) {
 		t.Fatalf("unable to cerate in memory KMS")
 	}
 	keyID := "funkey"
-	if err := kms.AddEncryptionKey(keyID); err != nil {
+	if _, err := kms.CreateEncryptionKey(keyID); err != nil {
 		t.Fatalf("unable to generate key: %v", err)
 	}
 
@@ -76,7 +76,7 @@ func TestMultipleRevisionKeys(t *testing.T) {
 		t.Fatalf("unable to cerate in memory KMS")
 	}
 	keyID := "funkey"
-	if err := kms.AddEncryptionKey(keyID); err != nil {
+	if _, err := kms.CreateEncryptionKey(keyID); err != nil {
 		t.Fatalf("unable to generate key: %v", err)
 	}
 
diff --git a/internal/revision/revision_test.go b/internal/revision/revision_test.go
index 379dc66..91817eb 100644
--- a/internal/revision/revision_test.go
+++ b/internal/revision/revision_test.go
@@ -128,7 +128,7 @@ func TestEncryptDecrypt(t *testing.T) {
 		t.Fatalf("unable to cerate in memory KMS")
 	}
 	keyID := "skeleton"
-	if err := kms.AddEncryptionKey(keyID); err != nil {
+	if _, err := kms.CreateEncryptionKey(keyID); err != nil {
 		t.Fatalf("unable to generate key: %v", err)
 	}
 	key := make([]byte, 32)
diff --git a/internal/setup/setup_test.go b/internal/setup/setup_test.go
index b2e77b4..809bfcf 100644
--- a/internal/setup/setup_test.go
+++ b/internal/setup/setup_test.go
@@ -231,9 +231,10 @@ func TestSetupWith(t *testing.T) {
 		if oe == nil {
 			t.Errorf("expected observability exporter to exist")
 		}
-
-		if _, ok := oe.(*observability.NoopExporter); !ok {
-			t.Errorf("expected %T to be GenericExporter", oe)
-		}
+		defer func() {
+			if err := oe.Close(); err != nil {
+				t.Fatal(err)
+			}
+		}()
 	})
 }