From db337debcdc6ca59cde20900d114f0f48bc12a33 Mon Sep 17 00:00:00 2001
From: Mariano Cano <mariano.cano@gmail.com>
Date: Tue, 5 Apr 2022 19:25:47 -0700
Subject: [PATCH] Load provisioner from the database instead of the extension.

---
 authority/provisioners.go | 34 +++++++++++++++++++++++++++++-----
 authority/tls.go          |  2 ++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/authority/provisioners.go b/authority/provisioners.go
index 496421f5..c0eec520 100644
--- a/authority/provisioners.go
+++ b/authority/provisioners.go
@@ -13,6 +13,7 @@ import (
 	"github.com/smallstep/certificates/authority/admin"
 	"github.com/smallstep/certificates/authority/config"
 	"github.com/smallstep/certificates/authority/provisioner"
+	"github.com/smallstep/certificates/db"
 	"github.com/smallstep/certificates/errs"
 	"go.step.sm/cli-utils/step"
 	"go.step.sm/cli-utils/ui"
@@ -44,13 +45,36 @@ func (a *Authority) GetProvisioners(cursor string, limit int) (provisioner.List,
 // LoadProvisionerByCertificate returns an interface to the provisioner that
 // provisioned the certificate.
 func (a *Authority) LoadProvisionerByCertificate(crt *x509.Certificate) (provisioner.Interface, error) {
+	// Default implementation looks at the provisioner extension.
+	loadProvisioner := func() (provisioner.Interface, error) {
+		p, ok := a.provisioners.LoadByCertificate(crt)
+		if !ok {
+			return nil, admin.NewError(admin.ErrorNotFoundType, "unable to load provisioner from certificate")
+		}
+		return p, nil
+	}
+
+	// Attempt to load the provisioner using the linked db
+	// TODO:(mariano)
+
+	// Attempt to load the provisioner from the db
+	if db, ok := a.db.(interface {
+		GetCertificateData(string) (*db.CertificateData, error)
+	}); ok {
+		if data, err := db.GetCertificateData(crt.SerialNumber.String()); err == nil && data.Provisioner != nil {
+			loadProvisioner = func() (provisioner.Interface, error) {
+				p, ok := a.provisioners.Load(data.Provisioner.ID)
+				if !ok {
+					return nil, admin.NewError(admin.ErrorNotFoundType, "unable to load provisioner from certificate")
+				}
+				return p, nil
+			}
+		}
+	}
+
 	a.adminMutex.RLock()
 	defer a.adminMutex.RUnlock()
-	p, ok := a.provisioners.LoadByCertificate(crt)
-	if !ok {
-		return nil, admin.NewError(admin.ErrorNotFoundType, "unable to load provisioner from certificate")
-	}
-	return p, nil
+	return loadProvisioner()
 }
 
 // LoadProvisionerByToken returns an interface to the provisioner that
diff --git a/authority/tls.go b/authority/tls.go
index 93bc0408..bebcdf1b 100644
--- a/authority/tls.go
+++ b/authority/tls.go
@@ -347,6 +347,8 @@ func (a *Authority) storeCertificate(prov provisioner.Interface, fullchain []*x5
 
 	// Store certificate in local db
 	switch s := a.db.(type) {
+	case linkedChainStorer:
+		return s.StoreCertificateChain(prov, fullchain...)
 	case certificateChainStorer:
 		return s.StoreCertificateChain(fullchain...)
 	default: