From 5a8f78d9d0e19076146304537075da755cf48bd9 Mon Sep 17 00:00:00 2001
From: Mariano Cano <mariano@smallstep.com>
Date: Tue, 5 Mar 2019 14:45:57 -0800
Subject: [PATCH] Add support to collection to load the encrypted keys.

---
 authority/provisioner/collection.go  | 38 ++++++++++++++++++++--------
 authority/provisioner/provisioner.go |  6 +++++
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/authority/provisioner/collection.go b/authority/provisioner/collection.go
index cd5d8934..d3bfb996 100644
--- a/authority/provisioner/collection.go
+++ b/authority/provisioner/collection.go
@@ -8,27 +8,27 @@ import (
 
 // Collection is a memory map of provisioners.
 type Collection struct {
-	byID *sync.Map
+	byID  *sync.Map
+	byKey *sync.Map
 }
 
 // NewCollection initializes a collection of provisioners.
 func NewCollection() *Collection {
 	return &Collection{
-		byID: new(sync.Map),
+		byID:  new(sync.Map),
+		byKey: new(sync.Map),
 	}
 }
 
 // Load a provisioner by the ID.
 func (c *Collection) Load(id string) (*Provisioner, bool) {
-	i, ok := c.byID.Load(id)
-	if !ok {
-		return nil, false
-	}
-	p, ok := i.(*Provisioner)
-	if !ok {
-		return nil, false
-	}
-	return p, true
+	return loadProvisioner(c.byID, id)
+}
+
+// LoadEncryptedKey returns a the encrypted key by KeyID. At this moment only
+// JWK encrypted keys are indexed by KeyID.
+func (c *Collection) LoadEncryptedKey(keyID string) (*Provisioner, bool) {
+	return loadProvisioner(c.byKey, keyID)
 }
 
 // Store adds a provisioner to the collection, it makes sure two provisioner
@@ -37,5 +37,21 @@ func (c *Collection) Store(p *Provisioner) error {
 	if _, loaded := c.byID.LoadOrStore(p.ID(), p); loaded == false {
 		return errors.New("cannot add multiple provisioners with the same id")
 	}
+	// Store EncryptedKey if defined
+	if kid, key, ok := p.EncryptedKey(); ok {
+		c.byKey.Store(kid, key)
+	}
 	return nil
 }
+
+func loadProvisioner(m *sync.Map, id string) (*Provisioner, bool) {
+	i, ok := m.Load(id)
+	if !ok {
+		return nil, false
+	}
+	p, ok := i.(*Provisioner)
+	if !ok {
+		return nil, false
+	}
+	return p, true
+}
diff --git a/authority/provisioner/provisioner.go b/authority/provisioner/provisioner.go
index e2895cdc..ce6a1bed 100644
--- a/authority/provisioner/provisioner.go
+++ b/authority/provisioner/provisioner.go
@@ -10,6 +10,7 @@ import (
 // Interface is the interface that all provisioner types must implement.
 type Interface interface {
 	ID() string
+	EncryptedKey() (kid string, key string, ok bool)
 	Init(claims *Claims) error
 	Authorize(token string) ([]SignOption, error)
 }
@@ -43,6 +44,11 @@ func (p *Provisioner) ID() string {
 	return p.base.ID()
 }
 
+// EncryptedKey returns the base provisioner encrypted key if it's defined.
+func (p *Provisioner) EncryptedKey() (string, string, bool) {
+	return p.base.EncryptedKey()
+}
+
 // Type return the provisioners type.
 func (p *Provisioner) Type() Type {
 	return p.typ