Repository: google/exposure-notifications-server
Pull request #765 (2020-07-23 02:02:23)

Managed revision keys
* add "in memory" implementation of the keys.KeyManager interface
  * used in these tests, but will be useful in other places as well
* create new RevisionKeys database table
* revision keys are generated AES keys that are encrypted using the configured KMS before storing in the databse
* revision tokens are encrypted/decrypted using envelope encryption based on these keys
* modify revision token to use a token manager
  * allowed keys are cached
  * the keys are only unwrapped when necessary
  * key rotation is enabled by minting a new revision key
  * a key can be destroyed when it should no longer be accepeted

Part of #734 - just a bit more to go after this.

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/765#" title="Author self-approved">mikehelmick</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)~~ [mikehelmick]

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":[]} -->
google-oss-robot: @mikehelmick: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-en-server-release-unit | fceaa011bfb195f1a683e366e5e45312da3cf71e | [link](https://oss-prow.knative.dev/view/gs/oss-prow/pr-logs/pull/google_exposure-notifications-server/765/pull-en-server-release-unit/1286119571492704256) | `/test pull-en-server-release-unit`

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
google-oss-robot: @mikehelmick: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-en-server-release-unit | f1c9fb2c747e5cdd161137700b333d44e2d5b998 | [link](https://oss-prow.knative.dev/view/gs/oss-prow/pr-logs/pull/google_exposure-notifications-server/765/pull-en-server-release-unit/1286121710600327168) | `/test pull-en-server-release-unit`

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
google-oss-robot: @mikehelmick: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-en-server-release-unit | d07d40a09d3b7c2d456e9939000cba92958fcf2b | [link](https://oss-prow.knative.dev/view/gs/oss-prow/pr-logs/pull/google_exposure-notifications-server/765/pull-en-server-release-unit/1286123472208007168) | `/test pull-en-server-release-unit`

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
google-oss-robot: @mikehelmick: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-en-server-release-unit | 158a042c0066c8d81784541b8c6b54a1ab57ea92 | [link](https://oss-prow.knative.dev/view/gs/oss-prow/pr-logs/pull/google_exposure-notifications-server/765/pull-en-server-release-unit/1286344553950875648) | `/test pull-en-server-release-unit`

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
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/765#" title="Author self-approved">mikehelmick</a>*, *<a href="https://github.com/google/exposure-notifications-server/pull/765#pullrequestreview-454328703" title="LGTM">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)~~ [mikehelmick,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":[]} -->
google-oss-robot: New changes are detected. LGTM label has been removed.
sethvargo: /lgtm

Patch:
diff --git a/internal/publish/config.go b/internal/publish/config.go
index e91d87f..f7c6565 100644
--- a/internal/publish/config.go
+++ b/internal/publish/config.go
@@ -51,6 +51,11 @@ type Config struct {
 	MaxMagnitudeSymptomOnsetDays uint          `env:"MAX_SYMPTOM_ONSET_DAYS, default=21"`
 	CreatedAtTruncateWindow      time.Duration `env:"TRUNCATE_WINDOW, default=1h"`
 
+	// Crypto key to use for wrapping/unwrapping the revision token cipher blocks.
+	RevisionTokenKeyID         string `env:"REVISION_TOKEN_KEY_ID"`
+	RevisionKeyWrapperAAD      string `env:"REVISION_KEY_WRAPPER_AAD"` // must be base64 encoded, may come from secret://
+	BypassRevisionCertificates bool   `env:"BYPASS_REVISION_CERTIFICATES, default=false"`
+
 	// Flags for local development and testing. This will cause still valid keys
 	// to not be embargoed.
 	// Normally "still valid" keys can be accepted, but are embargoed.
diff --git a/internal/revision/database/revision.go b/internal/revision/database/revision.go
new file mode 100644
index 0000000..994c0c7
--- /dev/null
+++ b/internal/revision/database/revision.go
@@ -0,0 +1,331 @@
+// 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 database contains the management of interactions with the database
+// for createion and storage of the wrapped keys that encrypet revision certificates.
+//
+// RevisionKey data is stored in the revisionkeys table.
+// * The most recently created 'allowed' key is considered to be the effective key.
+//   The effective key is used to encrypt outgoing revision tokens.
+// * Any still 'allowed' key can be used to decrypt incoming revision tokens.
+//
+// This package also supports the creation of new keys with a locally generated
+// AES key that is encrypted using the provided KMS and stored in the database
+// in it's encrypted form.
+package database
+
+import (
+	"context"
+	"crypto/rand"
+	"fmt"
+	"io"
+	"strconv"
+	"time"
+
+	"github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/logging"
+	"github.com/google/exposure-notifications-server/pkg/keys"
+	"github.com/jackc/pgx/v4"
+)
+
+// RevisionKey represents an individual revision key.
+type RevisionKey struct {
+	KeyID         int64
+	AAD           []byte // AAD for the wrapping/unwrapping of the cipher block.
+	WrappedCipher []byte
+	CreatedAt     time.Time
+	Allowed       bool
+
+	// The unwrapped cipher.
+	DEK []byte
+}
+
+// KeyIDString returns the keyID as a string that can be used in the encoded revision tokens.
+func (r *RevisionKey) KeyIDString() string {
+	return strconv.FormatInt(r.KeyID, 10)
+}
+
+// KMSConfig represents the configuration of the RevisionDB in terms how how it should utilize the
+// given KeyManager for wrapping/unwrapping keys.
+type KMSConfig struct {
+	WrapperKeyID string
+	WrapperAAD   []byte
+	KeyManager   keys.KeyManager
+}
+
+// Copy creates a deep copy of the KMSConfig
+// This is done to ensure that we have a clean copy of the ADD slice that can't be changed externally.
+func (k *KMSConfig) Copy() *KMSConfig {
+	c := KMSConfig{
+		WrapperKeyID: k.WrapperKeyID,
+		WrapperAAD:   make([]byte, 0, len(k.WrapperAAD)),
+		KeyManager:   k.KeyManager,
+	}
+	copy(c.WrapperAAD, k.WrapperAAD)
+	return &c
+}
+
+// RevisionDB wraps a database connection and provides functions for interacting with revision keys.
+type RevisionDB struct {
+	db     *database.DB
+	config *KMSConfig
+}
+
+// New creates a new `RevisionDB`
+func New(db *database.DB, c *KMSConfig) (*RevisionDB, error) {
+	if c.WrapperKeyID == "" {
+		return nil, fmt.Errorf("no KMS key ID passed in to revision.New")
+	}
+	if len(c.WrapperAAD) == 0 {
+		return nil, fmt.Errorf("no AAD provided")
+	}
+	if c.KeyManager == nil {
+		return nil, fmt.Errorf("no KeyManager provided")
+	}
+	return &RevisionDB{
+		db:     db,
+		config: c.Copy(),
+	}, nil
+}
+
+// DestroyKey zeros out the wrapped key and marks the key as allowed=false.
+func (rdb *RevisionDB) DestroyKey(ctx context.Context, keyID int64) error {
+	logger := logging.FromContext(ctx)
+	logger.Warnf("destroying key material for revision key ID %v", keyID)
+
+	if err := rdb.db.InTx(ctx, pgx.Serializable, func(tx pgx.Tx) error {
+		result, err := tx.Exec(ctx, `
+			UPDATE
+				RevisionKeys
+			SET
+				wrapped_cipher = $1, aad = $2, allowed = $3
+			WHERE
+				kid = $4
+		`, nil, nil, false, keyID)
+		if err != nil {
+			return fmt.Errorf("updating revisionkey: %w", err)
+		}
+		if result.RowsAffected() != 1 {
+			return fmt.Errorf("revision key was not updated as expected")
+		}
+		return nil
+	}); err != nil {
+		logger.Errorf("failed to destroy revision kid: %v: %v", keyID, err)
+	}
+	return nil
+}
+
+// GetAllowedRevisionKeyIDs returns just the IDs of still allowed keys.
+// Once the keys have been unwrapped, there is no reason to continue to unwrap them.
+//
+// The first return value is the ID of the currently effective key (most recently created, still active)
+// The second return value is a map of the currently allowed keys for decryption
+func (rdb *RevisionDB) GetAllowedRevisionKeyIDs(ctx context.Context) (int64, map[int64]struct{}, error) {
+	var effectiveID int64
+	keys := make(map[int64]struct{})
+	if err := rdb.db.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
+		// Ordering by created_at DESC puts the "effective key" first.
+		rows, err := tx.Query(ctx, `
+			SELECT
+				kid
+			FROM
+				revisionkeys
+			WHERE
+				allowed=$1
+			ORDER BY created_at DESC`, true)
+		if err != nil {
+			return nil
+		}
+		defer rows.Close()
+
+		for rows.Next() {
+			if err := rows.Err(); err != nil {
+				return fmt.Errorf("failed to iterate: %w", err)
+			}
+			var id int64
+			if err := rows.Scan(&id); err != nil {
+				return err
+			}
+			keys[id] = struct{}{}
+			// The first key we see (due to sort) is the effective key.
+			if effectiveID == 0 {
+				effectiveID = id
+			}
+		}
+
+		return nil
+	}); err != nil {
+		return 0, nil, fmt.Errorf("unable to read keys: %w1", err)
+	}
+	return effectiveID, keys, nil
+}
+
+// GetAllowedRevisionKeys returns all of the curently allowed revision keys.
+// This method will unwrap all of the keys so that they can be used to create and verify
+// revision tokens.
+//
+// The first return value is the ID of the effective RevisionKey
+// The second is a slice of all currently allowed RevisionKeys for decryption purposes. The returned
+// revision keys will be sorted in reverse time order by creation time.
+func (rdb *RevisionDB) GetAllowedRevisionKeys(ctx context.Context) (int64, []*RevisionKey, error) {
+	logger := logging.FromContext(ctx)
+	var effectiveID int64
+	keys := make([]*RevisionKey, 0)
+	if err := rdb.db.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
+		// Need to sort by created_at DESC so the first key encountered is the "effective" key.
+		rows, err := tx.Query(ctx, `
+			SELECT
+				kid, aad, wrapped_cipher, created_at, allowed
+			FROM
+				revisionkeys
+			WHERE
+				allowed=$1
+			ORDER BY created_at DESC`, true)
+		if err != nil {
+			return nil
+		}
+		defer rows.Close()
+
+		for rows.Next() {
+			if err := rows.Err(); err != nil {
+				return fmt.Errorf("failed to iterate: %w", err)
+			}
+			var r RevisionKey
+			if err := rows.Scan(&r.KeyID, &r.AAD, &r.WrappedCipher, &r.CreatedAt, &r.Allowed); err != nil {
+				return err
+			}
+			keys = append(keys, &r)
+			// The effective KEY is the first one due to ORDER BY clause.
+			if effectiveID == 0 {
+				effectiveID = r.KeyID
+			}
+		}
+
+		return nil
+	}); err != nil {
+		return 0, nil, fmt.Errorf("unable to read keys: %w1", err)
+	}
+
+	unwrappedKeys := make([]*RevisionKey, 0, len(keys))
+	// Attempt to unwrap all of the keys
+	for _, wk := range keys {
+		unwrapped, err := rdb.decrypt(ctx, wk.WrappedCipher)
+		if err != nil {
+			logger.Errorf("still allowed revision key that can't be unwrapped: kid: %v error: %v", wk.KeyID, err)
+			return 0, nil, fmt.Errorf("unable to unwrap revision key: %w", err)
+		}
+		wk.DEK = unwrapped
+		unwrappedKeys = append(unwrappedKeys, wk)
+	}
+
+	// Go through an unwrap all of the keys.
+	return effectiveID, unwrappedKeys, nil
+}
+
+// GetEffectiveRevisionKey returnes the revision key to use when encrypting revision tokens.
+// This is consided the most recently created key that is still "allowed"
+func (rdb *RevisionDB) GetEffectiveRevisionKey(ctx context.Context) (*RevisionKey, error) {
+	var revKey *RevisionKey
+	if err := rdb.db.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
+		rows, err := tx.Query(ctx, `
+			SELECT
+				kid, aad, wrapped_cipher, created_at, allowed
+			FROM
+				revisionkeys
+			WHERE
+				allowed=$1
+			ORDER BY created_at DESC
+			LIMIT 1`, true)
+		if err != nil {
+			return err
+		}
+		defer rows.Close()
+
+		if rows.Next() {
+			if err := rows.Err(); err != nil {
+				return fmt.Errorf("failed to iterate: %w", err)
+			}
+			var r RevisionKey
+			if err := rows.Scan(&r.KeyID, &r.AAD, &r.WrappedCipher, &r.CreatedAt, &r.Allowed); err != nil {
+				return err
+			}
+			revKey = &r
+		}
+		return nil
+	}); err != nil {
+		return nil, fmt.Errorf("unable to get revision key: %w", err)
+	}
+	if revKey == nil {
+		return nil, fmt.Errorf("no effective key found")
+	}
+
+	// Unwrap the DEK w/ the KeyManager.
+	unwrapped, err := rdb.decrypt(ctx, revKey.WrappedCipher)
+	if err != nil {
+		return nil, fmt.Errorf("unable to unwrap key: %w", err)
+	}
+	revKey.DEK = unwrapped
+
+	return revKey, nil
+}
+
+func (rdb *RevisionDB) decrypt(ctx context.Context, ciphertext []byte) ([]byte, error) {
+	return rdb.config.KeyManager.Decrypt(ctx, rdb.config.WrapperKeyID, ciphertext, rdb.config.WrapperAAD)
+}
+
+// CreateRevisionKey generates a new AES key and wraps it
+func (rdb *RevisionDB) CreateRevisionKey(ctx context.Context) (*RevisionKey, error) {
+	key := make([]byte, 32)
+	if _, err := io.ReadFull(rand.Reader, key); err != nil {
+		return nil, fmt.Errorf("unable to generate AES key: %w", err)
+	}
+	aad := make([]byte, 16)
+	if _, err := io.ReadFull(rand.Reader, aad); err != nil {
+		return nil, fmt.Errorf("unable to generate random data: %w", err)
+	}
+
+	// Wrap the key using the configured KMS.
+	wrapped, err := rdb.config.KeyManager.Encrypt(ctx, rdb.config.WrapperKeyID, key, rdb.config.WrapperAAD)
+	if err != nil {
+		return nil, fmt.Errorf("failed to wrap key: %w", err)
+	}
+
+	// Start building the RevisionKey
+	revKey := RevisionKey{
+		WrappedCipher: wrapped,
+		AAD:           aad,
+		CreatedAt:     time.Now().UTC(),
+		Allowed:       true,
+		DEK:           key,
+	}
+
+	if err := rdb.db.InTx(ctx, pgx.ReadCommitted, func(tx pgx.Tx) error {
+		row := tx.QueryRow(ctx, `
+			INSERT INTO
+				RevisionKeys
+				(aad, wrapped_cipher, created_at, allowed)
+			VALUES
+				($1, $2, $3, $4)
+			RETURNING kid`,
+			revKey.AAD, wrapped, revKey.CreatedAt, true)
+		if err := row.Scan(&revKey.KeyID); err != nil {
+			return fmt.Errorf("fetching kid: %w", err)
+		}
+		return nil
+	}); err != nil {
+		return nil, fmt.Errorf("unable to persist revision key: %w", err)
+	}
+
+	return &revKey, nil
+}
diff --git a/internal/revision/revision.go b/internal/revision/revision.go
index 4bc7782..ab20467 100644
--- a/internal/revision/revision.go
+++ b/internal/revision/revision.go
@@ -17,18 +17,147 @@
 package revision
 
 import (
+	"context"
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/rand"
 	"fmt"
 	"io"
 	"sort"
+	"strconv"
+	"sync"
+	"time"
 
+	"github.com/google/exposure-notifications-server/internal/logging"
 	"github.com/google/exposure-notifications-server/internal/pb"
 	"github.com/google/exposure-notifications-server/internal/publish/model"
+	"github.com/google/exposure-notifications-server/internal/revision/database"
 	"google.golang.org/protobuf/proto"
 )
 
+// TokenManager is responsible for creating and unlocking revision tokens.
+type TokenManager struct {
+	db *database.RevisionDB
+
+	// All encrypt/decrypt operations are done under read lock.
+	// Cache refresh esclates to a write lock.
+	mu sync.RWMutex
+
+	// A store of the currently allowed revision keys for decryption purposes.
+	allowed map[int64]*database.RevisionKey
+	// A poitners to the currently active key for encryption purposes.
+	effective *database.RevisionKey
+
+	// The allowed/effective keys are cached to avoid excessive decrypt calls to the KMS system to unwrap
+	// the individual revision keys.
+	// A cache refresh is initially a shallow refresh, if the IDs of allowed/effective keys haven't changed,
+	// we don't re-unwrap the keys. If there are any changes to the IDs, all of the daata is reloaded and
+	// the keys are unwrapped.
+	cacheDuration     time.Duration
+	cacheRefreshAfter time.Time
+}
+
+// New creates a new TokenManager that uses a database handle to manage a cache
+// of allowed revision keys.
+func New(ctx context.Context, db *database.RevisionDB, cacheDuration time.Duration) (*TokenManager, error) {
+	if cacheDuration > 60*time.Minute {
+		return nil, fmt.Errorf("cache duration must be <= 60 minutes, got: %v", cacheDuration)
+	}
+	now := time.Now()
+	tm := &TokenManager{
+		db:                db,
+		allowed:           make(map[int64]*database.RevisionKey),
+		cacheDuration:     cacheDuration,
+		cacheRefreshAfter: now.Add(-2 * cacheDuration),
+	}
+	if err := tm.maybeRefreshCache(ctx); err != nil {
+		return nil, err
+	}
+	return tm, nil
+}
+
+func (tm *TokenManager) expired() bool {
+	tm.mu.RLock()
+	defer tm.mu.RUnlock()
+	return time.Now().After(tm.cacheRefreshAfter)
+}
+
+func (tm *TokenManager) maybeRefreshCache(ctx context.Context) error {
+	if !tm.expired() {
+		return nil
+	}
+	// Escalate to a write lock and refresh the cache.
+	tm.mu.Lock()
+	defer tm.mu.Unlock()
+
+	reload, err := tm.isReloadNeeded(ctx)
+	if err != nil {
+		return fmt.Errorf("unable to read revsion keys: %w", err)
+	}
+	if !reload {
+		return nil
+	}
+
+	// At this point reload is needed. Trash the information currently stored
+	// so that if something failes on refresh, the cache has been invalidated.
+	tm.allowed = make(map[int64]*database.RevisionKey)
+	tm.effective = nil
+
+	// Go back and reload and unwrap the effective keys.
+	logger := logging.FromContext(ctx)
+	logger.Info("reloading revision key cache")
+
+	effectiveID, allowed, err := tm.db.GetAllowedRevisionKeys(ctx)
+	if err != nil {
+		logger.Errorf("unable to read revision keys: %v", err)
+		return fmt.Errorf("reading revision key cache: %w", err)
+	}
+
+	if len(allowed) == 0 {
+		return fmt.Errorf("no revision keys exist")
+	}
+
+	for _, rk := range allowed {
+		tm.allowed[rk.KeyID] = rk
+	}
+	tm.effective = tm.allowed[effectiveID]
+	// we did it! mark the next refresh time.
+	tm.cacheRefreshAfter = time.Now().Add(tm.cacheDuration)
+	return nil
+}
+
+// Determine if we actually need to reload and unwrap keys.
+// Must be called under write lock.
+func (tm *TokenManager) isReloadNeeded(ctx context.Context) (bool, error) {
+	// If there is no effective key, reload.
+	if tm.effective == nil {
+		return true, nil
+	}
+
+	effectiveID, allowedIDs, err := tm.db.GetAllowedRevisionKeyIDs(ctx)
+	if err != nil {
+		return true, err
+	}
+	if effectiveID != tm.effective.KeyID {
+		return true, nil
+	}
+
+	for k := range allowedIDs {
+		if rk := tm.allowed[k]; rk == nil {
+			// Found an allowed key that we haven't seen yet.
+			return true, nil
+		}
+	}
+	// remove any keys that are no longer allowed from the cache.
+	for k := range tm.allowed {
+		if _, ok := allowedIDs[k]; !ok {
+			delete(tm.allowed, k)
+		}
+	}
+
+	return false, nil
+}
+
 func buildTokenBufer(eKeys []*model.Exposure) *pb.RevisionTokenData {
 	// sort the keys.
 	sort.Slice(eKeys, func(i, j int) bool {
@@ -52,11 +181,25 @@ func buildTokenBufer(eKeys []*model.Exposure) *pb.RevisionTokenData {
 
 // MakeRevisionToken turns the TEK data from a given publish request
 // into an encrypted protocol buffer revision token.
-func MakeRevisionToken(eKeys []*model.Exposure, aad []byte, kid string, cryptoKey []byte) ([]byte, error) {
+// This is using envelope encryption, based on the currently active revision key.
+func (tm *TokenManager) MakeRevisionToken(ctx context.Context, eKeys []*model.Exposure, aad []byte) ([]byte, error) {
 	if len(eKeys) == 0 {
 		return nil, fmt.Errorf("no keys to build token for")
 	}
 
+	if err := tm.maybeRefreshCache(ctx); err != nil {
+		return nil, err
+	}
+	// Capture DEK and KID in read lock, but don't do encryption with the lock
+	var dek []byte
+	var kid string
+	{
+		tm.mu.RLock()
+		dek = tm.effective.DEK
+		kid = tm.effective.KeyIDString()
+		tm.mu.RUnlock()
+	}
+
 	tokenData := buildTokenBufer(eKeys)
 	plaintext, err := proto.Marshal(tokenData)
 	if err != nil {
@@ -64,7 +207,7 @@ func MakeRevisionToken(eKeys []*model.Exposure, aad []byte, kid string, cryptoKe
 	}
 
 	// encrypt the serialized proto.
-	block, err := aes.NewCipher(cryptoKey)
+	block, err := aes.NewCipher(dek)
 	if err != nil {
 		return nil, fmt.Errorf("bad cipher block: %w", err)
 	}
@@ -93,16 +236,33 @@ func MakeRevisionToken(eKeys []*model.Exposure, aad []byte, kid string, cryptoKe
 
 // UnmarshalRevisionToken unmarshals a revision token, decrypts the payload,
 // and returns the TEK data that was contained in the token if valid.
-func UnmarshalRevisionToken(tokenBytes []byte, aad []byte, cryptoKeys map[string][]byte) (*pb.RevisionTokenData, error) {
+//
+// The incoming key ID is used to determine if this token can still be unlocked.
+func (tm *TokenManager) UnmarshalRevisionToken(ctx context.Context, tokenBytes []byte, aad []byte) (*pb.RevisionTokenData, error) {
+	if err := tm.maybeRefreshCache(ctx); err != nil {
+		return nil, err
+	}
+
 	var revisionToken pb.RevisionToken
 	if err := proto.Unmarshal(tokenBytes, &revisionToken); err != nil {
 		return nil, fmt.Errorf("unable to unmarshal proto envelope: %w", err)
 	}
 	data := revisionToken.Data
+	kid, err := strconv.ParseInt(revisionToken.Kid, 10, 64)
+	if err != nil {
+		return nil, fmt.Errorf("invalid key id: %w", err)
+	}
 
-	dek, ok := cryptoKeys[revisionToken.Kid]
-	if !ok {
-		return nil, fmt.Errorf("token has invalid key id: %v", revisionToken.Kid)
+	var dek []byte
+	// Capture the DEK under read lock, but don't hold lock for decryption.
+	{
+		tm.mu.RLock()
+		defer tm.mu.RUnlock()
+		rk, ok := tm.allowed[kid]
+		if !ok {
+			return nil, fmt.Errorf("token has invalid key id: %v", revisionToken.Kid)
+		}
+		dek = rk.DEK
 	}
 
 	// Decrypt the data block.
diff --git a/migrations/000033_AddRevisionTokenKeys.down.sql b/migrations/000033_AddRevisionTokenKeys.down.sql
new file mode 100644
index 0000000..259443b
--- /dev/null
+++ b/migrations/000033_AddRevisionTokenKeys.down.sql
@@ -0,0 +1,19 @@
+-- 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.
+
+BEGIN;
+
+DROP TABLE RevisionKeys;
+
+END;
diff --git a/migrations/000033_AddRevisionTokenKeys.up.sql b/migrations/000033_AddRevisionTokenKeys.up.sql
new file mode 100644
index 0000000..ecab044
--- /dev/null
+++ b/migrations/000033_AddRevisionTokenKeys.up.sql
@@ -0,0 +1,27 @@
+-- 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.
+
+BEGIN;
+
+CREATE TABLE RevisionKeys (
+    kid SERIAL PRIMARY KEY,
+    aad bytea,
+    wrapped_cipher bytea,
+    created_at TIMESTAMPTZ NOT NULL,
+    allowed bool DEFAULT true
+);
+
+CREATE INDEX revision_keys_start_at ON RevisionKeys USING BTREE(allowed, created_at);
+
+END;
diff --git a/pkg/keys/in_memory.go b/pkg/keys/in_memory.go
new file mode 100644
index 0000000..47498ea
--- /dev/null
+++ b/pkg/keys/in_memory.go
@@ -0,0 +1,162 @@
+// 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 keys
+
+import (
+	"context"
+	"crypto"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"fmt"
+	"io"
+	"sync"
+)
+
+// 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 {
+	mu          sync.RWMutex
+	signingKeys map[string]*ecdsa.PrivateKey
+	cryptoKeys  map[string][]byte
+}
+
+// NewInMemory creates a new, local, in memory KeyManager.
+func NewInMemory(ctx context.Context) (*InMemory, error) {
+	return &InMemory{
+		signingKeys: make(map[string]*ecdsa.PrivateKey),
+		cryptoKeys:  make(map[string][]byte),
+	}, nil
+}
+
+// AddSigningKey generates a new ECDSA P256 Signing Key identified by
+// the provided keyID
+func (k *InMemory) AddSigningKey(keyID string) error {
+	k.mu.Lock()
+	defer k.mu.Unlock()
+
+	if _, ok := k.signingKeys[keyID]; ok {
+		return 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)
+	}
+
+	k.signingKeys[keyID] = pk
+	return nil
+}
+
+// AddEncryptionKey generates a new encryption key identified by
+// the provided keyID.
+func (k *InMemory) AddEncryptionKey(keyID string) error {
+	k.mu.Lock()
+	defer k.mu.Unlock()
+
+	if _, ok := k.cryptoKeys[keyID]; ok {
+		return 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)
+	}
+
+	k.cryptoKeys[keyID] = key
+
+	return nil
+}
+
+func (k *InMemory) NewSigner(ctx context.Context, keyID string) (crypto.Signer, error) {
+	var pk *ecdsa.PrivateKey
+	{
+		k.mu.RLock()
+		defer k.mu.RUnlock()
+
+		if k, ok := k.signingKeys[keyID]; !ok {
+			return nil, fmt.Errorf("key not found")
+		} else {
+			pk = k
+		}
+	}
+
+	return pk, nil
+}
+
+func (k *InMemory) getDEK(keyID string) ([]byte, error) {
+	k.mu.RLock()
+	defer k.mu.RUnlock()
+
+	if dek, ok := k.cryptoKeys[keyID]; ok {
+		return dek, nil
+	}
+	return nil, fmt.Errorf("key not found")
+}
+
+func (k *InMemory) Encrypt(ctx context.Context, keyID string, plaintext []byte, aad []byte) ([]byte, error) {
+	dek, err := k.getDEK(keyID)
+	if err != nil {
+		return nil, err
+	}
+
+	block, err := aes.NewCipher(dek)
+	if err != nil {
+		return nil, fmt.Errorf("bad cipher block: %w", err)
+	}
+	aesgcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, fmt.Errorf("failed to wrap cipher block: %w", err)
+	}
+	nonce := make([]byte, aesgcm.NonceSize())
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		return nil, fmt.Errorf("failed to generate nonce: %w", err)
+	}
+	ciphertext := aesgcm.Seal(nonce, nonce, plaintext, aad)
+
+	return ciphertext, nil
+}
+
+func (k *InMemory) Decrypt(ctx context.Context, keyID string, ciphertext []byte, aad []byte) ([]byte, error) {
+	dek, err := k.getDEK(keyID)
+	if err != nil {
+		return nil, err
+	}
+
+	block, err := aes.NewCipher(dek)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create cipher from dek: %w", err)
+	}
+
+	aesgcm, err := cipher.NewGCM(block)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create gcm from dek: %w", err)
+	}
+
+	size := aesgcm.NonceSize()
+	if len(ciphertext) < size {
+		return nil, fmt.Errorf("malformed ciphertext")
+	}
+	nonce, ciphertextPortion := ciphertext[:size], ciphertext[size:]
+
+	plaintext, err := aesgcm.Open(nil, nonce, ciphertextPortion, aad)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decrypt ciphertext with dek: %w", err)
+	}
+
+	return plaintext, nil
+}


Test patch:
diff --git a/internal/revision/database/revision_test.go b/internal/revision/database/revision_test.go
new file mode 100644
index 0000000..e1c01d9
--- /dev/null
+++ b/internal/revision/database/revision_test.go
@@ -0,0 +1,178 @@
+// 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 database
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/pkg/keys"
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+)
+
+var (
+	timeCmp = cmpopts.EquateApproxTime(time.Millisecond)
+)
+
+func TestRevisionKey(t *testing.T) {
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
+	ctx := context.Background()
+
+	kms, err := keys.NewInMemory(ctx)
+	if err != nil {
+		t.Fatalf("unable to cerate in memory KMS")
+	}
+	keyID := "funkey"
+	if err := kms.AddEncryptionKey(keyID); err != nil {
+		t.Fatalf("unable to generate key: %v", err)
+	}
+
+	cfg := KMSConfig{keyID, []byte("super"), kms}
+	revDB, err := New(testDB, &cfg)
+	if err != nil {
+		t.Fatalf("unable to provision revision DB: %v", err)
+	}
+
+	want, err := revDB.CreateRevisionKey(ctx)
+	if err != nil {
+		t.Fatalf("unable to create revision key: %v", err)
+	}
+
+	got, err := revDB.GetEffectiveRevisionKey(ctx)
+	if err != nil {
+		t.Fatalf("unable to read effective revision keyL: %v", err)
+	}
+
+	if diff := cmp.Diff(want, got, timeCmp); diff != "" {
+		t.Fatalf("mismatch (-want, +got):\n%s", diff)
+	}
+}
+
+func TestMultipleRevisionKeys(t *testing.T) {
+	t.Parallel()
+
+	testDB := database.NewTestDatabase(t)
+	ctx := context.Background()
+
+	kms, err := keys.NewInMemory(ctx)
+	if err != nil {
+		t.Fatalf("unable to cerate in memory KMS")
+	}
+	keyID := "funkey"
+	if err := kms.AddEncryptionKey(keyID); err != nil {
+		t.Fatalf("unable to generate key: %v", err)
+	}
+
+	cfg := KMSConfig{keyID, []byte("super"), kms}
+	revDB, err := New(testDB, &cfg)
+	if err != nil {
+		t.Fatalf("unable to provision revision DB: %v", err)
+	}
+
+	// Create two revision keys.
+	key1, err := revDB.CreateRevisionKey(ctx)
+	if err != nil {
+		t.Fatalf("failed to create revision key: %v", err)
+	}
+	time.Sleep(500 * time.Millisecond)
+	key2, err := revDB.CreateRevisionKey(ctx)
+	if err != nil {
+		t.Fatalf("failed to create revision key: %v", err)
+	}
+
+	// Check effective revision key (Should be second one)
+	{
+		got, err := revDB.GetEffectiveRevisionKey(ctx)
+		if err != nil {
+			t.Fatalf("unable to read effective keys: %v", err)
+		}
+		if diff := cmp.Diff(key2, got, timeCmp); diff != "" {
+			t.Fatalf("wrong effective key (-want, +got):\n%s", diff)
+		}
+	}
+
+	sorter := cmpopts.SortSlices(
+		func(a, b *RevisionKey) bool {
+			return a.CreatedAt.Before(b.CreatedAt)
+		})
+
+	// Check that both keys are available.
+	{
+		gotID, got, err := revDB.GetAllowedRevisionKeys(ctx)
+		if err != nil {
+			t.Fatalf("unable to read all keys: %v", err)
+		}
+		want := []*RevisionKey{key1, key2}
+		if diff := cmp.Diff(want, got, sorter, timeCmp); diff != "" {
+			t.Fatalf("mismatch (-want, +got):\n%s", diff)
+		}
+		if gotID != key2.KeyID {
+			t.Fatalf("wrong effective key ID want: %v got: %v", key2.KeyID, gotID)
+		}
+	}
+	// And check that both keys are available in the ID check.
+	{
+		want := map[int64]struct{}{
+			key1.KeyID: {},
+			key2.KeyID: {},
+		}
+		gotID, got, err := revDB.GetAllowedRevisionKeyIDs(ctx)
+		if err != nil {
+			t.Fatalf("unable to get allowed key IDs: %v", err)
+		}
+		if diff := cmp.Diff(want, got); diff != "" {
+			t.Fatalf("mismatch (-want, +got):\n%s", diff)
+		}
+		if gotID != key2.KeyID {
+			t.Fatalf("wrong effective key, want: %v got: %v", key2.KeyID, gotID)
+		}
+	}
+
+	// Destroy key2
+	if err := revDB.DestroyKey(ctx, key2.KeyID); err != nil {
+		t.Fatalf("unable to detroy key: %v", err)
+	}
+
+	// Check that effective revision key has changed.
+	{
+		got, err := revDB.GetEffectiveRevisionKey(ctx)
+		if err != nil {
+			t.Fatalf("unable to read effective keys: %v", err)
+		}
+		if diff := cmp.Diff(key1, got, timeCmp); diff != "" {
+			t.Fatalf("wrong effective key (-want, +got):\n%s", diff)
+		}
+	}
+
+	// Only key 2 should be available now.
+	{
+		gotID, got, err := revDB.GetAllowedRevisionKeys(ctx)
+		if err != nil {
+			t.Fatalf("unable to read all keys: %v", err)
+		}
+		want := []*RevisionKey{key1}
+		if diff := cmp.Diff(want, got, sorter, timeCmp); diff != "" {
+			t.Fatalf("mismatch (-want, +got):\n%s", diff)
+		}
+		if gotID != key1.KeyID {
+			t.Fatalf("wrong effective key ID: want: %v got: %v", key1.KeyID, gotID)
+		}
+	}
+}
diff --git a/internal/revision/revision_test.go b/internal/revision/revision_test.go
index 1d5e8f6..0332c00 100644
--- a/internal/revision/revision_test.go
+++ b/internal/revision/revision_test.go
@@ -15,12 +15,18 @@
 package revision
 
 import (
+	"context"
 	"crypto/rand"
+	"fmt"
 	"io"
 	"testing"
+	"time"
 
+	"github.com/google/exposure-notifications-server/internal/database"
 	"github.com/google/exposure-notifications-server/internal/pb"
 	"github.com/google/exposure-notifications-server/internal/publish/model"
+	revisiondb "github.com/google/exposure-notifications-server/internal/revision/database"
+	"github.com/google/exposure-notifications-server/pkg/keys"
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 )
@@ -92,11 +98,44 @@ func TestBuildTokenBuffer(t *testing.T) {
 }
 
 func TestEncryptDecrypt(t *testing.T) {
+	t.Parallel()
+	testDB := database.NewTestDatabase(t)
+	ctx := context.Background()
+
+	kms, err := keys.NewInMemory(ctx)
+	if err != nil {
+		t.Fatalf("unable to cerate in memory KMS")
+	}
+	keyID := "skeleton"
+	if err := kms.AddEncryptionKey(keyID); err != nil {
+		t.Fatalf("unable to generate key: %v", err)
+	}
 	key := make([]byte, 32)
 	if _, err := io.ReadFull(rand.Reader, key); err != nil {
 		t.Fatalf("unable to generate AES key: %v", err)
 	}
 
+	cfg := revisiondb.KMSConfig{
+		WrapperKeyID: keyID,
+		WrapperAAD:   []byte("super"),
+		KeyManager:   kms,
+	}
+	revDB, err := revisiondb.New(testDB, &cfg)
+	if err != nil {
+		t.Fatalf("unable to provision revision DB: %v", err)
+	}
+
+	// Add an allowed key.
+	firstKey, err := revDB.CreateRevisionKey(ctx)
+	if err != nil {
+		t.Fatalf("unable to create a revision key: %v", err)
+	}
+
+	tm, err := New(ctx, revDB, time.Second)
+	if err != nil {
+		t.Fatalf("unable to build token manager: %v", err)
+	}
+
 	// Build data to serialize and encrypt.
 	source := []*model.Exposure{
 		{
@@ -111,7 +150,6 @@ func TestEncryptDecrypt(t *testing.T) {
 		},
 	}
 	aad := []byte{16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1}
-	kid := "123"
 	// Expected result in the end
 	want := &pb.RevisionTokenData{
 		RevisableKeys: []*pb.RevisableKey{
@@ -128,15 +166,12 @@ func TestEncryptDecrypt(t *testing.T) {
 		},
 	}
 
-	encrypted, err := MakeRevisionToken(source, aad, kid, key)
+	encrypted, err := tm.MakeRevisionToken(ctx, source, aad)
 	if err != nil {
 		t.Fatalf("error encrypting and serializing data: %v", err)
 	}
 
-	validKeys := map[string][]byte{
-		kid: key,
-	}
-	got, err := UnmarshalRevisionToken(encrypted, aad, validKeys)
+	got, err := tm.UnmarshalRevisionToken(ctx, encrypted, aad)
 	if err != nil {
 		t.Fatalf("error decrypting token: %v", err)
 	}
@@ -144,4 +179,21 @@ func TestEncryptDecrypt(t *testing.T) {
 	if diff := cmp.Diff(want, got, cmpopts.IgnoreUnexported(pb.RevisionTokenData{}), cmpopts.IgnoreUnexported(pb.RevisableKey{})); diff != "" {
 		t.Fatalf("mismatch (-want, +got):\n%s", diff)
 	}
+
+	// Replace the effective key and destroy the original.
+	if _, err := revDB.CreateRevisionKey(ctx); err != nil {
+		t.Fatalf("unable to create second revision key: %v", err)
+	}
+	if err := revDB.DestroyKey(ctx, firstKey.KeyID); err != nil {
+		t.Fatalf("unable to destroy revision key: %v", err)
+	}
+
+	// force expire the cache.
+	tm.cacheRefreshAfter = time.Now().Add(-1 * time.Second)
+
+	// Decryption of the key should fail this time.
+	wantErr := fmt.Sprintf("token has invalid key id: %v", firstKey.KeyID)
+	if _, err := tm.UnmarshalRevisionToken(ctx, encrypted, aad); err == nil || err.Error() != wantErr {
+		t.Fatalf("want error: %v, got: %v", wantErr, err)
+	}
 }